使用 TorchTitan HuggingFace Transformers 建模后端以 4D 并行与 torch.compile 训练任意 HF 模型【免费下载链接】torchtitanA PyTorch native platform for training generative AI models项目地址: https://gitcode.com/GitHub_Trending/to/torchtitan本文介绍 TorchTitan 仓库中的transformers_modeling_backendHuggingFace Transformers 建模后端它允许把 HFHuggingFace生态中任意transformers模型直接接入 TorchTitan 的训练框架在原生 PyTorch 的 FSDP/TP/PP/CP 四维并行与torch.compile之下完成大规模训练。读完本文你将掌握该后端的启动配置、模型/并行度支持矩阵、FlexAttention 与 MoE 替换两大核心机制以及它当前的已知限制。该后端位于 torchtitan/experiments/transformers_modeling_backend/其设计目标写在 README.md 的第一句让 HF transformers 模型能够以4D 并行FSDP/TP/PP/CPtorch.compile的方式被训练并且整个后端建立在 TorchTitan 的SPMD 类型系统之上。快速开始最小可运行配置环境依赖后端对transformers库版本有明确要求README 中写明transformers5.9.0切换模块与配置后端通过 TorchTitan 的--module/--config机制接入无需改动训练主程序。按 README 给出的 diff把默认的 Llama3 模块替换为建模后端模块并选择debugmodel调试配置... - --module llama3 --module transformers_modeling_backend --config transformers_modeling_backend_debugmodel ...启动训练README 给出的启动命令需先准备好 tokenizer 资源LOG_RANK7 MODELtransformers_modeling_backend CONFIGtransformers_modeling_backend_debugmodel ./run_train.sh --compile.enable要点LOG_RANK7指定从最后一个 rank 读取日志。对于流水线并行PP场景loss 只存在于最后一个 pipeline stagerank 0 打印的是-1.0占位值因此要选择NGPU - 1这个 rank 才能看到真实 loss这一约定同样体现在 run_moe_tests.sh 的LOG_RANK$((ngpu - 1))中。--compile.enable开启torch.compile这也是该后端支持的核心特性之一。README 明确提示运行前必须已创建好 tokenizer。配置体系从hf_model到六个预置 flavor扩展配置入口TransformersBackendConfig后端在 config_registry.py 中定义了TransformersBackendConfig(Trainer.Config)它在 TorchTitan 标准Trainer.Config之上新增了一个核心字段dataclass(kw_onlyTrue, slotsTrue) class TransformersBackendConfig(Trainer.Config): hf_model: str HuggingFace model ID (e.g., Qwen/Qwen2.5-7B)hf_model即 HF 模型仓库 ID如Qwen/Qwen3-30B-A3B训练时通过--hf_model id或预置 config 指定。模型的实际架构、尺寸、MoE 参数全部由 HF 侧的AutoConfig.from_pretrained(hf_model)加载见 model.py 中的update_from_config因此更换模型通常只需更换hf_model无需修改任何并行化代码。六个预置 flavor仓库在 config_registry.py 中预置了以下配置flavor 定义在init.py 的flavors字典中配置名flavor用途transformers_modeling_backend_debugmodeldebugmodel稠密模型调试dim256, 2 层, 16 头causal掩码默认hf_modelQwen/Qwen3-4B-Instruct-2507transformers_modeling_backend_debugmodel_moedebugmodel_moeMoE 调试dim2048, 4 层, 8 专家, top-2默认hf_modelQwen/Qwen3-30B-A3Btransformers_modeling_backend_fullfull稠密全尺寸训练不覆盖任何 HF 配置全部取自 HF configtransformers_modeling_backend_full_moefull_moeMoE 全尺寸训练dim2048, 48 层, 128 专家, top-8transformers_modeling_backend_sft_fullsft_fullSFT 全尺寸block_causal掩码 initial_load_in_hfTrue加载 HF 预训练权重transformers_modeling_backend_sft_debugmodelsft_debugmodelSFT 调试dim256, 2 层,block_causal掩码用tests/assets/sft_test/data.json做 chat 数据其中debugmodel系列是 CI 与日常验证的主力hf_model使用真实 HF 模型 ID如 Qwen3-4B但 flavor 显式覆盖dim/n_layers/n_heads等尺寸把模型缩小为可快速训练的最小架构从而在 debug 规模上验证完整训练链路。两层配置模型TitanModelConfig/TitanMoeModelConfigflavor 的核心是init.py 中定义的两层 dataclass。HFTransformerModel.Config是桥接 TorchTitan 与 HF 的配置外壳内部持有model_config而TitanModelConfig稠密与其子类TitanMoeModelConfigMoE才是真正的模型参数描述。字段设计原则源码注释明确说明凡是能映射到 HF config key 的字段默认值一律为None这样AutoConfig.from_pretrained加载的值会被保留——如果给了非None默认值就会覆盖 HF 配置、强行改变模型架构例如rope_theta只有显式设置如 debugmodel 的尺寸覆盖才会注入。而 TorchTitan 独有的字段则保留具体默认值。TitanModelConfig关键字段字段默认值说明dim/n_layers/n_heads/n_kv_heads/vocab_size/intermediate_sizeNone映射到 HF 的hidden_size/num_hidden_layers/num_attention_heads/num_key_value_heads/vocab_size/intermediate_size映射表见 model.py 的_TT_TO_HF_MAPPINGSnorm_eps/rope_thetaNone映射到 HF 的rms_norm_eps/rope_thetamultiple_of256TorchTitan 独有显式覆盖dim时用于推导 FFN 尺寸的对齐因子ffn_dim_multiplierNoneTorchTitan 独有推导 FFN 尺寸的倍率max_seq_len2048训练时由training.max_context_length覆盖depth_initTrue是否启用深度相关的初始化缩放attn_mask_typecausalFlexAttention 掩码类型见下文「注意力」一节TitanMoeModelConfig在稠密基础上增加字段默认值说明num_experts/num_experts_per_tok/moe_intermediate_sizeNone路由专家总数 / top-k / 每个专家 MLP 的中间维度decoder_sparse_stepNone每隔 N 层把稠密 MLP 替换为 MoE 块1 每层都是 MoEnorm_topk_probNone是否把 top-k 路由分数归一化为和为 1HFnorm_topk_prob约定num_nextn_predict_layersNoneDeepSeek V3 风格的多 token 预测MTP头数None 或 0 表示禁用experts_implementationgrouped_mm专家前向 kernelgrouped_mm融合快速路径、batched_mm、eagerHF 原始 for 循环数值参考、native使用 HF 模型自带 kernel不可设置 kernel 的模型唯一合法选项load_balance_coeff1e-3无辅助损失的 MoE 负载均衡步长None禁用comm_backendstandard专家并行的 token 分发后端standardPyTorch all-to-all、deepepH100/NVLink 的 DeepEP、hybridepGB200/NVLink72 的 HybridEP核心机制一模型适配层HFTransformerModel后端在 model.py 中通过HFTransformerModel(BaseModel)把任意PreTrainedModel包装成 TorchTitan 认可的BaseModel。这一层解决的是两类接口不匹配1. 架构发现与模型类解析__init__中先取config.architectures[0]作为模型类名从globals()或transformers包动态导入失败时再通过model_type → MODEL_FOR_CAUSAL_LM_MAPPING_NAMES的 Auto 映射兜底解析。因此该后端面向CausalLM类模型且支持transformers包中注册的任意架构。2. 属性映射TT 风格访问 HF 模型_TT_TO_HF_MAPPINGS定义了 TorchTitan 属性名与 HF 属性名的对应关系如dim → hidden_size、n_kv_heads → num_key_value_heads。Config._create_getter_setter_dynamically会动态创建 property让self.dim读写self.hidden_size从而对上层暴露统一的 TorchTitan 风格接口。update_from_config在加载 HF config 后通过vars(hf_model_config)复制所有属性包括to_dict()会漏掉的运行时计算属性例如 DeepSeek V3 的qk_head_dim并保留attribute_map例如 DeepSeek V3 把num_local_experts映射到n_routed_experts最后重新注入用户在 flavor 中显式设置的覆盖值保证debugmodel的尺寸覆盖不被 HF config 冲掉。值得注意的两个细节均有源码注释佐证远程 config 黑名单_REMOTE_CONFIG_DENYLIST {deepseek_v2, deepseek_v3}。当模型仓库自带旧版远程 config 实现时混用会丢失本地模型代码期望的兼容属性因此这两个模型类型强制使用本地transformers的 config/model 对trust_remote_codeFalse。权重初始化补丁_patch_hf_llama_like对 HF 的_init_weights/_initialize_weights做 monkey-patch使其匹配 TorchTitan 的初始化方案attention/MLP 深度相关初始化init_std 0.02 / (2 * (layer_idx 1)) ** 0.5、MLA 低秩投影q_a_proj等按std0.02初始化、专家 3D 参数trunc_normal_等。这一步是 checkpoint 在 TorchTitan 格式与 HF 格式之间逐位可复现转换的前提。3. 结构规整SliceableModuleDictHF 模型用nn.ModuleList保存 decoder layers而 TorchTitan 的 PP/FSDP 期望ModuleDict语义。__init__末尾把nn.ModuleList转换为SliceableModuleDictkey 为字符串化的层索引见 model.py 的SliceableModuleDict既保持 state dict key 与 checkpoint 一致又支持按 slice 取层供 PP 切分使用。模型顶层则通过named_children()暴露tok_embeddings / layers / norm / lm_head / rotary_emb五个统一子模块兼容 Phi 的final_layernorm等命名差异。核心机制二注意力统一走 FlexAttention后端把所有HF 注意力路由到torch.nn.attention.flex_attention这是它区别于直接塞给 HF 训练的关键设计README「Attention」一节 model.py 的_configure_hf_attention注册自定义实现名_ATTN_IMPLEMENTATION flex_torchtitan到 HF 的AttentionInterface._global_mapping并设置config._attn_implementation从而绕过 HF 每个模型自带的_supports_flex_attn门禁统一走_flex_attention_torchtitanshim。掩码由attn_mask_type选择causal普通因果掩码get_causal_mask_mod()block_causal因果 同文档约束and_masks(get_causal_mask_mod(), get_document_mask_mod(positions))用于 packed / SFT 序列——README 明确说明is_causal无法表达跨样本掩码packed 训练必须用block_causal。掩码构建在get_attention_masks()中完成BLOCK_SIZE128 的 flex BlockMask并通过 HF 常规的attention_mask参数透传HF 对已是 4D/BlockMask 的输入原样返回。DSA 例外DeepSeek 稀疏注意力如 GLM-5model_typeglm_moe_dsa的 indexer 子注意力把掩码当稠密加法张量用对掩码调用.dim()、做逐元素相加flexBlockMask无法满足因此_uses_dsa()会检测index_topk配置并改用_build_dense_attention_mask()生成稠密 4D 加法掩码flex 以score_mask消费它causal/block_causal语义保持一致。在并行层面注意力运行在FlexInnerAttention之上HFFlexKernel把 flex HOP 包装成 TorchTitanModule在声明式 TP 下 q/k/v 携带 head 分片的 SPMD 标注输出获得对应标注CP 下则由 parallelize.py 的_wrap_flex_kernel_cp在每个 flex kernel 内部沿 CP 轴对 k/v 做 all-gather对应 TorchTitan 的flex_cp_allgather路径q 保持 seq 分片从而每个 rank 只对本地 seq 计算注意力、消费 Q 分片 / KV 全量的 BlockMask。跨文档泄漏防护有专门测试验证见 tests/test_flex_attention.py扰动前一个文档的 key/value 必须让后一个文档的输出逐位不变。核心机制三MoE 块替换为 TorchTitan grouped-experts MoE对于 MoE 模型后端并不直接训练 HF 自带的 MoE 实现而是把 MoE 块整体替换为 TorchTitan 的 grouped-experts MoEREADME「MoE」一节从而获得专家并行EP与grouped_mm融合快速路径。替换逻辑在 moe_replacement.py分两阶段Phase 1init 时prepare_native_moe_configs在 meta 设备上探测 HF MoE 块Qwen3MoeSparseMoeBlock等从gate/router、experts与 config 中解析出专家数兼容num_experts/n_routed_experts/num_local_experts三种命名、top-k、路由打分函数softmax / sigmoid通过e_score_correction_biasbuffer 探测、路由归一化、DeepSeek 的分组路由n_group/topk_group、shared expert 信息等构建 TitanMoE.Config存到层的_native_moe_config。Phase 2parallelize 时build_and_swap_native_moe在已知 EP/TP 度后设置set_moe_sharding_config、构建 Titan MoE、to_empty init_states物化参数再替换掉层的mlp属性。对 Gemma4 这种层级 MoErouter/experts 是 decoder layer 的兄弟节点、稠密 MLP 充当 shared expert替换后还会删除 HF 的独立 MoE 路径模块避免参数重复注册。替换的保真度由 MODEL_COMPATIBILITY.md 记录以未修改的 HF forward 为 gold standard在真实隐藏维度、缩减专家数E8、单层配置下比较输出。结果摘要KL 散度阈值PASS 1e-6WARN 1e-3FAIL ≥ 1e-3模型model_type状态KL div说明Qwen3-30B-A3Bqwen3_moePASS5.31e-07差异仅剩 dispatcher 精度OLMoE-1B-7BolmoePASS4.04e-07同上DeepSeek-V2-Litedeepseek_v2PASS2.64e-07同上Mixtral-8x7BmixtralWARN4.62e-05dispatcher 精度DeepSeek-V3deepseek_v3WARN3.66e-04dispatcher 精度GLM-4.7glm4_moeWARN5.06e-05dispatcher 精度GLM-5 (DSA)glm_moe_dsaWARN2.67e-04dispatcher 精度Gemma-4-26Bgemma4_textFAIL4.19e-03激活函数GeGLU vs SwiGLU 路由器架构差异文档还给出了 WARN 模型的三大数值差异根源均为 TitanLocalTokenDispatcher与 HF 的独立差异accumulation dtypebf16 vs f32、accumulation orderscatter_add按专家序 vs HF 按 token 序求和、topk 排序sortedFalsevssortedTrue三者修正后 max_diff0.00。这些结论为理解替换 MoE 是否安全提供了量化依据。支持的模型与并行度矩阵README 列出该后端经过验证的模型清单均为真实 HF 仓库 ID稠密模型支持 FSDP / CP / TP / PP /torch.compilemeta-llama/Llama-3.2-1Bmicrosoft/phi-2Qwen/Qwen2.5-7Bmistralai/Mistral-7B-v0.1ByteDance-Seed/Seed-Coder-8B-InstructQwen/Qwen3-4B-Instruct-2507arcee-ai/AFM-4.5Bibm-granite/granite-3b-code-base-2kbaidu/ERNIE-4.5-0.3B-Base-PTkyutai/helium-1-preview-2ballenai/OLMo-7B-hfmistralai/Ministral-8B-Instruct-2410MoE 模型并行支持见下方矩阵Qwen/Qwen3-30B-A3Bqwen3_moeGQAmistralai/Mixtral-8x7B-Instruct-v0.1mixtralGQAallenai/OLMoE-1B-7B-0924olmoeGQAdeepseek-ai/DeepSeek-V2-Litedeepseek_v2MLAdeepseek-ai/DeepSeek-V3deepseek_v3MLAzai-org/GLM-4.7glm4_moeMLAzai-org/GLM-5glm_moe_dsaMLA DSA 稀疏注意力google/gemma-4-26B-A4B-itgemma4_textREADME 给出的 MoE 并行度矩阵debug 规模、2 步 loss 递减验证通过FSDP 与所有列出的轴均可组合模型attnTPEPCP备注Qwen3-30B-A3BGQA是是是完整矩阵最高 TPEPCPMixtral-8x7B、OLMoE-1B-7BGQA是是是TPEP、EPCPDeepSeek-V2-Lite、V3、GLM-4.7MLA是是是*TPEP*flexCP 仅在 DeepSeek-V2-LiteEPCP上验证GLM-5MLA DSA否是否DSA indexer 在 TP 下 fail loudCP 未接线——仅 FSDP/EPGemma-4-26B-A4BGQA否是否注意力 TP 不可行num_global_key_value_heads2——仅 FSDP/EP注意MoE 路径的 PP 尚未接线见「后续工作」。并行化流水线SPMD 类型系统如何串起 4D 并行后端使用 TorchTitan SPMD 类型系统并非一句空话。parallelize.py 的parallelize_hf_transformers是并行化总入口其流程源码 docstring 明示为解绑 embedding/lm_head对tie_word_embeddingsTrue的模型如 Gemma4先克隆一份 lm_head 权重因为 FSDP2 不允许同一参数被两个 FSDP group 管理。MoE 替换若存在 MoE 层调用build_and_swap_native_moe换入 Titan MoE设置_sharding_config但不立即并行化。HF 模块转换为 Module protocol通过__class__交换把 HF 的nn.Module转成 TorchTitanModulemodule_conversion.py。设置 ShardingConfighf_sharding.py 的set_hf_sharding_configs依据模块角色embedding、attention、MLP、norm 等为每个模块声明式设置状态/激活张量的 TP 分片布局例如_hf_colwise_config列切分输出保留 HF batch 维、_hf_rowwise_config行切分可选 sequence-parallel 输出、flex kernel 的本地 SPMD region。单次model.parallelize(parallel_dims)TP、EP、MoE 全部在这一调用中完成状态分片与前向包装。AC / compile / FSDP先按需应用激活检查点ac_config.build(...).apply(model)再apply_compile编译整个 transformer 块含 Titan MoE源码注明此前 MoE-only 编译的 workaround 已因上游修复而废弃最后apply_fsdp应用 FSDP2。FSDP 的细节也值得注意当 EP 开启ep_degree 1时apply_fsdp使用带shard_placement_fn的 flat FSDP——专家参数路由到dp_mod_ep_mesh组合 EP 网格其余参数路由到普通dp_mesh在一个fully_shard调用内完成避免嵌套 FSDP hooks 在 recompute 阶段引发 SAC 算子计数不匹配EP 开启时还会设置显式的前向/后向预取链规避 EP 中 D2H 同步对 FSDP 隐式预取的干扰。PP 侧由 pipeline.py 的pipeline_hf_transformers负责generate_llm_fqn_per_model_part按embedding 计 1 层、normlm_head 计 1 层的权重把层列表均分到各 stage生成每个 stage 的模块名列表例如 2 个 stage 3 层时 embeddings 视为 2 层参与分配并复用 TorchTitan 的 1F1B、DualPipe、ZBV 等 schedulePP 占位模块使用nn.Identity替代Noneinit_states会跳过nn.Identity的初始化每个 stage 还单独携带rotary_emb。验证体系测试矩阵与数值回归后端自带一整套针对性的集成测试是理解其能力边界的最佳入口tests/run_moe_tests.shMoE 并行度矩阵测试。默认NGPU8, STEPS200使用MODULEtransformers_modeling_backend CONFIGtransformers_modeling_backend_debugmodel_moe覆盖 FSDP-only、FSDPEP2/4、FSDPTP2/4、TPEP、EPcompile、PPEP、HSDPEP、CP、CPEP、CPTPEP 等 20 种组合EP 从 fsdp×tp 乘积中借用不额外乘算 WORLD_SIZE8 卡时 4 卡测试并行成对跑以减半墙钟时间。脚本尾部还有针对 8 个 MoE 模型的模型×并行度全量扫描需要联网下载 HF config离线环境可设SKIP_MODEL_SWEEP1跳过Gemma4 因 HF Hub 门控需HF_TOKEN且默认注释。tests/test_flex_attention.pypacked 序列跨文档泄漏测试验证block_causal文档掩码真的阻止跨文档注意力。tests/test_moe_parallelism.py/tests/test_hf_sharding.py/tests/test_hf_build_forward_inputs.py/tests/cp_pp_numerical.py分别覆盖 MoE 并行、ShardingConfig TP 布局、前向输入构造与 CP/PP 数值一致性。MODEL_COMPATIBILITY.md上文已述的 MoE 数值等价性追踪表是替换 MoE 是否引入数值漂移的权威依据。已知问题与后续工作README 明确列出的已知问题当前版本尚未解决数值对齐使用 HF 建模时FSDP2对比FSDP2 PP2loss与grad_norm无法逐位一致但都能收敛而 TorchTitan 原生建模可以逐位一致。README 判断元凶大概率是加载seed_checkpoint时的register_buffer将在后续 PR 中解决。MFU 差距HF 建模的 MFU模型利用率低于 TorchTitan 原生建模。README 列出的后续工作MoE 路径缺失的PP 支持上文矩阵中 PP 未接线的正式出处加载 HF 权重当前从零初始化训练不过sft_fullflavor 已通过initial_load_in_hfTrue支持部分 HF 预训练权重加载见 config_registry.py增加LoRA 支持支持Titan RL仓库中 experiments/rl 强化学习训练管线。小结transformers_modeling_backend的价值在于把HuggingFace 模型生态与TorchTitan 原生训练能力解耦一方面用户只需给出hf_modelID 就能让任意 HF CausalLM 模型进入 FSDP/TP/PP/CP torch.compile的训练流水线模型适配属性映射、初始化对齐、模块规整全部由HFTransformerModel自动完成另一方面注意力统一走 FlexAttentioncausal/block_causal两种掩码DSA 模型自动退化为稠密掩码MoE 统一替换为 TorchTitan grouped-experts MoE 以获得 EP 与grouped_mm加速。对于想用 HF 模型快速验证并行策略、或需要在统一框架内训练多架构模型的研究与工程团队这是一个开箱即用的入口——同时也要留意其已知的 PP×MoE 缺失、HF 建模 MFU 偏低等边界。【免费下载链接】torchtitanA PyTorch native platform for training generative AI models项目地址: https://gitcode.com/GitHub_Trending/to/torchtitan创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
使用 TorchTitan HuggingFace Transformers 建模后端：以 4D 并行与 torch.compile 训练任意 HF 模型
使用 TorchTitan HuggingFace Transformers 建模后端以 4D 并行与 torch.compile 训练任意 HF 模型【免费下载链接】torchtitanA PyTorch native platform for training generative AI models项目地址: https://gitcode.com/GitHub_Trending/to/torchtitan本文介绍 TorchTitan 仓库中的transformers_modeling_backendHuggingFace Transformers 建模后端它允许把 HFHuggingFace生态中任意transformers模型直接接入 TorchTitan 的训练框架在原生 PyTorch 的 FSDP/TP/PP/CP 四维并行与torch.compile之下完成大规模训练。读完本文你将掌握该后端的启动配置、模型/并行度支持矩阵、FlexAttention 与 MoE 替换两大核心机制以及它当前的已知限制。该后端位于 torchtitan/experiments/transformers_modeling_backend/其设计目标写在 README.md 的第一句让 HF transformers 模型能够以4D 并行FSDP/TP/PP/CPtorch.compile的方式被训练并且整个后端建立在 TorchTitan 的SPMD 类型系统之上。快速开始最小可运行配置环境依赖后端对transformers库版本有明确要求README 中写明transformers5.9.0切换模块与配置后端通过 TorchTitan 的--module/--config机制接入无需改动训练主程序。按 README 给出的 diff把默认的 Llama3 模块替换为建模后端模块并选择debugmodel调试配置... - --module llama3 --module transformers_modeling_backend --config transformers_modeling_backend_debugmodel ...启动训练README 给出的启动命令需先准备好 tokenizer 资源LOG_RANK7 MODELtransformers_modeling_backend CONFIGtransformers_modeling_backend_debugmodel ./run_train.sh --compile.enable要点LOG_RANK7指定从最后一个 rank 读取日志。对于流水线并行PP场景loss 只存在于最后一个 pipeline stagerank 0 打印的是-1.0占位值因此要选择NGPU - 1这个 rank 才能看到真实 loss这一约定同样体现在 run_moe_tests.sh 的LOG_RANK$((ngpu - 1))中。--compile.enable开启torch.compile这也是该后端支持的核心特性之一。README 明确提示运行前必须已创建好 tokenizer。配置体系从hf_model到六个预置 flavor扩展配置入口TransformersBackendConfig后端在 config_registry.py 中定义了TransformersBackendConfig(Trainer.Config)它在 TorchTitan 标准Trainer.Config之上新增了一个核心字段dataclass(kw_onlyTrue, slotsTrue) class TransformersBackendConfig(Trainer.Config): hf_model: str HuggingFace model ID (e.g., Qwen/Qwen2.5-7B)hf_model即 HF 模型仓库 ID如Qwen/Qwen3-30B-A3B训练时通过--hf_model id或预置 config 指定。模型的实际架构、尺寸、MoE 参数全部由 HF 侧的AutoConfig.from_pretrained(hf_model)加载见 model.py 中的update_from_config因此更换模型通常只需更换hf_model无需修改任何并行化代码。六个预置 flavor仓库在 config_registry.py 中预置了以下配置flavor 定义在init.py 的flavors字典中配置名flavor用途transformers_modeling_backend_debugmodeldebugmodel稠密模型调试dim256, 2 层, 16 头causal掩码默认hf_modelQwen/Qwen3-4B-Instruct-2507transformers_modeling_backend_debugmodel_moedebugmodel_moeMoE 调试dim2048, 4 层, 8 专家, top-2默认hf_modelQwen/Qwen3-30B-A3Btransformers_modeling_backend_fullfull稠密全尺寸训练不覆盖任何 HF 配置全部取自 HF configtransformers_modeling_backend_full_moefull_moeMoE 全尺寸训练dim2048, 48 层, 128 专家, top-8transformers_modeling_backend_sft_fullsft_fullSFT 全尺寸block_causal掩码 initial_load_in_hfTrue加载 HF 预训练权重transformers_modeling_backend_sft_debugmodelsft_debugmodelSFT 调试dim256, 2 层,block_causal掩码用tests/assets/sft_test/data.json做 chat 数据其中debugmodel系列是 CI 与日常验证的主力hf_model使用真实 HF 模型 ID如 Qwen3-4B但 flavor 显式覆盖dim/n_layers/n_heads等尺寸把模型缩小为可快速训练的最小架构从而在 debug 规模上验证完整训练链路。两层配置模型TitanModelConfig/TitanMoeModelConfigflavor 的核心是init.py 中定义的两层 dataclass。HFTransformerModel.Config是桥接 TorchTitan 与 HF 的配置外壳内部持有model_config而TitanModelConfig稠密与其子类TitanMoeModelConfigMoE才是真正的模型参数描述。字段设计原则源码注释明确说明凡是能映射到 HF config key 的字段默认值一律为None这样AutoConfig.from_pretrained加载的值会被保留——如果给了非None默认值就会覆盖 HF 配置、强行改变模型架构例如rope_theta只有显式设置如 debugmodel 的尺寸覆盖才会注入。而 TorchTitan 独有的字段则保留具体默认值。TitanModelConfig关键字段字段默认值说明dim/n_layers/n_heads/n_kv_heads/vocab_size/intermediate_sizeNone映射到 HF 的hidden_size/num_hidden_layers/num_attention_heads/num_key_value_heads/vocab_size/intermediate_size映射表见 model.py 的_TT_TO_HF_MAPPINGSnorm_eps/rope_thetaNone映射到 HF 的rms_norm_eps/rope_thetamultiple_of256TorchTitan 独有显式覆盖dim时用于推导 FFN 尺寸的对齐因子ffn_dim_multiplierNoneTorchTitan 独有推导 FFN 尺寸的倍率max_seq_len2048训练时由training.max_context_length覆盖depth_initTrue是否启用深度相关的初始化缩放attn_mask_typecausalFlexAttention 掩码类型见下文「注意力」一节TitanMoeModelConfig在稠密基础上增加字段默认值说明num_experts/num_experts_per_tok/moe_intermediate_sizeNone路由专家总数 / top-k / 每个专家 MLP 的中间维度decoder_sparse_stepNone每隔 N 层把稠密 MLP 替换为 MoE 块1 每层都是 MoEnorm_topk_probNone是否把 top-k 路由分数归一化为和为 1HFnorm_topk_prob约定num_nextn_predict_layersNoneDeepSeek V3 风格的多 token 预测MTP头数None 或 0 表示禁用experts_implementationgrouped_mm专家前向 kernelgrouped_mm融合快速路径、batched_mm、eagerHF 原始 for 循环数值参考、native使用 HF 模型自带 kernel不可设置 kernel 的模型唯一合法选项load_balance_coeff1e-3无辅助损失的 MoE 负载均衡步长None禁用comm_backendstandard专家并行的 token 分发后端standardPyTorch all-to-all、deepepH100/NVLink 的 DeepEP、hybridepGB200/NVLink72 的 HybridEP核心机制一模型适配层HFTransformerModel后端在 model.py 中通过HFTransformerModel(BaseModel)把任意PreTrainedModel包装成 TorchTitan 认可的BaseModel。这一层解决的是两类接口不匹配1. 架构发现与模型类解析__init__中先取config.architectures[0]作为模型类名从globals()或transformers包动态导入失败时再通过model_type → MODEL_FOR_CAUSAL_LM_MAPPING_NAMES的 Auto 映射兜底解析。因此该后端面向CausalLM类模型且支持transformers包中注册的任意架构。2. 属性映射TT 风格访问 HF 模型_TT_TO_HF_MAPPINGS定义了 TorchTitan 属性名与 HF 属性名的对应关系如dim → hidden_size、n_kv_heads → num_key_value_heads。Config._create_getter_setter_dynamically会动态创建 property让self.dim读写self.hidden_size从而对上层暴露统一的 TorchTitan 风格接口。update_from_config在加载 HF config 后通过vars(hf_model_config)复制所有属性包括to_dict()会漏掉的运行时计算属性例如 DeepSeek V3 的qk_head_dim并保留attribute_map例如 DeepSeek V3 把num_local_experts映射到n_routed_experts最后重新注入用户在 flavor 中显式设置的覆盖值保证debugmodel的尺寸覆盖不被 HF config 冲掉。值得注意的两个细节均有源码注释佐证远程 config 黑名单_REMOTE_CONFIG_DENYLIST {deepseek_v2, deepseek_v3}。当模型仓库自带旧版远程 config 实现时混用会丢失本地模型代码期望的兼容属性因此这两个模型类型强制使用本地transformers的 config/model 对trust_remote_codeFalse。权重初始化补丁_patch_hf_llama_like对 HF 的_init_weights/_initialize_weights做 monkey-patch使其匹配 TorchTitan 的初始化方案attention/MLP 深度相关初始化init_std 0.02 / (2 * (layer_idx 1)) ** 0.5、MLA 低秩投影q_a_proj等按std0.02初始化、专家 3D 参数trunc_normal_等。这一步是 checkpoint 在 TorchTitan 格式与 HF 格式之间逐位可复现转换的前提。3. 结构规整SliceableModuleDictHF 模型用nn.ModuleList保存 decoder layers而 TorchTitan 的 PP/FSDP 期望ModuleDict语义。__init__末尾把nn.ModuleList转换为SliceableModuleDictkey 为字符串化的层索引见 model.py 的SliceableModuleDict既保持 state dict key 与 checkpoint 一致又支持按 slice 取层供 PP 切分使用。模型顶层则通过named_children()暴露tok_embeddings / layers / norm / lm_head / rotary_emb五个统一子模块兼容 Phi 的final_layernorm等命名差异。核心机制二注意力统一走 FlexAttention后端把所有HF 注意力路由到torch.nn.attention.flex_attention这是它区别于直接塞给 HF 训练的关键设计README「Attention」一节 model.py 的_configure_hf_attention注册自定义实现名_ATTN_IMPLEMENTATION flex_torchtitan到 HF 的AttentionInterface._global_mapping并设置config._attn_implementation从而绕过 HF 每个模型自带的_supports_flex_attn门禁统一走_flex_attention_torchtitanshim。掩码由attn_mask_type选择causal普通因果掩码get_causal_mask_mod()block_causal因果 同文档约束and_masks(get_causal_mask_mod(), get_document_mask_mod(positions))用于 packed / SFT 序列——README 明确说明is_causal无法表达跨样本掩码packed 训练必须用block_causal。掩码构建在get_attention_masks()中完成BLOCK_SIZE128 的 flex BlockMask并通过 HF 常规的attention_mask参数透传HF 对已是 4D/BlockMask 的输入原样返回。DSA 例外DeepSeek 稀疏注意力如 GLM-5model_typeglm_moe_dsa的 indexer 子注意力把掩码当稠密加法张量用对掩码调用.dim()、做逐元素相加flexBlockMask无法满足因此_uses_dsa()会检测index_topk配置并改用_build_dense_attention_mask()生成稠密 4D 加法掩码flex 以score_mask消费它causal/block_causal语义保持一致。在并行层面注意力运行在FlexInnerAttention之上HFFlexKernel把 flex HOP 包装成 TorchTitanModule在声明式 TP 下 q/k/v 携带 head 分片的 SPMD 标注输出获得对应标注CP 下则由 parallelize.py 的_wrap_flex_kernel_cp在每个 flex kernel 内部沿 CP 轴对 k/v 做 all-gather对应 TorchTitan 的flex_cp_allgather路径q 保持 seq 分片从而每个 rank 只对本地 seq 计算注意力、消费 Q 分片 / KV 全量的 BlockMask。跨文档泄漏防护有专门测试验证见 tests/test_flex_attention.py扰动前一个文档的 key/value 必须让后一个文档的输出逐位不变。核心机制三MoE 块替换为 TorchTitan grouped-experts MoE对于 MoE 模型后端并不直接训练 HF 自带的 MoE 实现而是把 MoE 块整体替换为 TorchTitan 的 grouped-experts MoEREADME「MoE」一节从而获得专家并行EP与grouped_mm融合快速路径。替换逻辑在 moe_replacement.py分两阶段Phase 1init 时prepare_native_moe_configs在 meta 设备上探测 HF MoE 块Qwen3MoeSparseMoeBlock等从gate/router、experts与 config 中解析出专家数兼容num_experts/n_routed_experts/num_local_experts三种命名、top-k、路由打分函数softmax / sigmoid通过e_score_correction_biasbuffer 探测、路由归一化、DeepSeek 的分组路由n_group/topk_group、shared expert 信息等构建 TitanMoE.Config存到层的_native_moe_config。Phase 2parallelize 时build_and_swap_native_moe在已知 EP/TP 度后设置set_moe_sharding_config、构建 Titan MoE、to_empty init_states物化参数再替换掉层的mlp属性。对 Gemma4 这种层级 MoErouter/experts 是 decoder layer 的兄弟节点、稠密 MLP 充当 shared expert替换后还会删除 HF 的独立 MoE 路径模块避免参数重复注册。替换的保真度由 MODEL_COMPATIBILITY.md 记录以未修改的 HF forward 为 gold standard在真实隐藏维度、缩减专家数E8、单层配置下比较输出。结果摘要KL 散度阈值PASS 1e-6WARN 1e-3FAIL ≥ 1e-3模型model_type状态KL div说明Qwen3-30B-A3Bqwen3_moePASS5.31e-07差异仅剩 dispatcher 精度OLMoE-1B-7BolmoePASS4.04e-07同上DeepSeek-V2-Litedeepseek_v2PASS2.64e-07同上Mixtral-8x7BmixtralWARN4.62e-05dispatcher 精度DeepSeek-V3deepseek_v3WARN3.66e-04dispatcher 精度GLM-4.7glm4_moeWARN5.06e-05dispatcher 精度GLM-5 (DSA)glm_moe_dsaWARN2.67e-04dispatcher 精度Gemma-4-26Bgemma4_textFAIL4.19e-03激活函数GeGLU vs SwiGLU 路由器架构差异文档还给出了 WARN 模型的三大数值差异根源均为 TitanLocalTokenDispatcher与 HF 的独立差异accumulation dtypebf16 vs f32、accumulation orderscatter_add按专家序 vs HF 按 token 序求和、topk 排序sortedFalsevssortedTrue三者修正后 max_diff0.00。这些结论为理解替换 MoE 是否安全提供了量化依据。支持的模型与并行度矩阵README 列出该后端经过验证的模型清单均为真实 HF 仓库 ID稠密模型支持 FSDP / CP / TP / PP /torch.compilemeta-llama/Llama-3.2-1Bmicrosoft/phi-2Qwen/Qwen2.5-7Bmistralai/Mistral-7B-v0.1ByteDance-Seed/Seed-Coder-8B-InstructQwen/Qwen3-4B-Instruct-2507arcee-ai/AFM-4.5Bibm-granite/granite-3b-code-base-2kbaidu/ERNIE-4.5-0.3B-Base-PTkyutai/helium-1-preview-2ballenai/OLMo-7B-hfmistralai/Ministral-8B-Instruct-2410MoE 模型并行支持见下方矩阵Qwen/Qwen3-30B-A3Bqwen3_moeGQAmistralai/Mixtral-8x7B-Instruct-v0.1mixtralGQAallenai/OLMoE-1B-7B-0924olmoeGQAdeepseek-ai/DeepSeek-V2-Litedeepseek_v2MLAdeepseek-ai/DeepSeek-V3deepseek_v3MLAzai-org/GLM-4.7glm4_moeMLAzai-org/GLM-5glm_moe_dsaMLA DSA 稀疏注意力google/gemma-4-26B-A4B-itgemma4_textREADME 给出的 MoE 并行度矩阵debug 规模、2 步 loss 递减验证通过FSDP 与所有列出的轴均可组合模型attnTPEPCP备注Qwen3-30B-A3BGQA是是是完整矩阵最高 TPEPCPMixtral-8x7B、OLMoE-1B-7BGQA是是是TPEP、EPCPDeepSeek-V2-Lite、V3、GLM-4.7MLA是是是*TPEP*flexCP 仅在 DeepSeek-V2-LiteEPCP上验证GLM-5MLA DSA否是否DSA indexer 在 TP 下 fail loudCP 未接线——仅 FSDP/EPGemma-4-26B-A4BGQA否是否注意力 TP 不可行num_global_key_value_heads2——仅 FSDP/EP注意MoE 路径的 PP 尚未接线见「后续工作」。并行化流水线SPMD 类型系统如何串起 4D 并行后端使用 TorchTitan SPMD 类型系统并非一句空话。parallelize.py 的parallelize_hf_transformers是并行化总入口其流程源码 docstring 明示为解绑 embedding/lm_head对tie_word_embeddingsTrue的模型如 Gemma4先克隆一份 lm_head 权重因为 FSDP2 不允许同一参数被两个 FSDP group 管理。MoE 替换若存在 MoE 层调用build_and_swap_native_moe换入 Titan MoE设置_sharding_config但不立即并行化。HF 模块转换为 Module protocol通过__class__交换把 HF 的nn.Module转成 TorchTitanModulemodule_conversion.py。设置 ShardingConfighf_sharding.py 的set_hf_sharding_configs依据模块角色embedding、attention、MLP、norm 等为每个模块声明式设置状态/激活张量的 TP 分片布局例如_hf_colwise_config列切分输出保留 HF batch 维、_hf_rowwise_config行切分可选 sequence-parallel 输出、flex kernel 的本地 SPMD region。单次model.parallelize(parallel_dims)TP、EP、MoE 全部在这一调用中完成状态分片与前向包装。AC / compile / FSDP先按需应用激活检查点ac_config.build(...).apply(model)再apply_compile编译整个 transformer 块含 Titan MoE源码注明此前 MoE-only 编译的 workaround 已因上游修复而废弃最后apply_fsdp应用 FSDP2。FSDP 的细节也值得注意当 EP 开启ep_degree 1时apply_fsdp使用带shard_placement_fn的 flat FSDP——专家参数路由到dp_mod_ep_mesh组合 EP 网格其余参数路由到普通dp_mesh在一个fully_shard调用内完成避免嵌套 FSDP hooks 在 recompute 阶段引发 SAC 算子计数不匹配EP 开启时还会设置显式的前向/后向预取链规避 EP 中 D2H 同步对 FSDP 隐式预取的干扰。PP 侧由 pipeline.py 的pipeline_hf_transformers负责generate_llm_fqn_per_model_part按embedding 计 1 层、normlm_head 计 1 层的权重把层列表均分到各 stage生成每个 stage 的模块名列表例如 2 个 stage 3 层时 embeddings 视为 2 层参与分配并复用 TorchTitan 的 1F1B、DualPipe、ZBV 等 schedulePP 占位模块使用nn.Identity替代Noneinit_states会跳过nn.Identity的初始化每个 stage 还单独携带rotary_emb。验证体系测试矩阵与数值回归后端自带一整套针对性的集成测试是理解其能力边界的最佳入口tests/run_moe_tests.shMoE 并行度矩阵测试。默认NGPU8, STEPS200使用MODULEtransformers_modeling_backend CONFIGtransformers_modeling_backend_debugmodel_moe覆盖 FSDP-only、FSDPEP2/4、FSDPTP2/4、TPEP、EPcompile、PPEP、HSDPEP、CP、CPEP、CPTPEP 等 20 种组合EP 从 fsdp×tp 乘积中借用不额外乘算 WORLD_SIZE8 卡时 4 卡测试并行成对跑以减半墙钟时间。脚本尾部还有针对 8 个 MoE 模型的模型×并行度全量扫描需要联网下载 HF config离线环境可设SKIP_MODEL_SWEEP1跳过Gemma4 因 HF Hub 门控需HF_TOKEN且默认注释。tests/test_flex_attention.pypacked 序列跨文档泄漏测试验证block_causal文档掩码真的阻止跨文档注意力。tests/test_moe_parallelism.py/tests/test_hf_sharding.py/tests/test_hf_build_forward_inputs.py/tests/cp_pp_numerical.py分别覆盖 MoE 并行、ShardingConfig TP 布局、前向输入构造与 CP/PP 数值一致性。MODEL_COMPATIBILITY.md上文已述的 MoE 数值等价性追踪表是替换 MoE 是否引入数值漂移的权威依据。已知问题与后续工作README 明确列出的已知问题当前版本尚未解决数值对齐使用 HF 建模时FSDP2对比FSDP2 PP2loss与grad_norm无法逐位一致但都能收敛而 TorchTitan 原生建模可以逐位一致。README 判断元凶大概率是加载seed_checkpoint时的register_buffer将在后续 PR 中解决。MFU 差距HF 建模的 MFU模型利用率低于 TorchTitan 原生建模。README 列出的后续工作MoE 路径缺失的PP 支持上文矩阵中 PP 未接线的正式出处加载 HF 权重当前从零初始化训练不过sft_fullflavor 已通过initial_load_in_hfTrue支持部分 HF 预训练权重加载见 config_registry.py增加LoRA 支持支持Titan RL仓库中 experiments/rl 强化学习训练管线。小结transformers_modeling_backend的价值在于把HuggingFace 模型生态与TorchTitan 原生训练能力解耦一方面用户只需给出hf_modelID 就能让任意 HF CausalLM 模型进入 FSDP/TP/PP/CP torch.compile的训练流水线模型适配属性映射、初始化对齐、模块规整全部由HFTransformerModel自动完成另一方面注意力统一走 FlexAttentioncausal/block_causal两种掩码DSA 模型自动退化为稠密掩码MoE 统一替换为 TorchTitan grouped-experts MoE 以获得 EP 与grouped_mm加速。对于想用 HF 模型快速验证并行策略、或需要在统一框架内训练多架构模型的研究与工程团队这是一个开箱即用的入口——同时也要留意其已知的 PP×MoE 缺失、HF 建模 MFU 偏低等边界。【免费下载链接】torchtitanA PyTorch native platform for training generative AI models项目地址: https://gitcode.com/GitHub_Trending/to/torchtitan创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
