cosmos3-post-training

用途

Cosmos3 Post-Training (SFT) 技能用于指导用户对 Cosmos3-Nano（或通过 LoRA 的 Cosmos3-Super）在示例 Bridge 视频数据集或自定义视频数据集上执行监督微调（SFT）。当用户询问 recipe TOML 中哪些字段需要覆盖（如 [model.parallelism].data_parallel_shard_degree、[dataloader_train].max_samples_per_batch、[optimizer].lr、[trainer].max_iter、[checkpoint].load_path 等）、应选择哪个实验 SKU、或将基础 Hugging Face 检查点转换为 DCP、或将训练后的 DCP 转换回 safetensors 时，均适用本技能。安装、--group=cu130-train / cu128-train、LD_LIBRARY_PATH 问题转交 cosmos3-setup；推理参数、并行预设、在线服务转交 cosmos3-inference；原始视频字幕生成或组装 SFT JSONL 见 docs/dataset_jsonl.md。

路径约定

所有路径相对于 cosmos3 包根目录（即本技能文件的 ../../../）。所有 uv run / python / torchrun / bash 命令也应从该目录运行。

工作流程概览

安装

安装训练扩展：uv sync --all-extras --group=cu130-train（旧驱动用 cu128-train），然后 source .venv/bin/activate && export LD_LIBRARY_PATH=.。

Step 1 — 准备数据与配置

对要运行的 recipe，将 HF 数据集下载到 examples/data/<dataset>/，将 Wan2.2 VAE 通过 uvx hf@latest download … 下载到 examples/checkpoints/wan22_vae/Wan2.2_VAE.pth。Reasoner recipe 在启动时从 HF Hub 流式读取数据集，无需 Step 1 下载。

Step 2 — 准备检查点

设置 BASE_CHECKPOINT_NAME（Cosmos3-Nano 或 Cosmos3-Super，与 recipe 匹配），并运行 python -m cosmos_framework.scripts.convert_model_to_dcp -o examples/checkpoints/$BASE_CHECKPOINT_NAME --checkpoint-path $BASE_CHECKPOINT_NAME。Reasoner recipe 可跳过（Qwen3-VL 骨干在启动时从 HF Hub 获取）。

Step 3 — 运行训练（方式 A，推荐）

在仓库根目录执行 bash examples/launch_sft_<recipe>.sh（如 launch_sft_vision_nano.sh）。启动脚本从 Steps 1+2 填充的默认 examples/ 位置解析 DATASET_PATH、BASE_CHECKPOINT_PATH、WAN_VAE_PATH；先在 shell 中 export 任一变量即可覆盖。

Step 3 — 运行训练（方式 B，原始 torchrun）

自行导出环境变量，然后执行 IMAGINAIRE_OUTPUT_ROOT=outputs/train PYTHONPATH=. torchrun --nproc_per_node=8 -m cosmos_framework.scripts.train --sft-toml=examples/toml/sft_config/<recipe>.toml。与方式 A 不同，原始 torchrun 不会从 examples/ 自动解析这三个环境变量，必须来自 shell，或手动编辑 TOML 内联字面路径。

输出

$RUN_DIR = $IMAGINAIRE_OUTPUT_ROOT/<job.project>/<job.group>/<job.name>。DCP 检查点位于 $RUN_DIR/checkpoints/iter_<N>/；最新 iter 名称在 $RUN_DIR/checkpoints/latest_checkpoint.txt。检查点旁的 $RUN_DIR/config.yaml 是推理时消费的配置。

推理

将 cosmos_framework.scripts.inference 指向 $RUN_DIR/checkpoints/iter_<N> 并结合 --config-file $RUN_DIR/config.yaml（预设与输入格式见 cosmos3-inference 技能）。

导出（可选）

python -m cosmos_framework.scripts.export_model --checkpoint-path $RUN_DIR/checkpoints/$(cat $RUN_DIR/checkpoints/latest_checkpoint.txt) --config-file $RUN_DIR/config.yaml -o $RUN_DIR/model 将可移植的 HF safetensors 检查点写入 $RUN_DIR/model。

文档中不明显的要点

训练扩展是独立 group：SFT 需要 cu130-train / cu128-train 安装组，而非仅推理的 cu130 / cu128。用错 group 重新 uv sync 会静默地不安装训练依赖。
Recipe = 配对的 examples/launch_sft_<r>.sh + examples/toml/sft_config/<r>.toml：.sh 直接声明 TOML_FILE（完整仓库相对路径）以及 ${DATASET_PATH:=…} / ${BASE_CHECKPOINT_PATH:=…} 默认值，与 Steps 1+2 落地位置对齐，然后 source examples/_sft_launcher_common.sh。用户 shell 中 export 的值优先于默认值。helper 转发到 cosmos_framework.scripts.train --sft-toml=$TOML_FILE，TAIL_OVERRIDES bash 数组项以 Hydra 风格的 key.path=value 追加在 -- 之后（最后应用在 cosmos_framework/configs/toml_config/sft_config.py 中经 pydantic 校验的 TOML schema 之上）。MASTER_PORT 在 helper 中默认 50012；仅在同一节点协同启动多任务时在 launcher 中设置（或 export）。
方式 A（配对启动 shell）vs 方式 B（原始 torchrun）：方式 A 从 examples/ 默认解析三个环境变量（DATASET_PATH / BASE_CHECKPOINT_PATH / WAN_VAE_PATH），未设置也能开箱运行。方式 B 需要自行 export 这些环境变量，或手工编辑 TOML 内联路径；TOML 使用 ${ENV:DATASET_PATH} 插值，在 TOML 加载时解析。
IMAGINAIRE_OUTPUT_ROOT 控制整个输出树：设置 IMAGINAIRE_OUTPUT_ROOT=outputs/train 使所有内容落到 outputs/train/<job.project>/<job.group>/<job.name>/（日志、config.yaml、checkpoints/iter_<N>、回调输出）。未设置时训练回退到 /tmp/imaginaire4-output/...。
W&B 默认禁用：每个 recipe TOML 设置 [job].wandb_mode = "disabled"。要记录到 W&B，在 TOML 中改为 "online" 并在启动前 export WANDB_API_KEY。
推理直接使用 DCP 检查点：标准流程将 cosmos_framework.scripts.inference 指向 $RUN_DIR/checkpoints/iter_<N> 并结合 --config-file $RUN_DIR/config.yaml。Hugging Face safetensors 导出（$RUN_DIR/model）是可选的，仅在需要可移植单文件检查点时使用。
并行度必须与拓扑匹配：TOML [model.parallelism] 块中 data_parallel_shard_degree × data_parallel_replicate_degree × context_parallel_shard_degree 必须等于 WORLD_SIZE。-1 会从 torchrun world size 自动选择 data_parallel_shard_degree。不匹配会导致 FSDP 初始化失败。
--dryrun：cosmos_framework.scripts.train 接受 --dryrun 以在不启动训练的情况下端到端校验配置。迭代 TOML 键或 Hydra 覆盖时使用。

技能介绍

用途