TorchTitan 训练 DeepSeek-V3 实践：模型注册表、融合内核优化与 HF 到 DCP 检查点转换

TorchTitan 训练 DeepSeek-V3 实践模型注册表、融合内核优化与 HF 到 DCP 检查点转换【免费下载链接】torchtitanA PyTorch native platform for training generative AI models项目地址: https://gitcode.com/GitHub_Trending/to/torchtitan本文基于 TorchTitan 仓库中 DeepSeek-V3 模型文档完整覆盖该模型在 TorchTitan 中的落地流程下载 tokenizer、通过run_train.sh启动 debugmodel/16B/671B 三档训练、启用融合 Triton 内核的性能优化选项以及 HuggingFace safetensors 到 DCP 格式的检查点离线转换并结合 模型注册表 与 训练配置 源码解读各配置的默认值与底层机制帮助你在多机多卡环境下完成 DeepSeek-V3 架构的预训练调试与权重迁移。一、DeepSeek-V3 模块在 TorchTitan 中的位置DeepSeek-V3 的实现在 torchtitan/models/deepseek_v3/ 目录下核心文件职责如下文件职责__init__.py模型注册表model_registry定义 debugmodel/16B/236B/671B 四档 flavor 及 MLAMoE 层构建逻辑model.pyDeepSeekV3Model、MLA 注意力Attention、DeepSeekV3TransformerBlock的定义moe.py / mtp.pyDeepSeek-V3 路由器DeepSeekV3RouterMTP多 token 预测MTPDecoder、MTPLossconfig_registry.py各 flavor 的Trainer.Config数据、优化器、并行度、编译等sharding.py / parallelize.pyFSDP/TP/SP/EP 分片策略与并行化入口state_dict_adapter.pyHF 命名与 TorchTitan 命名之间的 state dict 映射支撑检查点互转MTP.mdMTP 实现的完整设计文档本文聚焦主 README 的训练流程二、下载 Tokenizer2.1 两条下载命令671B 参数模型使用 DeepSeek-V3.1-Base 的官方 tokenizer自动下载tokenizer.json和tokenizer_config.json# DeepSeek 671B tokenizer python scripts/download_hf_assets.py --repo_id deepseek-ai/DeepSeek-V3.1-Base --assets tokenizer16B 参数模型则复用 deepseek-moe-16b-base 的 tokenizer# DeepSeek 16B tokenizer python scripts/download_hf_assets.py --repo_id deepseek-moe-16b-base --assets tokenizer2.2 为什么 16B 复用另一个仓库的 tokenizer原 README 明确说明TorchTitan 复用 deepseek-moe-16b-base 的 tokenizer 只是为了帮助用户测试和运行 16B 模型它不是DeepSeek-V3-16B 模型的官方 tokenizer。根本原因在于架构差异DeepSeek-V3 模型与 deepseek-moe 系列在注意力实现MLA、MoE router 实现等方面都不同导致无法直接把 deepseek-moe-16b 的权重加载进 DeepSeek-V3-16B。这一判断也能从 模型注册表源码得到印证——TorchTitan 的 16B 配置_16b采用的是 DeepSeek-V3 家族参数64 个专家、top_k6、MLA 注意力而非 deepseek-moe-16b 的结构。2.3 下载脚本的行为细节从 download_hf_assets.py 源码可以看到--assets tokenizer会匹配下载tokenizer.json、tokenizer_config.json、tokenizer.model、vocab.txt、vocab.json、merges.txt、special_tokens_map.json等模式脚本会根据repo_id中的模型名/之后的部分自动在local_dir下创建同名子目录存放文件。这解释了训练配置中hf_assets_path的取值——config_registry.py 中 671B 指向./assets/hf/DeepSeek-V3.1-Base16B 指向./assets/hf/deepseek-moe-16b-base即脚本按仓库名建目录后的产物位置。三、三档训练命令与 run_train.sh 启动机制3.1 启动脚本内部做了什么三条训练命令统一通过仓库根目录的 run_train.sh 执行该脚本的关键行为通过环境变量MODULE与CONFIG指定模型模块和训练配置默认llama3/llama3_debugmodelNGPU默认 8最终调用torchrun --nproc_per_node${NGPU} --rdzv_backend c10d -m torchtitan.train --module ${MODULE} --config ${CONFIG} $因此训练参数可以直接以 tyro 风格追加在命令末尾LOG_RANK默认 0指定只 tee 输出哪个 rank 的日志支持COMM_MODEfake_backend干跑模式使用伪造的进程组、单 GPU、无需 NCCL 初始化配合--training.steps 1可用来在纯 CPU/单机环境验证配置合法性。3.2 三条核心命令原 README 完整继承# 小型模型快速调试Quick debug run with small model MODULEdeepseek_v3 CONFIGdeepseek_v3_debugmodel ./run_train.sh# 16B 参数模型适配自较早的 16B 参数模型deepseek-moe-16b-base 同系列参数规模 MODULEdeepseek_v3 CONFIGdeepseek_v3_16b ./run_train.sh# 671B 参数模型 MODULEdeepseek_v3 CONFIGdeepseek_v3_671b ./run_train.sh3.3 各配置对应的 Trainer.Config 参数对照 config_registry.py 源码三个默认配置的实际参数如下配置项deepseek_v3_debugmodeldeepseek_v3_16bdeepseek_v3_671b模型 flavordebugmodeldim 2566 层8 专家16Bdim 204827 层64 专家671Bdim 716861 层256 专家tokenizer 资产路径./tests/assets/tokenizer仓库自带测试 tokenizer无需下载./assets/hf/deepseek-moe-16b-base./assets/hf/DeepSeek-V3.1-Base数据集c4_test仓库内小样本c4c4优化器/学习率AdamWlr8e-4AdamWlr2.2e-4AdamWlr2.2e-4LR 调度线性衰减warmup 2 步min_lr_factor0余弦衰减decay_ratio0.8min_lr_factor0.1余弦衰减warmup 2000 步decay_ratio0.8min_lr_factor0.1总步数10100010000微批 token 数/DP rank8 × max_context_length4 × max_context_length4 × max_context_length并行度EP1EP8PP 调度 Interleaved1F1BEP2PP 调度 Interleaved1F1B检查点每 10 步每 10 步每 500 步激活检查点SelectiveACSelectiveACSelectiveAC编译默认开启 loss 编译开启 loss 编译CUDA Graphs默认禁用禁用三档配置均使用ChunkedLossWrapper包裹CrossEntropyLoss并把词表大小传给 loss 以支持 loss-parallel 交叉熵路径。16B/671B 的注意力后端通过model_registry(..., attn_backendflex)指定为 flex attention。四档 flavor 的上下文长度上限均为 16384见 deepseekv3_configs 表deepseekv3_configs { debugmodel: (_debugmodel, 16384), 16B: (_16b, 16384), 236B: (_236b, 16384), 671B: (_671b, 16384), }需要说明的是注册表中实际还存在236Bflavordim 5120、60 层、160 专家、top_k6、Softmax 路由分数并带 8 组/限 3 组的 group-limited 路由但 README 默认推荐的是上面三档命令236B 可通过CONFIG指向对应注册配置使用。3.4 注册表中的架构要点从init.py 的 flavor 定义可以读出 DeepSeek-V3 的关键架构选择MLA多头潜在注意力所有 flavor 都有kv_lora_rank512、qk_nope_head_dim128、qk_rope_head_dim64、v_head_dim128debugmodel/16B 设q_lora_rank0Q 走单一线性wq236B/671B 设q_lora_rank1536Q 走低秩分解wq_a→q_norm→wq_b。model.py 中的 Attention 实现了 KV 压缩wkv_a压缩到 512 维 latentkv_norm归一化后wkv_b展开K 的 RoPE 分量通过对所有头 expand 复用来共享前密后稀的 FFN 布局n_dense_layers之前的层用稠密 FeedForwarddebugmodel/16B/236B 为 1 层671B 为 3 层之后全部为 MoE 层256 专家 × 671Btop_k88 组限 4 组的辅助损失路由aux_loss_coeff1e-3YaRN 长上下文 RoPE统一使用ComplexRoPEscalingyarn、rope_factor40.0、original_seq_len4096与 MLA 的mscale共同作用于 softmax 缩放见 model.py L92-L94可选变体config_registry.py还定义了deepseek_v3_debugmodel_mxfp8对 MoE grouped GEMM 及稠密线性做 MXFP8 量化pad_multiple128为 sm_100/B200 上 CuTeDSL 量化内核的硬要求、deepseek_v3_671b_float8float8 路径注释标明需要 torchao 且仅支持 NVIDIA SM89 或 AMD MI300其他后端构建时会报错应回退到普通 671B 配置、deepseek_v3_16b_hybridep/deepseek_v3_debugmodel_hybridepmoe_comm_backendhybridep并把non_blocking_capacity_factor1.0以及deepseek_v3_debugmodel_mtpnum_mtp_layers1、内部 loss 换为MTPLoss.Config其mtp_scale默认 0.3见 mtp.py。MTP 的输入构造、并行分片与损失计算细节可进一步阅读 MTP.md。四、性能优化选项融合 Triton 内核原 README 指出DeepSeek-V3 可以可选地启用三类融合 Triton 内核MLA Q/KV 组装、ComplexRoPE和SwiGLU。这些 override 通过 tyro 的--override.imports机制注入且保持现有模型参数名与检查点布局不变即转换后的模型可直接加载/保存原有格式的检查点MODULEdeepseek_v3 CONFIGdeepseek_v3_671b ./run_train.sh \ --override.imports torchtitan.overrides.fused_mla.fused_mla,torchtitan.overrides.fused_swiglu.fused_swiglu对应的实现位于 torchtitan/overrides/fused_mla.py 与 torchtitan/overrides/fused_swiglu.py模块内通过 override 注册替换默认的逐算子实现由于只替换计算、不改变参数布局这一优化对 checkpoint 兼容性零成本。run_train.sh会原样透传$因此该参数无需修改脚本即可生效。五、HuggingFace 到 DCP 检查点转换5.1 转换命令与适用范围TorchTitan 为 DeepSeek-V3 实现了StateDictAdapter用于 HuggingFace safetensors 到 DCPPyTorch Distributed Checkpoint格式的转换。原 README 明确了当前限制只支持从 HF 检查点到 DCP 检查点的离线转换使用 CPU plain tensor即单向、离线、非分布式。命令如下python scripts/checkpoint_conversion/convert_from_hf.py hf_checkpoints_dir dcp_output_dir --model_name deepseek_v3 --model_flavor 671B其中hf_checkpoints_dir为包含*.safetensors与 index 文件的 HF 权重目录dcp_output_dir为 DCP 输出目录--model_flavor需与注册表 flavor 名一致16B/671B 等。完整用法说明见 scripts/checkpoint_conversion/README.md。5.2 脚本执行流程源码走读convert_from_hf.py 的核心逻辑在convert_from_hf()函数中整个流程在torch.inference_mode()下运行importlib.import_module(ftorchtitan.models.{model_name})动态加载模型模块调用model_registry(model_flavor)拿到ModelSpec在torch.device(cpu)上build()模型——这就是 README 所说 using CPU plain tensor 的来源整个转换不依赖 GPU用ModelWrapper包装后取 TorchTitan 命名的空 state dictmodel._get_state_dict()调用sd_adapter.to_hf(state_dict)把 TT 命名映射回 HF 命名反向重命名 必要的 reshape/拼接得到 HF 命名的空 state dictdcp.load(hf_state_dict, storage_readerHuggingFaceStorageReader(pathinput_dir))从 HF safetensors 目录读取权重填入再sd_adapter.from_hf(hf_state_dict)映射回 TorchTitan 命名最后dcp.save(...)写出标准 DCP 目录。映射规则由 DeepSeekV3StateDictAdapter 定义负责处理 MLA 的wq_a/wq_b/wkv_a/wkv_b/wo、路由门、专家权重等在 HF 命名与 TorchTitan 命名间的差异。转换完成后得到的 DCP 目录即可被 TorchTitan 训练器通过--checkpoint.enable等参数直接恢复训练见 checkpoint.md。六、快速上手清单按 flavor 执行 tokenizer 下载命令debugmodel 可用仓库自带./tests/assets/tokenizer无需下载用MODULEdeepseek_v3 CONFIGconfig ./run_train.sh启动追加NGPU...、LOG_RANK...或 tyro 参数覆盖默认值无 GPU 时可用NGPU8 COMM_MODEfake_backend ./run_train.sh干跑验证配置大模型追求性能时为 671B 追加--override.imports启用 fused MLA 与 fused SwiGLU 内核已有 HF 权重时先用convert_from_hf.py离线转换为 DCP 格式再断点续训需要 MTP 训练、MXFP8/Float8 量化或 HybridEP 通信后端时参考 config_registry.py 中的对应变体配置与 MTP 设计文档。【免费下载链接】torchtitanA PyTorch native platform for training generative AI models项目地址: https://gitcode.com/GitHub_Trending/to/torchtitan创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考