快速入门：Cosmos3-Nano模型训练

快速入门：Cosmos3-Nano模型训练#

本文档介绍如何在 LoongForge 框架下快速启动 Cosmos3-Nano 的 SFT（监督微调）训练。示例: DROID Action-Policy SFT，采用 FSDP2 全分片、bf16 训练，与官方 cosmos-framework 的 launch_sft_action_policy_droid 对齐,对应的开源框架训练文档：cosmos-framework action_policy_droid_posttrain.md。

0. 资源准备#

Cosmos3-Nano 是一个视觉-语言联合建模的动作策略模型，训练依赖三类外部资源：主干权重（Cosmos3-Nano + Wan2.2 VAE 编码器）、Qwen3-VL tokenizer / processor，以及 DROID 后训练数据集。下面分别说明。

0.1 模型权重#

主权重来自 HuggingFace 上的 nvidia/Cosmos3-Nano。Cosmos3 训练流程使用 DCP（Distributed Checkpoint）格式，因此下载得到的原始 HF safetensors 权重需要先做一次离线转换，转换步骤参考官方框架文档 training.md Step 2。转换后的 DCP 权重目录通过 --pretrained-checkpoint 加载：

--pretrained-checkpoint $CHECKPOINT_PATH   # 转换后的 DCP 权重目录
--init-on-meta                             # 权重延迟到 FSDP wrap 之后再加载，降低初始化峰值显存

除主权重外，视频分支还需要一个 VAE 编码器，Cosmos3-Nano 复用 Wan2.2-TI2V-5B 中的 Wan2.2_VAE.pth，VAE 路径写在 configs/models/embodied/cosmos3/nano.yaml 的 vae_path 字段，训练脚本无需再显式指定：

0.2 Tokenizer#

Cosmos3-Nano 的语言路径基于 Qwen3-VL-8B-Instruct，tokenizer / processor 直接复用其 HuggingFace 目录，通过 --tokenizer-path 加载：

hf download Qwen/Qwen3-VL-8B-Instruct --local-dir /workspace/ckpt/Qwen3-VL-8B-Instruct
export TOKENIZER_PATH=/workspace/ckpt/Qwen3-VL-8B-Instruct

0.3 数据集#

示例使用 NVIDIA 官方发布的 DROID 后训练子集 nvidia/Cosmos3-DROID（success/ 分支，LeRobot 格式），下载到本地：

hf download nvidia/Cosmos3-DROID --repo-type dataset --local-dir /workspace/data/Cosmos3-DROID
export DATASET_PATH=/workspace/data/Cosmos3-DROID/success

1. 数据配置#

DROID 数据不需要额外的离线预处理，训练时通过内建的 cosmos3_droid 处理策略在线完成画面拼接、图像增强、动作对齐等步骤。启用方式是两个成对参数：

--dataset-format lerobot_datasets    # 数据集格式：LeRobot
--dataset-strategy cosmos3_droid     # Cosmos3 官方 DROID 数据处理策略

其余目标分辨率、动作 chunk 长度、CFG dropout 等已经写在 configs/models/embodied/cosmos3/nano.yaml 的 data 段（默认 target_h/target_w=480、action_chunk_length=32、action_fps=15.0），通常不需要覆盖。

2. 启动训练#

启动脚本：examples/embodied/cosmos3/run_cosmos3_nano_droid_fsdp.sh。默认为单机 8 卡 FSDP2 + bf16，训练 500 步、每卡 batch=2。

2.1 环境变量#

先统一设置路径：

cd /workspace/LoongForge

export LOONGFORGE_PATH=/workspace/LoongForge
export TOKENIZER_PATH=/workspace/ckpt/Qwen3-VL-8B-Instruct
export CHECKPOINT_PATH=/workspace/ckpt/Cosmos3-Nano-DCP   # 转换后的 DCP 权重目录
export DATASET_PATH=/workspace/data/Cosmos3-DROID/success
export OUTPUT_DIR=/workspace/outputs/cosmos3_nano_droid

2.2 启动脚本#

单机 8 卡 FSDP2 SFT：

bash examples/embodied/cosmos3/run_cosmos3_nano_droid_fsdp.sh

2.3 关键参数说明#

脚本内的参数按用途分组如下：

模型与分布式：

--model-name cosmos3_nano            # 通过 config_map 映射到 Cosmos3-Nano DROID 配方
--distributed-strategy fsdp          # 分布式策略：FSDP2 全分片
--dtype bfloat16                     # 训练精度：bf16
--init-on-meta                       # 通过 meta device 降低模型初始化时的峰值显存

数据：

--dataset-format lerobot_datasets    # 数据集格式：LeRobot
--dataset-strategy cosmos3_droid     # 官方 DROID 数据处理策略（画面拼接、图像增强等）
--dataset-path $DATASET_PATH         # DROID 数据目录
--tokenizer-path $TOKENIZER_PATH     # Qwen3-VL tokenizer / processor 目录
--num-workers 4                      # DataLoader worker 数

训练与优化器：

--trainer-type FinetuneTrainer
--train-iters 500                    # 训练步数
--per-device-batch-size 2            # 每卡 batch
--gradient-accumulation-steps 1
--disable-tf32                       # 关闭 TF32，保持与官方参考实现一致的数值精度
--pretrained-checkpoint $CHECKPOINT_PATH
--save-interval 100
--seed 42

分组学习率与优化器：

Cosmos3-Nano 的动作头（action2llm / llm2action / action_modality_embed）需要比视觉-语言主干更快的学习率，脚本用 --lr-group 按参数名前缀分组：

--lr-group net.action2llm=1e-3,net.llm2action=1e-3,net.action_modality_embed=1e-3,net=2e-4
                                     # 动作头 1e-3，其余 net 2e-4
--lr-decay-style lambda_linear       # 学习率衰减：线性
--lr-warmup-iters 0
--optimizer TorchFusedAdamW          # 优化器：fused AdamW
--clip-grad 1.0
--weight-decay 0.05
--adam-beta1 0.9
--adam-beta2 0.99
--adam-eps 1e-8

如需按实际训练规模调整，常用做法是：

  • 覆盖 --train-iters、--save-interval 控制训练时长与 checkpoint 频率

  • 覆盖 --per-device-batch-size 与 --gradient-accumulation-steps 调整 global batch

  • 覆盖 --lr-group 中 net= 项微调主干学习率