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/cosmosCosmos Tokenizer 是 NVIDIA Cosmos 世界基础模型平台中的视觉分词器组件负责把图像/视频压缩为连续潜变量Continuous Latent或离散 tokenDiscrete Token。本文以 cosmos1/models/tokenizer/nemo/README.md 为骨架完整讲解如何借助 NVIDIA NeMo Framework 对 Cosmos Tokenizer 进行后训练Post-training使预训练模型更好地建模自驾驶等场景中此前未见过的视频数据分布。读完本文你将掌握模型与后训练任务支持矩阵、从数据集组织到容器启动再到torchrun拉起训练的全流程、model/data/trainer/optim四大配置组件的底层默认参数以及如何通过 Weights Biaseswandb监控 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 网络定义同源因此后训练得到的权重与推理/下游扩散模型链路天然兼容。模型支持矩阵当前可后训练的两种 TokenizerNeMo 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 通道连续 latentCosmos-1.0-Tokenizer-DV8x16x16离散视频分词器时间 8x、空间 16x16总压缩率 2048xFSQ 量化输出 6 通道离散 code模型命名中的三个数字依次表示 T×H×W 方向的压缩倍数。在 configs.py 中可以找到两种模型的默认结构参数continuous_video设置latent_channels16、spatial_compression8、temporal_compression8discrete_video则设置spatial_compression16、temporal_compression8并使用 FSQ 量化器levels[8, 8, 8, 5, 5, 5]、embedding_dim6。关于空间与时间压缩因子的组合能力8x 或 16x 空间、4x 或 8x 时间可进一步阅读 tokenizer README。性能提示为获得最佳后训练性能官方推荐使用 H100-80GB 或 A100-80GB 级别的 GPU。后训练任务支持矩阵当前版本对 Cosmos Tokenizer 的后训练支持情况如下后训练任务支持状态通用后训练与验证General post-training and validationSupported已支持这与 POST_TRAINING.md 中描述的 Cosmos 整体后训练规划一致首版发布优先提供“通用后训练”脚本该文档中扩散 WFM 与自回归 WFM 的通用后训练均已支持指令控制、动作控制、相机控制、多视角生成等任务标记为 “Coming soon”。对于 Tokenizer 而言当前可直接投入使用的即是面向自定义视频分布的通用后训练。环境准备硬件、容器与凭据硬件与系统要求NVIDIA GPU 与驱动确保可访问 80GB 显存的 H100 或 A100以满足训练显存需求。容器化平台官方推荐使用 NVIDIA NeMo Docker Runtime也可选用 NVIDIA enrootNeMo 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/Cosmosdocker run --ipchost -it --gpusall \ -v $PATH_TO_COSMOS_REPO:/workspace/Cosmos \ nvcr.io/nvidia/nemo:24.12.01 bash其中$PATH_TO_COSMOS_REPO替换为本地仓库绝对路径。--ipchost用于共享主机 IPC 命名空间避免多进程 dataloader 的共享内存问题--gpusall暴露全部 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.mp4000000.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_height256计算得到时间随机裁剪TemporalRandomCrop(temporal_crop49)对时间维度做随机窗口采样归一化Normalize(mean0.5, std0.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.txtpip install megatron-energon4.0.0 pyav pip install githttps://github.com/NVIDIA/NeMo-Run.git pip install moviepy1.0.3 imageio其中megatron-energon负责 tar 分片数据集的读取与任务编码pyav用于视频解码NeMo-Run提供 CLI 参数化训练入口moviepy与imageio用于视频处理。2. 以 8 卡为例启动后训练Cosmos-1.0-Tokenizer-CV8x8x8export CKPT_PTHpath/to/your/HF/checkpoints/folder export DATApath/to/your/data # 可选用 Weights and Biases 监控训练进度 export WANDB_API_KEY/your/wandb/api/key export WANDB_PROJECT_NAMEcosmos-diffusion-nemo-post-training export WANDB_RUN_IDcosmos_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.modelCosmos-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 日志源码中loggerWandbLogger(...) if WANDB_API_KEY in os.environ else None。四大配置组件与可配置超参数train_tokenizer.py通过run.cli.factory(targetllm.train)暴露一个参数化训练入口支持四大配置组件任何超参数都可通过命令行覆盖形如组件.参数值。下表汇总了各组件的核心配置项及其源码默认值便于你按需调整model选择模型与指定 Checkpoint参数默认值说明model.modelCosmos-1.0-Tokenizer-CV8x8x8后训练目标模型可切换为 DV 系列model.jit_ckpt_pthNone须由命令行传入预训练 JIT checkpoint 路径训练从该权重初始化data数据加载与 dataloader 参数参数默认值说明data.pathNone须由命令行传入视频 tar 分片数据集根目录data.global_batch_size8全局 batch size跨所有 GPU 累积data.micro_batch_size1单卡微批大小data.num_workers1dataloader 工作进程数data.task_encoderImageTaskEncoder数据预处理管线见上文“第二步”源码中DiffusionDataModule同时承担训练与验证数据加载并通过MegatronDataSampler依据seq_len/micro_batch_size/global_batch_size自动推导梯度累积步数。trainer训练循环参数默认值说明trainer.devicesauto使用的设备数量通常由torchrun自动推导trainer.num_nodes环境变量SLURM_NNODES默认 1节点数适配 SLURM 集群trainer.acceleratorgpu加速器类型trainer.strategyddp_find_unused_parameters_trueDDP 策略兼容含未使用参数的网络trainer.max_epochs10000最大训练轮数配合按步保存实际以步数收敛为准trainer.precisionbf16混合精度与数据管线中的 bfloat16 一致trainer.val_check_interval100每 100 步做一次验证trainer.limit_val_batches1验证仅取 1 个 batchtrainer.log_every_n_steps1每步记录日志trainer.use_distributed_samplerFalse禁用分布式采样器由MegatronDataSampler接管此外trainer 挂载了三个回调ModelCheckpoint以global_step为监控指标每 100 训练步保存一次save_top_k3保留最新 3 份、PreemptionCallback应对集群抢占与TimingCallback计时统计。optim优化器参数参数默认值说明optim.optimizer_fntorch.optim.AdamW优化器类型lr1e-4学习率betas(0.5, 0.999)Adam 动量超参数beta1 取 0.5 而非常见 0.9适配生成式训练eps1e-8数值稳定项weight_decay0.01权重衰减fusedTrue启用 fused Adam 内核加速命令行覆盖示例model.jit_ckpt_pthyour/desired/path trainer.max_epochsyour/desired/epochs data.global_batch_sizeyour/desired/batch_size除上述四项外train_tokenizer.py还配置了resumeAutoResume存在 checkpoint 时自动续训、resume_past_endTrue允许越过已结束的 epoch 继续。按需调整这些参数后重新运行训练命令即可。训练过程监控与收敛判断训练启动后wandb 会记录 loss 曲线。根据训练脚本的默认配置日志每步写入一次。启动若干百次迭代后应当观察到 wandb 中loss开始下降。以官方后训练示例见 assets/loss.png为参考横轴为trainer/global_step0 至 500 步纵轴为loss约 0.050.1 区间训练初期 loss 从接近 0.1 快速下降至约 0.060.07随后在该区间内小幅波动并整体保持缓降、趋于稳定——这一模式可作为判断训练是否正常收敛的经验性参考。结合源码理解后训练流程为便于排查问题与二次开发这里补充几个与后训练强相关的源码细节训练入口结构train_tokenizer()工厂函数返回run.Partial(llm.train, ...)将所有配置组织成可序列化、可覆盖的参数树这正是命令行覆盖语法组件.参数值的底层机制入口逻辑位于if __name__ __main__: run.cli.main(llm.train, default_factorytrain_tokenizer)。模型前向与损失TokenizerModel采用与 continuous_video.py连续型AE 公式化输出 16 通道 latent和 discrete_video.py离散型FSQ 量化quant_loss计入总损失一致的编码器-量化-解码器结构encode得到 latent/tokendecode重建视频训练阶段同时返回重建帧与量化损失重建质量驱动分词器适配新数据分布。断点续训与保存AutoResume保证意外中断后从最近 checkpoint 恢复ModelCheckpoint每 100 步落盘save_top_k3控制磁盘占用。快速联调脚本内置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 做领域适配为下游扩散世界模型提供更贴合业务场景的压缩表示。【免费下载链接】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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 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/cosmosCosmos Tokenizer 是 NVIDIA Cosmos 世界基础模型平台中的视觉分词器组件负责把图像/视频压缩为连续潜变量Continuous Latent或离散 tokenDiscrete Token。本文以 cosmos1/models/tokenizer/nemo/README.md 为骨架完整讲解如何借助 NVIDIA NeMo Framework 对 Cosmos Tokenizer 进行后训练Post-training使预训练模型更好地建模自驾驶等场景中此前未见过的视频数据分布。读完本文你将掌握模型与后训练任务支持矩阵、从数据集组织到容器启动再到torchrun拉起训练的全流程、model/data/trainer/optim四大配置组件的底层默认参数以及如何通过 Weights Biaseswandb监控 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 网络定义同源因此后训练得到的权重与推理/下游扩散模型链路天然兼容。模型支持矩阵当前可后训练的两种 TokenizerNeMo 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 通道连续 latentCosmos-1.0-Tokenizer-DV8x16x16离散视频分词器时间 8x、空间 16x16总压缩率 2048xFSQ 量化输出 6 通道离散 code模型命名中的三个数字依次表示 T×H×W 方向的压缩倍数。在 configs.py 中可以找到两种模型的默认结构参数continuous_video设置latent_channels16、spatial_compression8、temporal_compression8discrete_video则设置spatial_compression16、temporal_compression8并使用 FSQ 量化器levels[8, 8, 8, 5, 5, 5]、embedding_dim6。关于空间与时间压缩因子的组合能力8x 或 16x 空间、4x 或 8x 时间可进一步阅读 tokenizer README。性能提示为获得最佳后训练性能官方推荐使用 H100-80GB 或 A100-80GB 级别的 GPU。后训练任务支持矩阵当前版本对 Cosmos Tokenizer 的后训练支持情况如下后训练任务支持状态通用后训练与验证General post-training and validationSupported已支持这与 POST_TRAINING.md 中描述的 Cosmos 整体后训练规划一致首版发布优先提供“通用后训练”脚本该文档中扩散 WFM 与自回归 WFM 的通用后训练均已支持指令控制、动作控制、相机控制、多视角生成等任务标记为 “Coming soon”。对于 Tokenizer 而言当前可直接投入使用的即是面向自定义视频分布的通用后训练。环境准备硬件、容器与凭据硬件与系统要求NVIDIA GPU 与驱动确保可访问 80GB 显存的 H100 或 A100以满足训练显存需求。容器化平台官方推荐使用 NVIDIA NeMo Docker Runtime也可选用 NVIDIA enrootNeMo 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/Cosmosdocker run --ipchost -it --gpusall \ -v $PATH_TO_COSMOS_REPO:/workspace/Cosmos \ nvcr.io/nvidia/nemo:24.12.01 bash其中$PATH_TO_COSMOS_REPO替换为本地仓库绝对路径。--ipchost用于共享主机 IPC 命名空间避免多进程 dataloader 的共享内存问题--gpusall暴露全部 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.mp4000000.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_height256计算得到时间随机裁剪TemporalRandomCrop(temporal_crop49)对时间维度做随机窗口采样归一化Normalize(mean0.5, std0.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.txtpip install megatron-energon4.0.0 pyav pip install githttps://github.com/NVIDIA/NeMo-Run.git pip install moviepy1.0.3 imageio其中megatron-energon负责 tar 分片数据集的读取与任务编码pyav用于视频解码NeMo-Run提供 CLI 参数化训练入口moviepy与imageio用于视频处理。2. 以 8 卡为例启动后训练Cosmos-1.0-Tokenizer-CV8x8x8export CKPT_PTHpath/to/your/HF/checkpoints/folder export DATApath/to/your/data # 可选用 Weights and Biases 监控训练进度 export WANDB_API_KEY/your/wandb/api/key export WANDB_PROJECT_NAMEcosmos-diffusion-nemo-post-training export WANDB_RUN_IDcosmos_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.modelCosmos-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 日志源码中loggerWandbLogger(...) if WANDB_API_KEY in os.environ else None。四大配置组件与可配置超参数train_tokenizer.py通过run.cli.factory(targetllm.train)暴露一个参数化训练入口支持四大配置组件任何超参数都可通过命令行覆盖形如组件.参数值。下表汇总了各组件的核心配置项及其源码默认值便于你按需调整model选择模型与指定 Checkpoint参数默认值说明model.modelCosmos-1.0-Tokenizer-CV8x8x8后训练目标模型可切换为 DV 系列model.jit_ckpt_pthNone须由命令行传入预训练 JIT checkpoint 路径训练从该权重初始化data数据加载与 dataloader 参数参数默认值说明data.pathNone须由命令行传入视频 tar 分片数据集根目录data.global_batch_size8全局 batch size跨所有 GPU 累积data.micro_batch_size1单卡微批大小data.num_workers1dataloader 工作进程数data.task_encoderImageTaskEncoder数据预处理管线见上文“第二步”源码中DiffusionDataModule同时承担训练与验证数据加载并通过MegatronDataSampler依据seq_len/micro_batch_size/global_batch_size自动推导梯度累积步数。trainer训练循环参数默认值说明trainer.devicesauto使用的设备数量通常由torchrun自动推导trainer.num_nodes环境变量SLURM_NNODES默认 1节点数适配 SLURM 集群trainer.acceleratorgpu加速器类型trainer.strategyddp_find_unused_parameters_trueDDP 策略兼容含未使用参数的网络trainer.max_epochs10000最大训练轮数配合按步保存实际以步数收敛为准trainer.precisionbf16混合精度与数据管线中的 bfloat16 一致trainer.val_check_interval100每 100 步做一次验证trainer.limit_val_batches1验证仅取 1 个 batchtrainer.log_every_n_steps1每步记录日志trainer.use_distributed_samplerFalse禁用分布式采样器由MegatronDataSampler接管此外trainer 挂载了三个回调ModelCheckpoint以global_step为监控指标每 100 训练步保存一次save_top_k3保留最新 3 份、PreemptionCallback应对集群抢占与TimingCallback计时统计。optim优化器参数参数默认值说明optim.optimizer_fntorch.optim.AdamW优化器类型lr1e-4学习率betas(0.5, 0.999)Adam 动量超参数beta1 取 0.5 而非常见 0.9适配生成式训练eps1e-8数值稳定项weight_decay0.01权重衰减fusedTrue启用 fused Adam 内核加速命令行覆盖示例model.jit_ckpt_pthyour/desired/path trainer.max_epochsyour/desired/epochs data.global_batch_sizeyour/desired/batch_size除上述四项外train_tokenizer.py还配置了resumeAutoResume存在 checkpoint 时自动续训、resume_past_endTrue允许越过已结束的 epoch 继续。按需调整这些参数后重新运行训练命令即可。训练过程监控与收敛判断训练启动后wandb 会记录 loss 曲线。根据训练脚本的默认配置日志每步写入一次。启动若干百次迭代后应当观察到 wandb 中loss开始下降。以官方后训练示例见 assets/loss.png为参考横轴为trainer/global_step0 至 500 步纵轴为loss约 0.050.1 区间训练初期 loss 从接近 0.1 快速下降至约 0.060.07随后在该区间内小幅波动并整体保持缓降、趋于稳定——这一模式可作为判断训练是否正常收敛的经验性参考。结合源码理解后训练流程为便于排查问题与二次开发这里补充几个与后训练强相关的源码细节训练入口结构train_tokenizer()工厂函数返回run.Partial(llm.train, ...)将所有配置组织成可序列化、可覆盖的参数树这正是命令行覆盖语法组件.参数值的底层机制入口逻辑位于if __name__ __main__: run.cli.main(llm.train, default_factorytrain_tokenizer)。模型前向与损失TokenizerModel采用与 continuous_video.py连续型AE 公式化输出 16 通道 latent和 discrete_video.py离散型FSQ 量化quant_loss计入总损失一致的编码器-量化-解码器结构encode得到 latent/tokendecode重建视频训练阶段同时返回重建帧与量化损失重建质量驱动分词器适配新数据分布。断点续训与保存AutoResume保证意外中断后从最近 checkpoint 恢复ModelCheckpoint每 100 步落盘save_top_k3控制磁盘占用。快速联调脚本内置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 做领域适配为下游扩散世界模型提供更贴合业务场景的压缩表示。【免费下载链接】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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考