NVIDIA Cosmos Tokenizer 基于 NeMo Framework 的后训练（Post-training）实战指南-二趣网

NVIDIA Cosmos Tokenizer 基于 NeMo Framework 的后训练（Post-training）实战指南

【免费下载链接】cosmosNVIDIA Cosmos is an open platform of world models, datasets, and tools that enables developers to build Physical AI for robots, autonomous vehicles, smart infrastructure, and more.项目地址: https://gitcode.com/GitHub_Trending/cosmos7/cosmos

Cosmos Tokenizer 是 NVIDIA Cosmos 世界基础模型平台中的视觉分词器组件，负责把图像/视频压缩为连续潜变量（Continuous Latent）或离散 token（Discrete Token）。本文以 cosmos1/models/tokenizer/nemo/README.md 为骨架，完整讲解如何借助 NVIDIA NeMo Framework 对 Cosmos Tokenizer 进行后训练（Post-training），使预训练模型更好地建模自驾驶等场景中此前未见过的视频数据分布。读完本文，你将掌握：模型与后训练任务支持矩阵、从数据集组织到容器启动再到torchrun拉起训练的全流程、model/data/trainer/optim四大配置组件的底层默认参数，以及如何通过 Weights & Biases（wandb）监控 loss 收敛。

Cosmos Tokenizer 后训练：解决什么问题

Cosmos Tokenizer 在预训练阶段学习的是一般性的视频压缩规律，而 Physical AI 业务（尤其是自动驾驶）的客户视频数据往往带有独特的视觉与时间模式——例如特定传感器视角、特定路况、特定光照条件，这些模式在预训练语料中可能覆盖不足。对 Tokenizer 做后训练，本质上是让分词器适配你自有数据的特有分布，从而更精确地压缩与重建目标场景的视频。

这一点至关重要，因为下游的扩散模型（diffusion model）直接消费 Tokenizer 输出的 latent 或 token：只有当分词器忠实保留场景的物理细节（物体形状、运动轨迹、遮挡关系）时，扩散模型才可能生成逼真的物理场景，最终提升自动驾驶系统的性能与安全性。从源码结构看，该职责由 cosmos1/models/tokenizer/nemo/train_tokenizer.py 中的TokenizerModel承载，它与仓库内原生 PyTorch 的 continuous_video.py / discrete_video.py 网络定义同源，因此后训练得到的权重与推理/下游扩散模型链路天然兼容。

模型支持矩阵：当前可后训练的两种 Tokenizer

NeMo Framework 目前支持对以下两种 Cosmos Tokenizer 模型进行后训练，可从 Hugging Face 下载对应 checkpoint（在仓库 tokenizer README 的“Download Pre-trained Checkpoints”一节中给出了snapshot_download方式的批量下载脚本，模型目录下提供encoder.jit、decoder.jit、autoencoder.jit三份 JIT 文件）：

模型名称	压缩特性	说明
Cosmos-1.0-Tokenizer-CV8x8x8	连续视频分词器，时间 8x、空间 8x8，总压缩率 512x	输出 16 通道连续 latent
Cosmos-1.0-Tokenizer-DV8x16x16	离散视频分词器，时间 8x、空间 16x16，总压缩率 2048x	FSQ 量化，输出 6 通道离散 code

模型命名中的三个数字依次表示 T×H×W 方向的压缩倍数。在 configs.py 中可以找到两种模型的默认结构参数：continuous_video设置latent_channels=16、spatial_compression=8、temporal_compression=8；discrete_video则设置spatial_compression=16、temporal_compression=8，并使用 FSQ 量化器（levels=[8, 8, 8, 5, 5, 5]、embedding_dim=6）。关于空间与时间压缩因子的组合能力（8x 或 16x 空间、4x 或 8x 时间），可进一步阅读 tokenizer README。

性能提示：为获得最佳后训练性能，官方推荐使用 H100-80GB 或 A100-80GB 级别的 GPU。

后训练任务支持矩阵

当前版本对 Cosmos Tokenizer 的后训练支持情况如下：

后训练任务	支持状态
通用后训练与验证（General post-training and validation）	Supported（已支持）

这与 POST_TRAINING.md 中描述的 Cosmos 整体后训练规划一致：首版发布优先提供“通用后训练”脚本（该文档中扩散 WFM 与自回归 WFM 的通用后训练均已支持），指令控制、动作控制、相机控制、多视角生成等任务标记为 “Coming soon”。对于 Tokenizer 而言，当前可直接投入使用的即是面向自定义视频分布的通用后训练。

环境准备：硬件、容器与凭据

硬件与系统要求

NVIDIA GPU 与驱动：确保可访问 80GB 显存的 H100 或 A100，以满足训练显存需求。
容器化平台：官方推荐使用 NVIDIA NeMo Docker Runtime（也可选用 NVIDIA enroot），NeMo Framework 容器同时支持后训练与推理。
Hugging Face User Access Token：下载 Cosmos 预训练模型与训练所必需。
Weights and Biases API Key：用于训练日志与指标追踪。

克隆仓库与启动 NeMo 容器

将本仓库克隆到本地（如需从远端获取源码，可使用git clone https://gitcode.com/GitHub_Trending/cosmos7/cosmos.git），然后启动 NeMo 容器并把仓库挂载进/workspace/Cosmos：

docker run --ipc=host -it --gpus=all \ -v $PATH_TO_COSMOS_REPO:/workspace/Cosmos \ nvcr.io/nvidia/nemo:24.12.01 bash

其中$PATH_TO_COSMOS_REPO替换为本地仓库绝对路径。--ipc=host用于共享主机 IPC 命名空间（避免多进程 dataloader 的共享内存问题），--gpus=all暴露全部 GPU。容器启动后，后续所有命令均在容器内执行。

下载预训练 Checkpoint

按照上文“模型支持矩阵”中的链接，从 Hugging Face 下载Cosmos-1.0-Tokenizer-CV8x8x8（或Cosmos-1.0-Tokenizer-DV8x16x16）的 checkpoint 到本地目录，并记录该目录路径，稍后作为model.jit_ckpt_pth传入。仓库 tokenizer README 中给出了完整的批量下载示例代码，下载后目录结构形如checkpoints/Cosmos-1.0-Tokenizer-CV8x8x8/{encoder.jit, decoder.jit, autoencoder.jit}。

后训练三步走

后训练一个 Cosmos Tokenizer，使其更擅长压缩你的 Physical AI 场景视频，共分三个步骤：准备数据集 → 预处理数据 → 后训练模型。其中数据预处理环节由训练脚本内置的 TaskEncoder 自动完成（见下文“第二步”）。

第一步：准备数据集

将数据组织为包含多个视频 tar 分片（shard）的文件夹，每个 tar 内含 MP4 格式视频（建议分辨率至少 720p）。推荐目录结构如下：

000000.tar ├── 1.mp4 └── 2.mp4 000001.tar ├── 3.mp4 └── 4.mp4

000000.tar、000001.tar分别代表独立分片，可按需添加更多分片。该结构对应训练脚本中data.path指向的根目录。

第二步：数据预处理（ImageTaskEncoder 自动管线）

数据读取与预处理由 train_tokenizer.py 中的ImageTaskEncoder（继承自 megatron-energon 的DefaultTaskEncoder）完成，运行时无需单独编写预处理代码。从源码可以看到其encode_sample的处理链：

帧截取：sample.image.frames[:33, :, :256, :256]——取前 33 帧、裁剪到 256×256 空间尺寸；
维度重排：rearrange(frames, "t c h w -> c t h w")——转为C×T×H×W布局；
空间随机裁剪：RandomCrop，裁剪尺寸由get_crop_size_info(crop_height)依据crop_height=256计算得到；
时间随机裁剪：TemporalRandomCrop(temporal_crop=49)，对时间维度做随机窗口采样；
归一化：Normalize(mean=0.5, std=0.5)——将像素值线性映射到[-1, 1]区间；
精度转换：转为torch.bfloat16。

每个样本最终组织为{VIDEO_KEY: frames, MASK_KEY: mask, "aspect_ratio": "1,1"}结构，其中 mask 为与帧同形的全 1 张量（训练脚本中视频帧恒有效）。这套管线同时定义了模型训练时的输入张量规格[B, C, T, H, W] = [B, 3, 33, 256, 256]，与后文FakeDataset中生成的数据形状一致。

第三步：安装依赖并启动后训练

1. 安装依赖（requirements.txt）：

pip install megatron-energon==4.0.0 pyav pip install git+https://github.com/NVIDIA/NeMo-Run.git pip install moviepy==1.0.3 imageio

其中megatron-energon负责 tar 分片数据集的读取与任务编码，pyav用于视频解码，NeMo-Run提供 CLI 参数化训练入口，moviepy与imageio用于视频处理。

2. 以 8 卡为例启动后训练Cosmos-1.0-Tokenizer-CV8x8x8：

export CKPT_PTH="<path/to/your/HF/checkpoints/folder>" export DATA="<path/to/your/data>" # 可选：用 Weights and Biases 监控训练进度 export WANDB_API_KEY="</your/wandb/api/key>" export WANDB_PROJECT_NAME="cosmos-diffusion-nemo-post-training" export WANDB_RUN_ID="cosmos_diffusion_7b_text2world" torchrun --nproc-per-node 8 cosmos1/models/tokenizer/nemo/train_tokenizer.py --yes \ data.path=$DATA \ model.jit_ckpt_pth=$CKPT_PTH \ model.model="Cosmos-1.0-Tokenizer-CV8x8x8"

说明：

--nproc-per-node 8表示每个节点使用 8 个 GPU 进程，可与torchrun的其他参数（如--nnodes）组合扩展到多节点；
--yes跳过配置确认交互，直接按给定覆盖项启动；
data.path指向第一步中 tar 分片所在目录；
model.jit_ckpt_pth指向 Hugging Face 下载的 checkpoint 目录；
model.model指定要后训练的模型名称，可替换为"Cosmos-1.0-Tokenizer-DV8x16x16"以训练离散分词器；
若未设置WANDB_API_KEY，脚本会自动退化为不启用 wandb 日志（源码中logger=WandbLogger(...) if "WANDB_API_KEY" in os.environ else None）。

四大配置组件与可配置超参数

train_tokenizer.py通过@run.cli.factory(target=llm.train)暴露一个参数化训练入口，支持四大配置组件，任何超参数都可通过命令行覆盖（形如组件.参数=值）。下表汇总了各组件的核心配置项及其源码默认值，便于你按需调整：

model：选择模型与指定 Checkpoint

参数	默认值	说明
`model.model`	`Cosmos-1.0-Tokenizer-CV8x8x8`	后训练目标模型（可切换为 DV 系列）
`model.jit_ckpt_pth`	`None`（须由命令行传入）	预训练 JIT checkpoint 路径，训练从该权重初始化

data：数据加载与 dataloader 参数

参数	默认值	说明
`data.path`	`None`（须由命令行传入）	视频 tar 分片数据集根目录
`data.global_batch_size`	`8`	全局 batch size（跨所有 GPU 累积）
`data.micro_batch_size`	`1`	单卡微批大小
`data.num_workers`	`1`	dataloader 工作进程数
`data.task_encoder`	`ImageTaskEncoder`	数据预处理管线（见上文“第二步”）

源码中DiffusionDataModule同时承担训练与验证数据加载，并通过MegatronDataSampler依据seq_len/micro_batch_size/global_batch_size自动推导梯度累积步数。

trainer：训练循环

参数	默认值	说明
`trainer.devices`	`auto`	使用的设备数量，通常由`torchrun`自动推导
`trainer.num_nodes`	环境变量`SLURM_NNODES`（默认 1）	节点数，适配 SLURM 集群
`trainer.accelerator`	`gpu`	加速器类型
`trainer.strategy`	`ddp_find_unused_parameters_true`	DDP 策略，兼容含未使用参数的网络
`trainer.max_epochs`	`10000`	最大训练轮数（配合按步保存，实际以步数收敛为准）
`trainer.precision`	`bf16`	混合精度，与数据管线中的 bfloat16 一致
`trainer.val_check_interval`	`100`	每 100 步做一次验证
`trainer.limit_val_batches`	`1`	验证仅取 1 个 batch
`trainer.log_every_n_steps`	`1`	每步记录日志
`trainer.use_distributed_sampler`	`False`	禁用分布式采样器，由`MegatronDataSampler`接管

此外，trainer 挂载了三个回调：ModelCheckpoint（以global_step为监控指标，每 100 训练步保存一次，save_top_k=3保留最新 3 份）、PreemptionCallback（应对集群抢占）与TimingCallback（计时统计）。

optim：优化器参数

参数	默认值	说明
`optim.optimizer_fn`	`torch.optim.AdamW`	优化器类型
`lr`	`1e-4`	学习率
`betas`	`(0.5, 0.999)`	Adam 动量超参数（beta1 取 0.5 而非常见 0.9，适配生成式训练）
`eps`	`1e-8`	数值稳定项
`weight_decay`	`0.01`	权重衰减
`fused`	`True`	启用 fused Adam 内核加速

命令行覆盖示例：

model.jit_ckpt_pth=<your/desired/path> trainer.max_epochs=<your/desired/epochs> data.global_batch_size=<your/desired/batch_size>

除上述四项外，train_tokenizer.py还配置了resume（AutoResume：存在 checkpoint 时自动续训、resume_past_end=True允许越过已结束的 epoch 继续）。按需调整这些参数后重新运行训练命令即可。

训练过程监控与收敛判断

训练启动后，wandb 会记录 loss 曲线。根据训练脚本的默认配置，日志每步写入一次。启动若干百次迭代后，应当观察到 wandb 中loss开始下降。以官方后训练示例（见 assets/loss.png）为参考，横轴为trainer/global_step（0 至 500 步），纵轴为loss（约 0.05～0.1 区间）：训练初期 loss 从接近 0.1 快速下降至约 0.06～0.07，随后在该区间内小幅波动并整体保持缓降、趋于稳定——这一模式可作为判断训练是否正常收敛的经验性参考。

结合源码理解后训练流程

为便于排查问题与二次开发，这里补充几个与后训练强相关的源码细节：

训练入口结构：train_tokenizer()工厂函数返回run.Partial(llm.train, ...)，将所有配置组织成可序列化、可覆盖的参数树，这正是命令行覆盖语法组件.参数=值的底层机制；入口逻辑位于if __name__ == "__main__": run.cli.main(llm.train, default_factory=train_tokenizer)。
模型前向与损失：TokenizerModel采用与 continuous_video.py（连续型，AE 公式化，输出 16 通道 latent）和 discrete_video.py（离散型，FSQ 量化，quant_loss计入总损失）一致的编码器-量化-解码器结构：encode得到 latent/token，decode重建视频，训练阶段同时返回重建帧与量化损失，重建质量驱动分词器适配新数据分布。
断点续训与保存：AutoResume保证意外中断后从最近 checkpoint 恢复；ModelCheckpoint每 100 步落盘，save_top_k=3控制磁盘占用。
快速联调：脚本内置FakeDataset（直接生成[2, 3, 33, 256, 256]的 bfloat16 随机张量与全 1 mask），可用于在真实数据就绪前验证训练管线是否可跑通。

总结

本文围绕 cosmos1/models/tokenizer/nemo/README.md 完整呈现了基于 NeMo Framework 后训练 Cosmos Tokenizer 的实战路径：从“为何要后训练（适配自驾驶等 Physical AI 数据分布）”出发，梳理了当前支持的两种模型（CV8x8x8、DV8x16x16）与通用后训练任务；随后按“准备 tar 分片数据集 → ImageTaskEncoder 自动预处理 → 安装依赖并torchrun拉起训练”三步落地，并结合 train_tokenizer.py 源码逐项解读了model / data / trainer / optim四大配置组件的默认值与覆盖方式；最后给出 wandb loss 收敛判读方法与关键源码细节。照此流程，你即可基于自有视频数据对 Cosmos Tokenizer 做领域适配，为下游扩散世界模型提供更贴合业务场景的压缩表示。

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

企业官网建设流程全解析