企业官网建设流程全解析

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_npu==2.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_npu=2.10.0、Python=3.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.sh

uv 环境管理工具安装

本样例使用 uv 管理 Python 依赖，其版本要求与索引配置已固化在 pyproject.toml 的[tool.uv]段（required-version >= 0.11.3）。如果当前环境已经安装 uv，可跳过此步：

wget -qO- https://astral.sh/uv/install.sh | sh

Python 运行环境安装

在cosmos-framework根目录下执行uv sync，指定 Python 3.13：

cd cosmos-framework uv sync --python 3.13

从依赖清单可以进一步理解适配的关键点（见 pyproject.toml）：

• 核心依赖直接声明torch_npu==2.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（安全护栏相关）、serve（gradio/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.pth

Wan2.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 时回退 64GB；
• adapt_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.py：

• import pynvml改为 try/except 可选导入（NPU 环境通常无 NVML）；
• torch.cuda.set_device(local_rank)→torch.npu.set_device(local_rank)；
• dist.init_process_group(backend="nccl"→backend="hccl"，即用昇腾 HCCL 通信库替代 NCCL；
• torch.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。该补丁的核心改动包括：

1. VLM tokenizer：当Path(checkpoint_path, "text_tokenizer").is_dir()成立时，直接以 checkpoint 目录作为tokenizer_type构建 processor，避免从远端仓库下载 tokenizer；
2. Wan2.2 VAE：检测 checkpoint 目录下的Wan2.2_VAE.pth，存在则改写 VAE 配置指向本地文件（bucket_name=""、vae_path=本地路径）；
3. AVAE 音频 tokenizer：from_checkpoint默认值由 False 改为 True，优先使用 checkpoint 内捆绑的sound_tokenizer/，保证离线本地推理自包含。

pynvml 可选化

adapt_optional_nvml处理cosmos_framework/utils/device.py：将import pynvml包上 try/except，except 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/MQA（supports_gqa_mqa=True），不支持 MLA（supports_mla=False）；
• 通过统一的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 传入的是 batch=1 的 packed BSND 张量，先squeeze(0)去掉单例 batch 维，再按 FIA 的 TND 布局传递actual_seq_lengths（注意 Cosmos 的 cumulative seqlens 含前导 0，需切片[1:]只保留累计结束偏移）；causal 模式使用 2048×2048 的三角掩码（_tnd_causal_mask），并通过pre_tokens=65535、next_tokens=65535、sparse_mode=3（causal）或 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_RESOLUTION=480 export COSMOS_SEED=0 export COSMOS_NPUS=1

输入 JSON 配置说明

推理命令中的-i inputs/omni/t2v.json用于指定单条样例输入。常用字段如下：

• model_mode：任务类型，例如text2video、image2video、video2video；
• prompt：文本提示词，T2V 只需要配置该字段即可；
• vision_path：I2V/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-preset=latency \ -i inputs/omni/t2v.json \ -o outputs/t2v \ --checkpoint-path ${COSMOS_CHECKPOINT} \ --resolution=${COSMOS_RESOLUTION} \ --seed=${COSMOS_SEED} \ --no-guardrails

I2V 图生视频

torchrun --nproc-per-node=${COSMOS_NPUS} -m cosmos_framework.scripts.inference \ --parallelism-preset=latency \ -i inputs/omni/i2v.json \ -o outputs/i2v \ --checkpoint-path ${COSMOS_CHECKPOINT} \ --resolution=${COSMOS_RESOLUTION} \ --seed=${COSMOS_SEED} \ --no-guardrails

V2V 视频生视频

torchrun --nproc-per-node=${COSMOS_NPUS} -m cosmos_framework.scripts.inference \ --parallelism-preset=latency \ -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-preset=latency：单卡低延迟预设，适合功能验证；多卡并行场景改用throughput预设；
• --no-guardrails：关闭文本/人脸安全护栏（对应 pyproject 中的guardrail可选依赖组），避免额外依赖与推理开销；
• 基础场景默认单卡即可完成，此时COSMOS_NPUS=1。

多卡并行

当前适配支持 CP（Context Parallel）、CFGP（Classifier-Free Guidance Parallel）和 FSDP 三种多卡推理方式。运行前请将COSMOS_NPUS设置为实际使用的 NPU 数量，并保证各并行度与进程数匹配。

FSDP 的 DP 通信组与 CP/CFGP 通信组相互独立，通信域大小满足：

dp-shard-size × dp-replicate-size = WORLD_SIZE WORLD_SIZE % (cp-size × cfgp-size) = 0

CP 上下文并行

CP 将长序列的 Attention 计算切分到多卡上并行执行，适合超长视频/长上下文生成：

torchrun --nproc-per-node=${COSMOS_NPUS} -m cosmos_framework.scripts.inference \ --parallelism-preset=throughput \ --dp-shard-size=1 --cp-size=${COSMOS_NPUS} --cfgp-size=1 \ -i inputs/omni/t2v.json \ -o outputs/t2v_cp \ --checkpoint-path ${COSMOS_CHECKPOINT} \ --resolution=${COSMOS_RESOLUTION} \ --seed=${COSMOS_SEED} \ --no-guardrails

CFGP 无条件引导并行

CFG（Classifier-Free Guidance）需要同时计算有条件和无条件两个去噪分支，CFGP 将两个分支分配到不同设备组并行执行。单独启用时设置COSMOS_NPUS=2：

torchrun --nproc-per-node=${COSMOS_NPUS} -m cosmos_framework.scripts.inference \ --parallelism-preset=throughput \ --dp-shard-size=1 --cp-size=1 --cfgp-size=2 \ -i inputs/omni/t2v.json \ -o outputs/t2v_cfgp \ --checkpoint-path ${COSMOS_CHECKPOINT} \ --resolution=${COSMOS_RESOLUTION} \ --seed=${COSMOS_SEED} \ --no-guardrails

FSDP 参数分片

FSDP 按进程数切分模型参数，将权重/优化器状态分布到多卡，优先降低单卡权重显存占用：

torchrun --nproc-per-node=${COSMOS_NPUS} -m cosmos_framework.scripts.inference \ --parallelism-preset=throughput \ --dp-shard-size=${COSMOS_NPUS} --cp-size=1 --cfgp-size=1 \ -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 视频生视频三类场景的示例生成结果，可用于快速对照查看生成效果：

建议验证时按以下顺序逐步推进：先以COSMOS_NPUS=1跑通 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