Cosmos3-Nano 世界基础模型昇腾 NPU 适配与推理实战从环境搭建到多卡并行【免费下载链接】cann-recipes-embodied-ai本项目针对具身智能业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-ai本文围绕 CANN 适配样例仓库中world_model/cosmos3目录提供的 Cosmos3 昇腾 NPU 适配方案展开完整介绍如何在昇腾 A3 环境上拉取 Cosmos3 框架、配置 CANN 9.0 依赖、执行npu_adapt.sh适配脚本并运行 Cosmos3-Nano 的文生视频T2V、图生视频I2V、视频生视频V2V推理与 CP/CFGP/FSDP 多卡并行验证。读完本文你将掌握 Cosmos3 世界基础模型在 CANN 平台上的完整落地路径并理解 FIA attention、HCCL 分布式通信等关键适配点的底层原理。Cosmos3 整体介绍功能介绍Cosmos3 是一个世界基础模型World Foundation Models框架面向物理 AI、机器人、自动驾驶、视频生成与多模态理解等场景。Cosmos3-Nano 支持文生视频T2V、图生视频I2V、视频生视频V2V等推理任务并可结合结构化输入完成世界模型生成与理解。本样例基于 Cosmos3 框架完成昇腾 NPU 适配提供以下能力依赖配置适配后的 pyproject.toml 将torch_npu2.10.0纳入依赖并指定 PyTorch 走 CPU 索引、torch_npu 走华为云 PyPI 镜像索引实现一键式环境还原设备适配脚本npu_adapt.sh 对 Cosmos3 源码做机械化的 CUDA→NPU 替换与后端注册覆盖推理入口、flags、模型加载、分布式运行时、序列打包、注意力后端等十余处FIA attention 后端新增 CANN 注意力后端见 cann 子包通过torch_npu.npu_fused_infer_attention_score提供 BNSD/TND 两种布局的融合注意力计算本地权重路径适配inference_local_checkpoint.patch 支持完全离线的本地 checkpoint 加载避免重复下载 tokenizer、Wan2.2 VAE 与 AVAE 资源基础推理验证命令覆盖 T2V/I2V/V2V 三种输入模式与 CP/CFGP/FSDP 三种并行策略。代码仓拉取与适配文件覆盖本样例需要两个代码仓CANN 适配仓当前仓库与 Cosmos3 原始代码仓。使用时先将适配文件覆盖到 Cosmos3 仓根目录再在 Cosmos3 仓内完成环境安装与推理。建议让cann-recipes-embodied-ai与cosmos-framework保持同级目录# 进入需要放置代码仓的本地目录建议让 cann-recipes-embodied-ai 与 cosmos-framework 保持同级 git clone https://gitcode.com/cann/cann-recipes-embodied-ai.git # 从 NVIDIA 官方渠道克隆 cosmos-framework 代码仓并切换到适配基线提交 a61b292 git clone cosmos-framework 仓库地址 cd cosmos-framework git checkout a61b292 # 回到两个代码仓的共同上级目录 cd ../ # 将 CANN 仓中的 Cosmos3 适配文件覆盖到 Cosmos3 仓根目录 cp -rf cann-recipes-embodied-ai/world_model/cosmos3/* ./cosmos-framework完成覆盖后cosmos-framework根目录下应包含npu_adapt.sh、pyproject.toml等适配文件。需要特别说明的是基线提交 a61b292 是硬性约束npu_adapt.sh在运行前会执行verify_expected_commit校验只有当前cosmos-framework的 HEAD 与预期提交一致才会继续执行源码改写否则直接报错退出这是为了防止 sed/perl 改写脚本作用在错误的源码布局上。Cosmos3 在昇腾 A3 上的运行环境配置与昇腾平台相关的环境配置安装 CANN 软件包。本样例依赖 CANN 开发套件包cann-toolkit与 CANN 二进制算子包cann-kernels支持的软件版本为CANN 9.0.0、torch_npu2.10.0、Python3.13。从昇腾社区软件包下载页面获取Ascend-cann-toolkit_${version}_linux-${aarch}.run与Atlas-A3-cann-kernels_${version}_linux-${aarch}.run两个软件包并参考 CANN 官方安装文档依次安装。其中${version}表示 CANN 包版本号如 9.0.0${aarch}表示 CPU 架构如 aarch64、x86_64。安装完成后每次新建终端时首先 source 环境变量脚本${cann_install_path}为 CANN 包的实际安装目录# 方式1默认路径安装以 root 用户为例 source /usr/local/Ascend/ascend-toolkit/set_env.sh # 方式2指定路径进行安装 source ${cann_install_path}/ascend-toolkit/set_env.shuv 环境管理工具安装本样例使用 uv 管理 Python 依赖其版本要求与索引配置已固化在 pyproject.toml 的[tool.uv]段required-version 0.11.3。如果当前环境已经安装 uv可跳过此步wget -qO- https://astral.sh/uv/install.sh | shPython 运行环境安装在cosmos-framework根目录下执行uv sync指定 Python 3.13cd cosmos-framework uv sync --python 3.13从依赖清单可以进一步理解适配的关键点见 pyproject.toml核心依赖直接声明torch_npu2.10.0与 CANN 9.0.0 配套[tool.uv.sources]中torch/torchvision走pytorch-cpu索引避免默认拉取 CUDA 版torch_npu走华为云 PyPI 镜像索引二者均为explicit true只有显式声明才会生效override-dependencies将 numpy 锁定在2.0.0,2.3兼容 robosuite/numba 传递依赖同时通过pynvml; sys_platform never方式隔离 NVML 包配合适配脚本将 pynvml 变为可选导入可选依赖组guardrail安全护栏相关、servegradio/ray 服务、train训练相关可按需安装。模型权重下载本样例使用 Cosmos3-Nano 权重进行推理验证。从nvidia/Cosmos3-Nano模型仓库下载权重并将推理命令中的COSMOS_CHECKPOINT指向本地权重目录。此外视频生成还需要 Wan2.2 VAE 权重。从Wan-AI/Wan2.2-TI2V-5B仓库下载Wan2.2_VAE.pth即可——只需要该 VAE 权重文件无需下载完整 Wan2.2-TI2V-5B 模型——并将其放置到 Cosmos3-Nano 本地权重目录下# 示例将 Cosmos3-Nano 权重放置在本地目录 /mnt/workspace/cosmos3/cosmos3-nano export COSMOS_CHECKPOINT/mnt/workspace/cosmos3/cosmos3-nano # Wan2.2 VAE 权重应放置为如下路径 ls ${COSMOS_CHECKPOINT}/Wan2.2_VAE.pthWan2.2_VAE.pth放在 checkpoint 目录下并非随意为之适配脚本会应用 inference_local_checkpoint.patch其中显式检查Path(checkpoint_path) / Wan2.2_VAE.pth是否存在存在时会将 VAE 配置的bucket_name置空并令vae_path指向该本地文件从而完全跳过从远端下载 VAE。如果部署环境无法直接访问 Hugging Face可在可联网环境下载完整 Cosmos3-Nano 权重目录和Wan2.2_VAE.pth后拷贝到昇腾服务器再使用本地路径作为--checkpoint-path。执行 NPU 适配脚本完成文件覆盖后在cosmos-framework根目录执行cd cosmos-framework bash npu_adapt.sh脚本整体结构见 npu_adapt.sh。其开头注释明确说明脚本只包含从基线a61b292到a93d52c的简单/机械化改写有意排除了 uv.lock 与复杂特性补丁如本地 checkpoint/tokenizer 加载、Qwen3-VL CPU 预处理绕过后者通过补丁文件单独应用。基线提交校验与文件保护verify_expected_commit要求COSMOS_ROOT默认.必须是 git 工作树且 HEAD 等于a61b292否则拒绝运行。所有改写目标文件在修改前都会经过ensure_file检查缺失时打印[WARN] Skip missing file并跳过保证脚本幂等可重入。设备默认值与入口适配ensure_inference_default_npu在 cosmos_framework/scripts/inference.py 的 import 区插入import torch_npu与torch.set_default_device(npu)使入口进程默认在 NPU 上分配张量adapt_flags修改cosmos_framework/utils/flags.py将COSMOS_TRAINING默认值由 True 改为 False、COSMOS_DEVICE默认设备由cuda改为npu并新增Device.NPU npu枚举成员adapt_model_loader在cosmos_framework/utils/vfm/model_loader.py中为 NCCL 后端判断之外增加backend.endswith(hccl)分支返回torch.device(npu, torch.npu.current_device())。显存探测与编译开关adapt_device_memory将cosmos_framework/inference/args.py中基于 pynvml/CUDA 的_get_device_memory_bytes()替换为基于torch.npu.get_device_properties(0).total_memory的实现无 NPU 时回退 64GBadapt_compile_defaults将cosmos_framework/inference/common/args.py中use_torch_compile与use_cuda_graphs默认值改为 False。这是因为 NPU 适配阶段优先保证功能正确性torch.compile 与 CUDA Graph 机制不在默认启用范围adapt_common_init/adapt_common_inference将torch.cuda的 set_device、set_per_process_memory_fraction、错误标志张量设备等统一替换为torch.npu对应接口。分布式运行时 HCCL 化adapt_distributed_runtime是单机多卡的关键改造发生在cosmos_framework/utils/distributed.pyimport pynvml改为 try/except 可选导入NPU 环境通常无 NVMLtorch.cuda.set_device(local_rank)→torch.npu.set_device(local_rank)dist.init_process_group(backendnccl→backendhccl即用昇腾 HCCL 通信库替代 NCCLtorch.cuda.current_device()→torch.npu.current_device()并将日志文案由 Training with {get_world_size()} GPUs 改为 Running with {get_world_size()} NPUs。推理运行时adapt_inference_runtime则对cosmos_framework/inference/inference.py做了更彻底的替换torch.cuda.Stream/Event/stream/device_count/current_stream全部改为torch.npu版本device: Any cuda改为npu。张量搬移层面adapt_transfer_and_vision_inputs将 transfer/vision 输入处理中的.cuda()全部改为.npu()adapt_sequence_packing_device_move将PackedSequence.to_cuda()重命名为to_npu()adapt_omni_device_condition放行Device.NPU上的 OmniMoT 初始化并将torch.cuda.empty_cache()换为torch.npu.empty_cache()。MoE Triton 导入保护adapt_moe_triton_import_guard针对cosmos_framework/model/vfm/vlm/qwen3_vl_moe/moe_kernels.py将import triton改为 try/except 形式并设置_HAS_TRITON标志当 Triton 不可用时将_fill_indices_kernel置为 None避免 Qwen3-VL MoE 模块在无 Triton 的 NPU 环境中因导入失败而中断。CANN FIA attention 后端注册adapt_attention_cann_registration将本仓新增的 CANN 注意力后端接入 Cosmos3 的注意力后端分发机制在cosmos_framework/model/attention/backends.py中引入cann_attention_check并在后端字典注册cann: cann_attention_check在设备为npu时强制backend_list [cann]跳过 CUDA 架构标签arch_tag驱动的后端筛选在cosmos_framework/model/attention/frontend.py中注册cann: cann_attention分发入口移除工具函数中对torch.npu.is_available()返回 80 的 arch_tag 特判。本地 checkpoint 资源加载补丁apply_inference_local_checkpoint_patch调用patch -p1应用 inference_local_checkpoint.patch。该补丁的核心改动包括VLM tokenizer当Path(checkpoint_path, text_tokenizer).is_dir()成立时直接以 checkpoint 目录作为tokenizer_type构建 processor避免从远端仓库下载 tokenizerWan2.2 VAE检测 checkpoint 目录下的Wan2.2_VAE.pth存在则改写 VAE 配置指向本地文件bucket_name、vae_path本地路径AVAE 音频 tokenizerfrom_checkpoint默认值由 False 改为 True优先使用 checkpoint 内捆绑的sound_tokenizer/保证离线本地推理自包含。pynvml 可选化adapt_optional_nvml处理cosmos_framework/utils/device.py将import pynvml包上 try/exceptexcept pynvml.NVMLError放宽为except Exception并仅在pynvml is not None时调用nvmlShutdown()彻底解耦对 NVIDIA 管理库的依赖。CANN FIA Attention 后端原理新增的 CANN 注意力后端位于 cosmos_framework/model/attention/cann/ 子包是 Cosmos3 昇腾适配中最具代表性的算子层工作。后端能力声明checks.py 中的cann_attention_check声明了该后端的支持范围支持数据类型torch.float16、torch.bfloat16支持 GQA/MQAsupports_gqa_mqaTrue不支持 MLAsupports_mlaFalse通过统一的attention_tensor_checks完成张量形状、requires_grad 等一致性校验。融合算子调用与布局转换functions.py 中cann_attention是统一入口优先读取cumulative_seqlen_Q/cumulative_seqlen_KV判断是否为 varlen变长序列打包模式scale缺省时按query.shape[-1] ** -0.5计算return_lse请求时返回(output, None)占位。两条计算路径都落在torch_npu.npu_fused_infer_attention_score上BNSD 标准路径Cosmos3 frontend 传入的是 BSND 布局batch、sequence、num_heads、head_dim而 FIA 标准路径期望 BNSD因此先permute(0, 2, 1, 3)再做融合注意力最后转回 BSND 返回TND 变长路径varlen 模式下 Cosmos3 传入的是 batch1 的 packed BSND 张量先squeeze(0)去掉单例 batch 维再按 FIA 的 TND 布局传递actual_seq_lengths注意 Cosmos 的 cumulative seqlens 含前导 0需切片[1:]只保留累计结束偏移causal 模式使用 2048×2048 的三角掩码_tnd_causal_mask并通过pre_tokens65535、next_tokens65535、sparse_mode3causal或 0 配置稀疏注意力输出按有效长度截断计算后再在序列末尾补零回原始总长度并unsqueeze(0)还原 batch 维。从源码结构看该实现把 Cosmos3 的 BSND 打包张量与昇腾 FIA 的 BNSD/TND 两种输入布局解耦是模型侧无需改动注意力逻辑即可获得融合算子加速的关键。推理验证示例完成适配后可在cosmos-framework根目录下执行推理命令进行基础场景验证。COSMOS_CHECKPOINT用于指定本地权重目录如不设置可直接将命令中的--checkpoint-path替换为实际权重路径。export COSMOS_CHECKPOINT/mnt/workspace/cosmos3/cosmos3-nano export COSMOS_RESOLUTION480 export COSMOS_SEED0 export COSMOS_NPUS1输入 JSON 配置说明推理命令中的-i inputs/omni/t2v.json用于指定单条样例输入。常用字段如下model_mode任务类型例如text2video、image2video、video2videoprompt文本提示词T2V 只需要配置该字段即可vision_pathI2V/V2V 的输入图片或视频路径仅图生视频、视频生视频需要。vision_path可以写远程 URL也可以写本地文件路径。若服务器无法访问远端资源或遇到证书、代理、内网限制等网络问题请先手动下载输入图片/视频到本地然后在 JSON 中改成本地绝对路径例如{ model_mode: image2video, prompt: A robot arm moves smoothly in a lab., vision_path: /mnt/workspace/cosmos3/inputs/robot_153.jpg }T2V 示例 JSON 可简化为{ model_mode: text2video, name: t2v, prompt: A realistic video of molten metal being poured in a steel mill. }T2V 文生视频torchrun --nproc-per-node${COSMOS_NPUS} -m cosmos_framework.scripts.inference \ --parallelism-presetlatency \ -i inputs/omni/t2v.json \ -o outputs/t2v \ --checkpoint-path ${COSMOS_CHECKPOINT} \ --resolution${COSMOS_RESOLUTION} \ --seed${COSMOS_SEED} \ --no-guardrailsI2V 图生视频torchrun --nproc-per-node${COSMOS_NPUS} -m cosmos_framework.scripts.inference \ --parallelism-presetlatency \ -i inputs/omni/i2v.json \ -o outputs/i2v \ --checkpoint-path ${COSMOS_CHECKPOINT} \ --resolution${COSMOS_RESOLUTION} \ --seed${COSMOS_SEED} \ --no-guardrailsV2V 视频生视频torchrun --nproc-per-node${COSMOS_NPUS} -m cosmos_framework.scripts.inference \ --parallelism-presetlatency \ -i inputs/omni/v2v.json \ -o outputs/v2v \ --checkpoint-path ${COSMOS_CHECKPOINT} \ --resolution${COSMOS_RESOLUTION} \ --seed${COSMOS_SEED} \ --no-guardrails命令要点说明torchrun --nproc-per-node${COSMOS_NPUS}以多进程方式拉起分布式推理进程数等于 NPU 数量--parallelism-presetlatency单卡低延迟预设适合功能验证多卡并行场景改用throughput预设--no-guardrails关闭文本/人脸安全护栏对应 pyproject 中的guardrail可选依赖组避免额外依赖与推理开销基础场景默认单卡即可完成此时COSMOS_NPUS1。多卡并行当前适配支持 CPContext Parallel、CFGPClassifier-Free Guidance Parallel和 FSDP 三种多卡推理方式。运行前请将COSMOS_NPUS设置为实际使用的 NPU 数量并保证各并行度与进程数匹配。并行方式主要参数适用目的与约束CPContext Parallel--cp-size沿 token 序列切分 Attention 计算适合长序列并降低激活显存当前 CP 范围为 132CFGPClassifier-Free Guidance Parallel--cfgp-size将有条件与无条件 CFG 分支分配到不同设备CFGP 仅支持 1 或 2更大规模可与 CP/FSDP 组合FSDP--dp-shard-size按进程数切分模型参数优先降低单卡权重显存FSDP 的 DP 通信组与 CP/CFGP 通信组相互独立通信域大小满足dp-shard-size × dp-replicate-size WORLD_SIZE WORLD_SIZE % (cp-size × cfgp-size) 0CP 上下文并行CP 将长序列的 Attention 计算切分到多卡上并行执行适合超长视频/长上下文生成torchrun --nproc-per-node${COSMOS_NPUS} -m cosmos_framework.scripts.inference \ --parallelism-presetthroughput \ --dp-shard-size1 --cp-size${COSMOS_NPUS} --cfgp-size1 \ -i inputs/omni/t2v.json \ -o outputs/t2v_cp \ --checkpoint-path ${COSMOS_CHECKPOINT} \ --resolution${COSMOS_RESOLUTION} \ --seed${COSMOS_SEED} \ --no-guardrailsCFGP 无条件引导并行CFGClassifier-Free Guidance需要同时计算有条件和无条件两个去噪分支CFGP 将两个分支分配到不同设备组并行执行。单独启用时设置COSMOS_NPUS2torchrun --nproc-per-node${COSMOS_NPUS} -m cosmos_framework.scripts.inference \ --parallelism-presetthroughput \ --dp-shard-size1 --cp-size1 --cfgp-size2 \ -i inputs/omni/t2v.json \ -o outputs/t2v_cfgp \ --checkpoint-path ${COSMOS_CHECKPOINT} \ --resolution${COSMOS_RESOLUTION} \ --seed${COSMOS_SEED} \ --no-guardrailsFSDP 参数分片FSDP 按进程数切分模型参数将权重/优化器状态分布到多卡优先降低单卡权重显存占用torchrun --nproc-per-node${COSMOS_NPUS} -m cosmos_framework.scripts.inference \ --parallelism-presetthroughput \ --dp-shard-size${COSMOS_NPUS} --cp-size1 --cfgp-size1 \ -i inputs/omni/t2v.json \ -o outputs/t2v_fsdp \ --checkpoint-path ${COSMOS_CHECKPOINT} \ --resolution${COSMOS_RESOLUTION} \ --seed${COSMOS_SEED} \ --no-guardrails多卡并行之所以能开箱即用与上文adapt_distributed_runtime的 HCCL 化改造直接相关CP/CFGP 通信组与 FSDP 的 DP 通信组均建立在hccl后端之上torch.npu.set_device(local_rank)保证每个 rank 张量落在自己的 NPU 上。同类思路也可参考本仓库 Cosmos 系列多卡并行优化说明其中对 CFG 并行、上下文并行与融合算子Flash Attention、RMSNorm、Rotary的改造与本样例互为印证。样例输出与效果验证在昇腾 NPU 上运行上述基础场景后生成结果会输出到-o指定的目录如outputs/t2v。原仓库 README 的样例输出展示章节提供了 T2V 文生视频、I2V 图生视频、V2V 视频生视频三类场景的示例生成结果可用于快速对照查看生成效果场景说明T2V 文生视频纯文本提示词驱动的视频生成结果I2V 图生视频以单张图片为条件输入的视频生成结果V2V 视频生视频以视频片段为输入的续写/变换结果建议验证时按以下顺序逐步推进先以COSMOS_NPUS1跑通 T2V 单卡链路确认权重加载、VAE 推理与视频落盘正常再依次验证 I2V检查vision_path本地路径可读与 V2V最后切换到throughput预设并按上节公式配置并行度逐步扩展到多卡场景。总结与引用本样例以适配文件覆盖 脚本化源码改写 新增 CANN 后端三层结构将 Cosmos3-Nano 世界基础模型的 T2V/I2V/V2V 推理完整迁移到昇腾 A3 NPU 平台环境侧锁定 CANN 9.0.0 / torch_npu 2.10.0 / Python 3.13设备侧完成 CUDA→NPU 与 NCCL→HCCL 的全面替换算子侧通过npu_fused_infer_attention_score提供 FIA attention 加速推理侧通过本地 checkpoint 补丁实现完全离线运行并行侧支持 CP/CFGP/FSDP 的灵活组合。读者可依据本文从零搭建环境并完成多场景验证亦可深入 npu_adapt.sh 与 cann attention 实现 理解每个适配点的源码级细节。如需引用 Cosmos3 模型本身可参考上游论文《Cosmos 3: Omnimodal World Models for Physical AI》NVIDIA 团队2026 年其 BibTeX 引用信息收录于原 README 的 citation 章节。【免费下载链接】cann-recipes-embodied-ai本项目针对具身智能业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-ai创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
Cosmos3-Nano 世界基础模型昇腾 NPU 适配与推理实战：从环境搭建到多卡并行
Cosmos3-Nano 世界基础模型昇腾 NPU 适配与推理实战从环境搭建到多卡并行【免费下载链接】cann-recipes-embodied-ai本项目针对具身智能业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-ai本文围绕 CANN 适配样例仓库中world_model/cosmos3目录提供的 Cosmos3 昇腾 NPU 适配方案展开完整介绍如何在昇腾 A3 环境上拉取 Cosmos3 框架、配置 CANN 9.0 依赖、执行npu_adapt.sh适配脚本并运行 Cosmos3-Nano 的文生视频T2V、图生视频I2V、视频生视频V2V推理与 CP/CFGP/FSDP 多卡并行验证。读完本文你将掌握 Cosmos3 世界基础模型在 CANN 平台上的完整落地路径并理解 FIA attention、HCCL 分布式通信等关键适配点的底层原理。Cosmos3 整体介绍功能介绍Cosmos3 是一个世界基础模型World Foundation Models框架面向物理 AI、机器人、自动驾驶、视频生成与多模态理解等场景。Cosmos3-Nano 支持文生视频T2V、图生视频I2V、视频生视频V2V等推理任务并可结合结构化输入完成世界模型生成与理解。本样例基于 Cosmos3 框架完成昇腾 NPU 适配提供以下能力依赖配置适配后的 pyproject.toml 将torch_npu2.10.0纳入依赖并指定 PyTorch 走 CPU 索引、torch_npu 走华为云 PyPI 镜像索引实现一键式环境还原设备适配脚本npu_adapt.sh 对 Cosmos3 源码做机械化的 CUDA→NPU 替换与后端注册覆盖推理入口、flags、模型加载、分布式运行时、序列打包、注意力后端等十余处FIA attention 后端新增 CANN 注意力后端见 cann 子包通过torch_npu.npu_fused_infer_attention_score提供 BNSD/TND 两种布局的融合注意力计算本地权重路径适配inference_local_checkpoint.patch 支持完全离线的本地 checkpoint 加载避免重复下载 tokenizer、Wan2.2 VAE 与 AVAE 资源基础推理验证命令覆盖 T2V/I2V/V2V 三种输入模式与 CP/CFGP/FSDP 三种并行策略。代码仓拉取与适配文件覆盖本样例需要两个代码仓CANN 适配仓当前仓库与 Cosmos3 原始代码仓。使用时先将适配文件覆盖到 Cosmos3 仓根目录再在 Cosmos3 仓内完成环境安装与推理。建议让cann-recipes-embodied-ai与cosmos-framework保持同级目录# 进入需要放置代码仓的本地目录建议让 cann-recipes-embodied-ai 与 cosmos-framework 保持同级 git clone https://gitcode.com/cann/cann-recipes-embodied-ai.git # 从 NVIDIA 官方渠道克隆 cosmos-framework 代码仓并切换到适配基线提交 a61b292 git clone cosmos-framework 仓库地址 cd cosmos-framework git checkout a61b292 # 回到两个代码仓的共同上级目录 cd ../ # 将 CANN 仓中的 Cosmos3 适配文件覆盖到 Cosmos3 仓根目录 cp -rf cann-recipes-embodied-ai/world_model/cosmos3/* ./cosmos-framework完成覆盖后cosmos-framework根目录下应包含npu_adapt.sh、pyproject.toml等适配文件。需要特别说明的是基线提交 a61b292 是硬性约束npu_adapt.sh在运行前会执行verify_expected_commit校验只有当前cosmos-framework的 HEAD 与预期提交一致才会继续执行源码改写否则直接报错退出这是为了防止 sed/perl 改写脚本作用在错误的源码布局上。Cosmos3 在昇腾 A3 上的运行环境配置与昇腾平台相关的环境配置安装 CANN 软件包。本样例依赖 CANN 开发套件包cann-toolkit与 CANN 二进制算子包cann-kernels支持的软件版本为CANN 9.0.0、torch_npu2.10.0、Python3.13。从昇腾社区软件包下载页面获取Ascend-cann-toolkit_${version}_linux-${aarch}.run与Atlas-A3-cann-kernels_${version}_linux-${aarch}.run两个软件包并参考 CANN 官方安装文档依次安装。其中${version}表示 CANN 包版本号如 9.0.0${aarch}表示 CPU 架构如 aarch64、x86_64。安装完成后每次新建终端时首先 source 环境变量脚本${cann_install_path}为 CANN 包的实际安装目录# 方式1默认路径安装以 root 用户为例 source /usr/local/Ascend/ascend-toolkit/set_env.sh # 方式2指定路径进行安装 source ${cann_install_path}/ascend-toolkit/set_env.shuv 环境管理工具安装本样例使用 uv 管理 Python 依赖其版本要求与索引配置已固化在 pyproject.toml 的[tool.uv]段required-version 0.11.3。如果当前环境已经安装 uv可跳过此步wget -qO- https://astral.sh/uv/install.sh | shPython 运行环境安装在cosmos-framework根目录下执行uv sync指定 Python 3.13cd cosmos-framework uv sync --python 3.13从依赖清单可以进一步理解适配的关键点见 pyproject.toml核心依赖直接声明torch_npu2.10.0与 CANN 9.0.0 配套[tool.uv.sources]中torch/torchvision走pytorch-cpu索引避免默认拉取 CUDA 版torch_npu走华为云 PyPI 镜像索引二者均为explicit true只有显式声明才会生效override-dependencies将 numpy 锁定在2.0.0,2.3兼容 robosuite/numba 传递依赖同时通过pynvml; sys_platform never方式隔离 NVML 包配合适配脚本将 pynvml 变为可选导入可选依赖组guardrail安全护栏相关、servegradio/ray 服务、train训练相关可按需安装。模型权重下载本样例使用 Cosmos3-Nano 权重进行推理验证。从nvidia/Cosmos3-Nano模型仓库下载权重并将推理命令中的COSMOS_CHECKPOINT指向本地权重目录。此外视频生成还需要 Wan2.2 VAE 权重。从Wan-AI/Wan2.2-TI2V-5B仓库下载Wan2.2_VAE.pth即可——只需要该 VAE 权重文件无需下载完整 Wan2.2-TI2V-5B 模型——并将其放置到 Cosmos3-Nano 本地权重目录下# 示例将 Cosmos3-Nano 权重放置在本地目录 /mnt/workspace/cosmos3/cosmos3-nano export COSMOS_CHECKPOINT/mnt/workspace/cosmos3/cosmos3-nano # Wan2.2 VAE 权重应放置为如下路径 ls ${COSMOS_CHECKPOINT}/Wan2.2_VAE.pthWan2.2_VAE.pth放在 checkpoint 目录下并非随意为之适配脚本会应用 inference_local_checkpoint.patch其中显式检查Path(checkpoint_path) / Wan2.2_VAE.pth是否存在存在时会将 VAE 配置的bucket_name置空并令vae_path指向该本地文件从而完全跳过从远端下载 VAE。如果部署环境无法直接访问 Hugging Face可在可联网环境下载完整 Cosmos3-Nano 权重目录和Wan2.2_VAE.pth后拷贝到昇腾服务器再使用本地路径作为--checkpoint-path。执行 NPU 适配脚本完成文件覆盖后在cosmos-framework根目录执行cd cosmos-framework bash npu_adapt.sh脚本整体结构见 npu_adapt.sh。其开头注释明确说明脚本只包含从基线a61b292到a93d52c的简单/机械化改写有意排除了 uv.lock 与复杂特性补丁如本地 checkpoint/tokenizer 加载、Qwen3-VL CPU 预处理绕过后者通过补丁文件单独应用。基线提交校验与文件保护verify_expected_commit要求COSMOS_ROOT默认.必须是 git 工作树且 HEAD 等于a61b292否则拒绝运行。所有改写目标文件在修改前都会经过ensure_file检查缺失时打印[WARN] Skip missing file并跳过保证脚本幂等可重入。设备默认值与入口适配ensure_inference_default_npu在 cosmos_framework/scripts/inference.py 的 import 区插入import torch_npu与torch.set_default_device(npu)使入口进程默认在 NPU 上分配张量adapt_flags修改cosmos_framework/utils/flags.py将COSMOS_TRAINING默认值由 True 改为 False、COSMOS_DEVICE默认设备由cuda改为npu并新增Device.NPU npu枚举成员adapt_model_loader在cosmos_framework/utils/vfm/model_loader.py中为 NCCL 后端判断之外增加backend.endswith(hccl)分支返回torch.device(npu, torch.npu.current_device())。显存探测与编译开关adapt_device_memory将cosmos_framework/inference/args.py中基于 pynvml/CUDA 的_get_device_memory_bytes()替换为基于torch.npu.get_device_properties(0).total_memory的实现无 NPU 时回退 64GBadapt_compile_defaults将cosmos_framework/inference/common/args.py中use_torch_compile与use_cuda_graphs默认值改为 False。这是因为 NPU 适配阶段优先保证功能正确性torch.compile 与 CUDA Graph 机制不在默认启用范围adapt_common_init/adapt_common_inference将torch.cuda的 set_device、set_per_process_memory_fraction、错误标志张量设备等统一替换为torch.npu对应接口。分布式运行时 HCCL 化adapt_distributed_runtime是单机多卡的关键改造发生在cosmos_framework/utils/distributed.pyimport pynvml改为 try/except 可选导入NPU 环境通常无 NVMLtorch.cuda.set_device(local_rank)→torch.npu.set_device(local_rank)dist.init_process_group(backendnccl→backendhccl即用昇腾 HCCL 通信库替代 NCCLtorch.cuda.current_device()→torch.npu.current_device()并将日志文案由 Training with {get_world_size()} GPUs 改为 Running with {get_world_size()} NPUs。推理运行时adapt_inference_runtime则对cosmos_framework/inference/inference.py做了更彻底的替换torch.cuda.Stream/Event/stream/device_count/current_stream全部改为torch.npu版本device: Any cuda改为npu。张量搬移层面adapt_transfer_and_vision_inputs将 transfer/vision 输入处理中的.cuda()全部改为.npu()adapt_sequence_packing_device_move将PackedSequence.to_cuda()重命名为to_npu()adapt_omni_device_condition放行Device.NPU上的 OmniMoT 初始化并将torch.cuda.empty_cache()换为torch.npu.empty_cache()。MoE Triton 导入保护adapt_moe_triton_import_guard针对cosmos_framework/model/vfm/vlm/qwen3_vl_moe/moe_kernels.py将import triton改为 try/except 形式并设置_HAS_TRITON标志当 Triton 不可用时将_fill_indices_kernel置为 None避免 Qwen3-VL MoE 模块在无 Triton 的 NPU 环境中因导入失败而中断。CANN FIA attention 后端注册adapt_attention_cann_registration将本仓新增的 CANN 注意力后端接入 Cosmos3 的注意力后端分发机制在cosmos_framework/model/attention/backends.py中引入cann_attention_check并在后端字典注册cann: cann_attention_check在设备为npu时强制backend_list [cann]跳过 CUDA 架构标签arch_tag驱动的后端筛选在cosmos_framework/model/attention/frontend.py中注册cann: cann_attention分发入口移除工具函数中对torch.npu.is_available()返回 80 的 arch_tag 特判。本地 checkpoint 资源加载补丁apply_inference_local_checkpoint_patch调用patch -p1应用 inference_local_checkpoint.patch。该补丁的核心改动包括VLM tokenizer当Path(checkpoint_path, text_tokenizer).is_dir()成立时直接以 checkpoint 目录作为tokenizer_type构建 processor避免从远端仓库下载 tokenizerWan2.2 VAE检测 checkpoint 目录下的Wan2.2_VAE.pth存在则改写 VAE 配置指向本地文件bucket_name、vae_path本地路径AVAE 音频 tokenizerfrom_checkpoint默认值由 False 改为 True优先使用 checkpoint 内捆绑的sound_tokenizer/保证离线本地推理自包含。pynvml 可选化adapt_optional_nvml处理cosmos_framework/utils/device.py将import pynvml包上 try/exceptexcept pynvml.NVMLError放宽为except Exception并仅在pynvml is not None时调用nvmlShutdown()彻底解耦对 NVIDIA 管理库的依赖。CANN FIA Attention 后端原理新增的 CANN 注意力后端位于 cosmos_framework/model/attention/cann/ 子包是 Cosmos3 昇腾适配中最具代表性的算子层工作。后端能力声明checks.py 中的cann_attention_check声明了该后端的支持范围支持数据类型torch.float16、torch.bfloat16支持 GQA/MQAsupports_gqa_mqaTrue不支持 MLAsupports_mlaFalse通过统一的attention_tensor_checks完成张量形状、requires_grad 等一致性校验。融合算子调用与布局转换functions.py 中cann_attention是统一入口优先读取cumulative_seqlen_Q/cumulative_seqlen_KV判断是否为 varlen变长序列打包模式scale缺省时按query.shape[-1] ** -0.5计算return_lse请求时返回(output, None)占位。两条计算路径都落在torch_npu.npu_fused_infer_attention_score上BNSD 标准路径Cosmos3 frontend 传入的是 BSND 布局batch、sequence、num_heads、head_dim而 FIA 标准路径期望 BNSD因此先permute(0, 2, 1, 3)再做融合注意力最后转回 BSND 返回TND 变长路径varlen 模式下 Cosmos3 传入的是 batch1 的 packed BSND 张量先squeeze(0)去掉单例 batch 维再按 FIA 的 TND 布局传递actual_seq_lengths注意 Cosmos 的 cumulative seqlens 含前导 0需切片[1:]只保留累计结束偏移causal 模式使用 2048×2048 的三角掩码_tnd_causal_mask并通过pre_tokens65535、next_tokens65535、sparse_mode3causal或 0 配置稀疏注意力输出按有效长度截断计算后再在序列末尾补零回原始总长度并unsqueeze(0)还原 batch 维。从源码结构看该实现把 Cosmos3 的 BSND 打包张量与昇腾 FIA 的 BNSD/TND 两种输入布局解耦是模型侧无需改动注意力逻辑即可获得融合算子加速的关键。推理验证示例完成适配后可在cosmos-framework根目录下执行推理命令进行基础场景验证。COSMOS_CHECKPOINT用于指定本地权重目录如不设置可直接将命令中的--checkpoint-path替换为实际权重路径。export COSMOS_CHECKPOINT/mnt/workspace/cosmos3/cosmos3-nano export COSMOS_RESOLUTION480 export COSMOS_SEED0 export COSMOS_NPUS1输入 JSON 配置说明推理命令中的-i inputs/omni/t2v.json用于指定单条样例输入。常用字段如下model_mode任务类型例如text2video、image2video、video2videoprompt文本提示词T2V 只需要配置该字段即可vision_pathI2V/V2V 的输入图片或视频路径仅图生视频、视频生视频需要。vision_path可以写远程 URL也可以写本地文件路径。若服务器无法访问远端资源或遇到证书、代理、内网限制等网络问题请先手动下载输入图片/视频到本地然后在 JSON 中改成本地绝对路径例如{ model_mode: image2video, prompt: A robot arm moves smoothly in a lab., vision_path: /mnt/workspace/cosmos3/inputs/robot_153.jpg }T2V 示例 JSON 可简化为{ model_mode: text2video, name: t2v, prompt: A realistic video of molten metal being poured in a steel mill. }T2V 文生视频torchrun --nproc-per-node${COSMOS_NPUS} -m cosmos_framework.scripts.inference \ --parallelism-presetlatency \ -i inputs/omni/t2v.json \ -o outputs/t2v \ --checkpoint-path ${COSMOS_CHECKPOINT} \ --resolution${COSMOS_RESOLUTION} \ --seed${COSMOS_SEED} \ --no-guardrailsI2V 图生视频torchrun --nproc-per-node${COSMOS_NPUS} -m cosmos_framework.scripts.inference \ --parallelism-presetlatency \ -i inputs/omni/i2v.json \ -o outputs/i2v \ --checkpoint-path ${COSMOS_CHECKPOINT} \ --resolution${COSMOS_RESOLUTION} \ --seed${COSMOS_SEED} \ --no-guardrailsV2V 视频生视频torchrun --nproc-per-node${COSMOS_NPUS} -m cosmos_framework.scripts.inference \ --parallelism-presetlatency \ -i inputs/omni/v2v.json \ -o outputs/v2v \ --checkpoint-path ${COSMOS_CHECKPOINT} \ --resolution${COSMOS_RESOLUTION} \ --seed${COSMOS_SEED} \ --no-guardrails命令要点说明torchrun --nproc-per-node${COSMOS_NPUS}以多进程方式拉起分布式推理进程数等于 NPU 数量--parallelism-presetlatency单卡低延迟预设适合功能验证多卡并行场景改用throughput预设--no-guardrails关闭文本/人脸安全护栏对应 pyproject 中的guardrail可选依赖组避免额外依赖与推理开销基础场景默认单卡即可完成此时COSMOS_NPUS1。多卡并行当前适配支持 CPContext Parallel、CFGPClassifier-Free Guidance Parallel和 FSDP 三种多卡推理方式。运行前请将COSMOS_NPUS设置为实际使用的 NPU 数量并保证各并行度与进程数匹配。并行方式主要参数适用目的与约束CPContext Parallel--cp-size沿 token 序列切分 Attention 计算适合长序列并降低激活显存当前 CP 范围为 132CFGPClassifier-Free Guidance Parallel--cfgp-size将有条件与无条件 CFG 分支分配到不同设备CFGP 仅支持 1 或 2更大规模可与 CP/FSDP 组合FSDP--dp-shard-size按进程数切分模型参数优先降低单卡权重显存FSDP 的 DP 通信组与 CP/CFGP 通信组相互独立通信域大小满足dp-shard-size × dp-replicate-size WORLD_SIZE WORLD_SIZE % (cp-size × cfgp-size) 0CP 上下文并行CP 将长序列的 Attention 计算切分到多卡上并行执行适合超长视频/长上下文生成torchrun --nproc-per-node${COSMOS_NPUS} -m cosmos_framework.scripts.inference \ --parallelism-presetthroughput \ --dp-shard-size1 --cp-size${COSMOS_NPUS} --cfgp-size1 \ -i inputs/omni/t2v.json \ -o outputs/t2v_cp \ --checkpoint-path ${COSMOS_CHECKPOINT} \ --resolution${COSMOS_RESOLUTION} \ --seed${COSMOS_SEED} \ --no-guardrailsCFGP 无条件引导并行CFGClassifier-Free Guidance需要同时计算有条件和无条件两个去噪分支CFGP 将两个分支分配到不同设备组并行执行。单独启用时设置COSMOS_NPUS2torchrun --nproc-per-node${COSMOS_NPUS} -m cosmos_framework.scripts.inference \ --parallelism-presetthroughput \ --dp-shard-size1 --cp-size1 --cfgp-size2 \ -i inputs/omni/t2v.json \ -o outputs/t2v_cfgp \ --checkpoint-path ${COSMOS_CHECKPOINT} \ --resolution${COSMOS_RESOLUTION} \ --seed${COSMOS_SEED} \ --no-guardrailsFSDP 参数分片FSDP 按进程数切分模型参数将权重/优化器状态分布到多卡优先降低单卡权重显存占用torchrun --nproc-per-node${COSMOS_NPUS} -m cosmos_framework.scripts.inference \ --parallelism-presetthroughput \ --dp-shard-size${COSMOS_NPUS} --cp-size1 --cfgp-size1 \ -i inputs/omni/t2v.json \ -o outputs/t2v_fsdp \ --checkpoint-path ${COSMOS_CHECKPOINT} \ --resolution${COSMOS_RESOLUTION} \ --seed${COSMOS_SEED} \ --no-guardrails多卡并行之所以能开箱即用与上文adapt_distributed_runtime的 HCCL 化改造直接相关CP/CFGP 通信组与 FSDP 的 DP 通信组均建立在hccl后端之上torch.npu.set_device(local_rank)保证每个 rank 张量落在自己的 NPU 上。同类思路也可参考本仓库 Cosmos 系列多卡并行优化说明其中对 CFG 并行、上下文并行与融合算子Flash Attention、RMSNorm、Rotary的改造与本样例互为印证。样例输出与效果验证在昇腾 NPU 上运行上述基础场景后生成结果会输出到-o指定的目录如outputs/t2v。原仓库 README 的样例输出展示章节提供了 T2V 文生视频、I2V 图生视频、V2V 视频生视频三类场景的示例生成结果可用于快速对照查看生成效果场景说明T2V 文生视频纯文本提示词驱动的视频生成结果I2V 图生视频以单张图片为条件输入的视频生成结果V2V 视频生视频以视频片段为输入的续写/变换结果建议验证时按以下顺序逐步推进先以COSMOS_NPUS1跑通 T2V 单卡链路确认权重加载、VAE 推理与视频落盘正常再依次验证 I2V检查vision_path本地路径可读与 V2V最后切换到throughput预设并按上节公式配置并行度逐步扩展到多卡场景。总结与引用本样例以适配文件覆盖 脚本化源码改写 新增 CANN 后端三层结构将 Cosmos3-Nano 世界基础模型的 T2V/I2V/V2V 推理完整迁移到昇腾 A3 NPU 平台环境侧锁定 CANN 9.0.0 / torch_npu 2.10.0 / Python 3.13设备侧完成 CUDA→NPU 与 NCCL→HCCL 的全面替换算子侧通过npu_fused_infer_attention_score提供 FIA attention 加速推理侧通过本地 checkpoint 补丁实现完全离线运行并行侧支持 CP/CFGP/FSDP 的灵活组合。读者可依据本文从零搭建环境并完成多场景验证亦可深入 npu_adapt.sh 与 cann attention 实现 理解每个适配点的源码级细节。如需引用 Cosmos3 模型本身可参考上游论文《Cosmos 3: Omnimodal World Models for Physical AI》NVIDIA 团队2026 年其 BibTeX 引用信息收录于原 README 的 citation 章节。【免费下载链接】cann-recipes-embodied-ai本项目针对具身智能业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-ai创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
