Megatron Core HybridModel 迁移实战指南从 GPTModel 检查点加载、转换与继续训练【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM本篇指南围绕 Megatron Core 的HybridModel展开完整讲解如何将一个基于GPTModel训练的同构 Transformer 模型迁移为模式驱动的混合架构包括替换模型入口、通过加载时翻译或离线转换两种方式复用分布式检查点、在保留架构的前提下继续训练或恢复训练以及流水线并行布局的重新表达。读完本文你将掌握--hybrid-layer-pattern模式语法、gpt_hybrid_conversion.py转换器的使用边界以及从 GPT 检查点无缝续跑 HybridModel 的完整实操路径。1. 什么是 HybridModel标准GPTModel的每个 decoder 层在一个层索引下同时包含 self-attention 子层与 MLP/MoE 子层。HybridModel则构建一个有序的层栈其中每个位置只代表一种层家族layer family层的顺序完全由--hybrid-layer-pattern描述。模式中的一个符号就是 HybridModel 的一层因此保留原始架构时一个 GPT transformer block 会变成两个 HybridModel 层。--hybrid-layer-pattern支持的符号定义对应 Symbols 类 中的LAYER_CONFIG_MAP符号层家族对应配置类MMamba-2 状态空间层MambaLayerConfigGGated Delta Network (GDN) 层GDNLayerConfig*标准 Self-attention 层AttentionLayerConfigDDeepSeek Sparse Attention (DSA) 层DSALayerConfig-稠密 MLP 层MLPLayerConfigEMixture-of-Experts (MoE) 层MoELayerConfigMulti-Latent Attention (MLA) 层MLALayerConfig\|流水线段边界不计入层数—/引入重复的 Multi-Token Prediction (MTP) 模式—说明迁移文档中的符号表只列出M/G/*/D/-/E六种。从 层工具源码 可见仓库还支持MLA符号另外_validate_pattern明确禁止同一模型中混用标准 Attention 与 MLA/DSA见 hybrid_layer_allocation.py。1.1 从 GPT 架构映射到 Hybrid 模式由于一个 GPT 层贡献恰好一个 attention 子模块和一个 MLP 子模块保留架构的映射关系如下源架构等价的 HybridModel 模式Hybrid 层数两个稠密 GPT block*-*-4两个全 MoE GPT block*E*E4例如源 GPT 第 0 层被拆分为 Hybrid 第 0、1 层它的 attention 参数进入第一个*MLP 参数进入第一个-或E源 GPT 第 1 层映射到下一对以此类推。参数按出现次序配对而不是按数值层号硬编码这正是gpt_hybrid_conversion.py中build_layer_index_mapping的核心逻辑第i个*位置取 GPT 第i层的 attention第i个-/E位置取 GPT 第i层的 MLP见 转换工具源码。模式还可以描述执行布局|标记流水线段边界/引入重复的 MTP 模式。例如*-*-|*-*-把 4 个 GPT 等价 block 放到两个流水线段上。分隔符不计入层数。从 parse_hybrid_pattern 的实现可见统一模式的完整语法为main_pattern/mtp_pattern/mtp_pattern/...且所有 MTP 子模式必须完全一致例如M*M*/MM/MM表示主解码器M*M*加 2 层 MTP 深度。1.2 HybridModel 提供的能力不同的层家族可以组合进同一个模型无需强制每个 decoder block 结构一致attention 与稠密/专家 MLP 层可以独立放置与配置Mamba、GDN、标准或 DeepSeek attention、稠密 MLP 与 MoE 层共享同一个模式驱动的模型接口HybridStack通过hybrid_stack_spec统一装配见 hybrid_layer_specs.py包含 Mamba 的模式可以用亚二次subquadratic序列混合与固定大小的循环推理状态替代部分二次复杂度 attention 层流水线与虚拟流水线切分可以直接用模型模式表达而不再依赖独立的层布局参数同一个模型抽象既能描述纯 Transformer重复*-、纯 Mamba 模型也能描述异构架构。需要强调的是这些能力并不自动带来吞吐或质量的提升。保留架构的*-或*E迁移应当验证数值等价性而引入了新层家族的模式应视为新架构需要独立做基准测试这一点在 HybridModel 类文档 中同样没有承诺任何性能收益。2. 如何转换检查点将GPTModel权重带入HybridModel运行有两条路径两者都停留在 Megatron 的分布式检查点distributed-checkpoint格式内不需要以 Hugging Face 作为中间格式并且后续加载时都可以跨不同的 TP/PP/EP/FSDP 布局重新分片方案 A——加载时翻译无需单独步骤直接用 hybrid 运行去加载 GPT 检查点。hybrid 模型会在加载期间把自己的检查点 state dict 重定向到 GPT 检查点的 key 上磁盘上不会写出第二份拷贝。同时支持torch_dist检查点与 Megatron-FSDP 的fsdp_dtensor检查点含优化器状态。该路径也支持包含 GPT 没有对应物的层家族的模式——例如 MambaM位置会保留自己的全新初始化。方案 B——离线转换为新检查点使用 tools/checkpoint/gpt_hybrid_conversion.py 写出一份独立的HybridModel检查点其 key 已经符合 hybrid 布局。适用于需要持久化 hybrid 检查点、需要一份保留架构的*-或*E拷贝、或希望在训练前先检查目标结果的场景。该路径只支持*-与*E两种布局。2.1 方案 A加载时翻译加载时翻译由 megatron/core/dist_checkpointing/gpt_checkpoint_interop.py 处理当非 hybridGPT检查点被加载进HybridModel运行时自动触发对于torch_dist运行的模型与优化器分片 state dict 会被改写成 GPT 检查点的同构层homogeneous-layer格式——把decoder.layers.g.mlp...中的层索引从存储 key 中剥离、将匹配的 GPT 层索引变为前置的分片轴恰好镜像TransformerBlock.sharded_state_dict在non_homogeneous_layersFalse时保存的格式见 retarget_sharded_state_dict_to_gpt_checkpoint 与 模块头注释对于fsdp_dtensor在 Torch DCP 规划之前把显式的参数名映射改写为 GPT 的 key见 retarget_fsdp_state_dict_to_gpt_checkpoint。检查点被直接读取权重与优化器状态会重新分片到当前 TP/PP/EP/ETP/FSDP 布局不运行任何转换工具磁盘上的 GPT 检查点永不被修改。反向不匹配会直接报错把 hybrid 运行保存的检查点加载进非 hybrid 运行会抛出RuntimeError并引导你使用 hybrid 训练入口。选择检查点语义当 hybrid 运行加载 GPT 检查点时必须设置--hybrid-layer-pattern以便把检查点层与 hybrid 层位置配对。将--load或--pretrained-checkpoint指向 GPT 检查点根目录见第 3 节。--finetune可选并保留其常规的检查点加载含义GPT-to-Hybrid 翻译不会替用户选择它不带--finetune时直接--load会按常规检查点与并行布局兼容规则恢复迭代、优化器、调度器、RNG 与 rerun 状态带--finetune时迭代、调度器、RNG 与 rerun 状态全部从头开始。模型权重仍会加载翻译后的优化器状态也会加载除非设置--no-load-optim现有的--pretrained-checkpoint回退路径当--load目录中没有检查点时使用微调语义。需要完整恢复语义时请直接使用--load。默认情况下GPT 运行的优化器状态也会被翻译并加载——Adam 动量与 fp32 master 参数会随 attention 与 MLP 层一起带入从而支持保留架构的继续训练。传入--no-load-optim可跳过这一步让每一层的优化器状态从头开始。对于torch_dist加载优化器状态要求 GPT 检查点使用 model-space 的分布式优化器格式fully_reshardable或fully_sharded_model_space即用--dist-ckpt-optim-fully-reshardable保存。bucket-space 格式按扁平 buffer 布局组织优化器状态 key额外的 hybrid 层会打乱该布局因此运行会报错并提示你重新保存检查点或传入--no-load-optim。对于 Megatron FSDP请使用--ckpt-format fsdp_dtensor保存与加载。加载器会重定向显式的 DTensor 模型 key 以及分布式优化器使用的模型参数名因此模型权重与优化器状态可以在自动 GPT-to-Hybrid 加载期间跨不同的 FSDP、TP、EP、ETP 布局重新分片。该路径面向 Megatron FSDPTorch FSDP2 的torch_dcp格式不受自动翻译支持。没有 GPT 对应物的层例如 MambaM位置在检查点中没有优化器状态其动量从头开始运行会打印一条警告说明受影响层数。支持的模式与 key 映射主模式/MTP 后缀之前的部分忽略|流水线分隔符可包含符号权重来源*配对 GPT 层的self_attention子模块-或E配对 GPT 层的mlp子模块MoE 张量也位于mlp.*之下M无 GPT 来源Mamba 层保留全新初始化decoder.final_norm从 GPT 的decoder.final_layernorm加载embedding 与输出权重原样复制。由于每个 GPT 层恰好提供一个 attention 与一个 MLP 子模块加载器会拒绝以下模式包含 MTP 层/...后缀——它们没有 GPT 源权重使用无法翻译的层类型如 GDNG或 DeepSeek Sparse AttentionD——其权重布局与 GPT attention 不同在同一模式中混用稠密-与 MoEEMLP 位置*与 MLP 位置数量不相等或为零。检查点的num_layers必须等于模式中*位置的数量不匹配会被拒绝。警告优化器翻译覆盖 Adam 动量与 fp32 master 参数。仅加载权重请传--no-load-optim需要迭代、调度器、RNG 与 rerun 状态重新开始而非恢复时使用--finetune。2.2 方案 B用gpt_hybrid_conversion.py离线转换选择保留架构的模式对于含N个 GPT 层的源检查点稠密 GPT 模型使用*-重复N次每层 MLP 均为 MoE 的 GPT 模型使用*E重复N次。转换器按出现次序而非仅按数值层号映射参数源参数目标参数GPT 第i层的 attention第i个*层GPT 第i层的 MLP 或 MoE第i个-或E层embedding 与输出权重原样复制模型角色不变decoder.final_layernorm重命名为decoder.final_norm这一映射逻辑在 convert_gpt_to_hybrid 中实现decoder.layers.下的 key 先按is_attention_param/is_mlp_param归类再通过replace_layer_num把层号改写为映射后的 hybrid 层号final_layernorm统一替换为final_norm。MoE 张量因两侧都位于mlp.{router,experts,shared_experts}.*之下而可原样往返无需折叠专家或重初始化 router见转换工具头注释。检查前置条件源检查点必须使用以下分布式检查点格式之一torch_distfsdp_dtensor优先把包含latest_checkpointed_iteration.txt的顶层检查点根目录作为--load-dir。如果--load-dir直接指向含metadata.json的目录转换器会写出一份没有 tracker 文件的扁平目标而标准训练入口期望检查点根目录与 tracker。在仓库根目录、Megatron 环境下运行转换器即可。普通python进程就足够不需要torchrun也不需要 GPU。工具会在 CPU 上聚合完整的逻辑张量因此宿主机需要有足够内存容纳未分片的源与目标模型 state dict。目标目录必须与源目录不同。警告这是权重转换不是可恢复的完整训练状态转换。分片的优化器、RNG、rerun 以及 Transformer Engine_extra_state张量不会被转换。部分非张量条目可能残留在common.pt中但它们不构成转换后的优化器或 RNG 状态。转换后的模型请使用全新的优化器与 RNG 状态启动。运行转换以下示例将一个四层稠密 GPT 模型转换为等价的八层 HybridModel模式*-*-*-*-uv run python tools/checkpoint/gpt_hybrid_conversion.py \ --direction gpt-to-hybrid \ --load-dir /path/to/gpt-checkpoints \ --save-dir /path/to/hybrid-checkpoints \ --hybrid-layer-pattern *-*-*-*- \ --reset-iterations务必给模式加引号因为*和|对 shell 有特殊含义。--input-format auto与--output-format auto是默认值工具自动检测源后端并写出相同后端也可用--input-format/--output-format显式指定torch_dist或fsdp_dtensor。--reset-iterations会重置检查点迭代数、已消费样本计数以及缓存的train_iters与train_samples当新运行需要保留这些调度元数据时请省略它。模式中*位置数与-或E位置数都必须等于源 GPT 层数。模式校验器会拒绝 GDN、DSA 以及稠密/MoE 混合布局对应 validate_pattern_gpt_compatible 的实现。当检查点中缓存了训练参数时工具还会拒绝交错 MoE、实验性或线性 attention、异构 block 规格、Multi-Latent Attention 与 MTP 检查点对应 _GPT_COMPAT_REJECT_FIELDS 中的字段清单moe_layer_freq、experimental_attention_variant、linear_attention_freq、heterogeneous_block_specs、heterogeneous_layers_config_path、multi_latent_attention、mtp_num_layers等。当common.pt中没有缓存的args、或较旧检查点缺少某个字段时该源特性校验不完整需要手动核对这些特性。转换只识别标准的 attention 与 MLP/MoE state-dict key其他层局部张量会被省略。文档化的hybrid_stack_spec使用 Transformer Engine 的 fused layernorm/linear 布局本地或其他非 TE 的源布局需要兼容的自定义 Hybrid stack 与 key 转换。务必执行下文所述的严格加载检查。不要在转换期间追加 MTP/...后缀。转换器只映射第一个/之前的主模式因此不会创建 MTP 参数。当源路径是含latest_checkpointed_iteration.txt的检查点根目录时输出会包含一个迭代目录与匹配的 tracker 文件。保存的全形状张量可以由后续 Megatron 加载按不同的 tensor、pipeline、expert 或 FSDP 布局重新分片。此外该工具还支持反向hybrid-to-gpt方向从模块头注释可见完整用法此时 SSM 层被丢弃并打印警告旧版mp_rank_XX/model_optim_rng.pt布局不受支持需要先转成torch_dist。3. 如何训练模型3.1 更新训练命令从训练 GPT 模型的命令出发做以下改动将pretrain_gpt.py替换为pretrain_hybrid.py删除--num-layers改用与转换时相同的、有序的主层符号序列。可以添加或移动流水线|分隔符。命令行解析器会从模式中推导num_layers见 arguments.py当--hybrid-layer-pattern存在时num_layers自动取get_hybrid_total_layer_count的结果并忽略--num-layers用--spec megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec选择 HybridModel 栈规格。从 hybrid_builder 的实现可以看到该--spec通过import_module动态导入未提供时会直接报错把检查点输入指向预训练权重并把新训练的检查点写到独立目录方案 A加载时翻译将--load直接指向GPT检查点以获得恢复语义或用--pretrained-checkpoint获得微调语义。优化器状态默认加载只有想要全新优化器时才加--no-load-optim。无需离线转换翻译过程也不要求--finetune。方案 B离线转换将--pretrained-checkpoint指向转换后的hybrid检查点并将--ckpt-format设置为转换器的torch_dist或fsdp_dtensor输出格式。一个最小的方案 A 迁移——直接加载 GPT 检查点及其优化器状态进行保留架构的继续训练——如下- torchrun --nproc_per_node8 pretrain_gpt.py \ - --num-layers 4 \ - --load /path/to/gpt-checkpoints \ - --save /path/to/gpt-checkpoints torchrun --nproc_per_node8 pretrain_hybrid.py \ --hybrid-layer-pattern *-*-*-*- \ --spec megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec \ --load /path/to/gpt-checkpoints \ # 首次启动之后切换到 save 目录 --save /path/to/new-training-checkpoints方案 B 则改为把--pretrained-checkpoint指向转换后的 hybrid 检查点。保持既有的架构、优化器、精度、数据与基础 TP/DP/EP/CP 参数不变除非本指南指明需要改动。模式驱动的流水线布局与 GPT 专属的数据集特性需要另行评估。警告pretrain_hybrid.py在设置--fim-data时不会选择GPTFIMDataset。使用 fill-in-the-middle 数据的 GPT 训练流程在迁移前需要自定义数据集路径或等价的 Hybrid 入口支持。当--load目录为空时--pretrained-checkpoint以微调语义加载预训练权重迭代从零开始、RNG 状态不恢复。方案 A 下优化器状态仍会热启动除非设置--no-load-optim方案 B 总是以全新优化器开始。任务向--load写入检查点后后续启动会正常恢复新的 HybridModel 训练状态。3.2 从零开始训练若要从头初始化每一层使用相同的pretrain_hybrid.py、--hybrid-layer-pattern与--spec参数但省略--pretrained-checkpoint。把--load与--save指向新运行目录以便后续启动恢复。与检查点转换不同从零训练可以使用所选 HybridModel 栈规格支持的所有兼容层家族并且可以包含 MTP 后缀例如M*M*/MM/MM。模式约束仍然适用——例如标准 attention*与 DSAD不能出现在同一个模型中。3.3 注意层索引扩展任何按 decoder 层索引的下标、映射或回调都必须改用 HybridModel 索引。以*E*E为例源第 0 层的 attention 是 Hybrid 第 0 层、其 MoE 是 Hybrid 第 1 层源第 1 层的 attention 是 Hybrid 第 2 层、其 MoE 是 Hybrid 第 3 层。只关注 attention 的列表需要为中间的 MLP 或 MoE 位置补上非活动条目。这一每类层各自编号的对应关系正是get_layer_maps_from_layer_type_listhybrid_layer_allocation.py所建立的每个全局层索引都被映射到该符号第几个出现的局部索引。3.4 配置流水线并行对于流水线并行直接添加|分隔符而无需改变有序层符号。例如转换后的模式*-*-*-*-可以按两个流水线段训练为*-*-|*-*-。管道分隔的段数必须能被--pipeline-model-parallel-size整除该约束在 select_pipeline_segment 中强制校验。模式取代了传统的流水线布局控制参数。请删除--num-layers-per-virtual-pipeline-stage--num-virtual-stages-per-pipeline-rank--pipeline-model-parallel-layout--account-for-embedding-in-pipeline-split--account-for-loss-in-pipeline-split当模式包含|时还需删除--decoder-first-pipeline-num-layers--decoder-last-pipeline-num-layers虚拟流水线切分改用额外的管道分隔段来表达。注意声明式的HybridModelBuilder目前拒绝虚拟流水线并行pretrain_hybrid.py的 CLI builder 支持 pipe 定义的虚拟 stage但自定义 builder 用户必须避免 VPP或使用显式支持它的路径。从源码看无|模式 pp_size 1的情况仍保留旧的自动切层逻辑按num_layers % pp_size均匀切分或通过--num-layers-in-first/last-pipeline-stage不均匀切分但会打印弃用警告建议显式加|分隔符select_pipeline_segment。3.5 更新自定义 provider 与转换映射自定义 provider 与转换映射还需适配以下 API 与 state-dict 差异构建或注册HybridModel而非GPTModel提供hybrid_stack_spec而非 GPT transformer-layer spec将程序化num_layers设为主模式中的层符号数与 CLI 路径不同自定义 provider 可能不会自动推导它把 attention 与 MLP/MoE 参数映射到各自独立的 Hybrid 层索引在 HybridModel 映射中使用decoder.final_norm而非decoder.final_layernorm把 attention-window schedule 等逐层设置扩展到完整的 HybridModel 模式。3.6 在扩展规模前验证在启动长任务之前严格加载转换后的检查点确认没有模型 key 或张量形状缺失/意外即开启 strict-load对*-或*E迁移在固定 batch 上把 logits 与源 GPT 模型在预期精度容差内做对比运行若干训练迭代检查 loss、梯度范数与逐层参数计数保存并重新加载一个新检查点确认新优化器与 RNG 状态能正确恢复。4. 仓库中的验证与测试资产以上迁移路径并非孤立文档仓库提供了多层次的测试与可运行样例用于印证转换工具单元测试tests/unit_tests/tools/checkpoint/test_gpt_hybrid_conversion.py 验证模式解析含 MTP/pipe 分隔符剥离、GPT↔Hybrid 层索引映射、final_layernorm↔final_norm重命名、SSM 参数初始化形状与 dtype以及 GPT→Hybrid→GPT 往返后 attention/MLP 权重保持test_gpt_hybrid_conversion_parallelism.py 覆盖并行布局相关行为加载时互操作测试tests/unit_tests/dist_checkpointing/models/test_gpt_hybrid_interop.py 直接验证gpt_checkpoint_interop.py的加载时翻译逻辑功能测试样例tests/functional_tests/test_cases/hybrid/下包含大量 HybridModel 训练/推理配置例如hybrid_mr_mcore_te_tp1_pp4_cp1_dgx_a100_1N8G、hybrid_mr_mcore_te_tp1_pp2_vpp2_*演示 pipe 定义 VPP、hybrid_static_inference_tp1_pp1_2B_logitsmatch演示 logits 一致性校验等可作为迁移后训练配置的参考模板Mamba 示例脚本examples/mamba/train.sh 展示了--hybrid-layer-pattern与--spec在实际训练命令中的用法。5. 总结从GPTModel到HybridModel的迁移在 Megatron Core 中是零中间格式的要么在加载时由 gpt_checkpoint_interop.py 自动把 hybrid 的 sharded state dict 重定向到 GPT 检查点的同构层格式支持torch_dist与fsdp_dtensor含优化器状态要么用 gpt_hybrid_conversion.py 离线生成独立的 hybrid 检查点。无论走哪条路模式中的*与-/E都按出现次序与 GPT 层一一配对M位置保持全新初始化G/D与 MTP 后缀因缺乏 GPT 源权重而被拒绝。训练侧只需把入口换成pretrain_hybrid.py、用--hybrid-layer-pattern取代--num-layers与各类流水线布局参数并在规模化之前完成严格加载与数值等价性验证——这既是迁移的安全边界也是充分发挥 HybridModel 架构灵活性的起点。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
Megatron Core HybridModel 迁移实战指南：从 GPTModel 检查点加载、转换与继续训练
Megatron Core HybridModel 迁移实战指南从 GPTModel 检查点加载、转换与继续训练【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM本篇指南围绕 Megatron Core 的HybridModel展开完整讲解如何将一个基于GPTModel训练的同构 Transformer 模型迁移为模式驱动的混合架构包括替换模型入口、通过加载时翻译或离线转换两种方式复用分布式检查点、在保留架构的前提下继续训练或恢复训练以及流水线并行布局的重新表达。读完本文你将掌握--hybrid-layer-pattern模式语法、gpt_hybrid_conversion.py转换器的使用边界以及从 GPT 检查点无缝续跑 HybridModel 的完整实操路径。1. 什么是 HybridModel标准GPTModel的每个 decoder 层在一个层索引下同时包含 self-attention 子层与 MLP/MoE 子层。HybridModel则构建一个有序的层栈其中每个位置只代表一种层家族layer family层的顺序完全由--hybrid-layer-pattern描述。模式中的一个符号就是 HybridModel 的一层因此保留原始架构时一个 GPT transformer block 会变成两个 HybridModel 层。--hybrid-layer-pattern支持的符号定义对应 Symbols 类 中的LAYER_CONFIG_MAP符号层家族对应配置类MMamba-2 状态空间层MambaLayerConfigGGated Delta Network (GDN) 层GDNLayerConfig*标准 Self-attention 层AttentionLayerConfigDDeepSeek Sparse Attention (DSA) 层DSALayerConfig-稠密 MLP 层MLPLayerConfigEMixture-of-Experts (MoE) 层MoELayerConfigMulti-Latent Attention (MLA) 层MLALayerConfig\|流水线段边界不计入层数—/引入重复的 Multi-Token Prediction (MTP) 模式—说明迁移文档中的符号表只列出M/G/*/D/-/E六种。从 层工具源码 可见仓库还支持MLA符号另外_validate_pattern明确禁止同一模型中混用标准 Attention 与 MLA/DSA见 hybrid_layer_allocation.py。1.1 从 GPT 架构映射到 Hybrid 模式由于一个 GPT 层贡献恰好一个 attention 子模块和一个 MLP 子模块保留架构的映射关系如下源架构等价的 HybridModel 模式Hybrid 层数两个稠密 GPT block*-*-4两个全 MoE GPT block*E*E4例如源 GPT 第 0 层被拆分为 Hybrid 第 0、1 层它的 attention 参数进入第一个*MLP 参数进入第一个-或E源 GPT 第 1 层映射到下一对以此类推。参数按出现次序配对而不是按数值层号硬编码这正是gpt_hybrid_conversion.py中build_layer_index_mapping的核心逻辑第i个*位置取 GPT 第i层的 attention第i个-/E位置取 GPT 第i层的 MLP见 转换工具源码。模式还可以描述执行布局|标记流水线段边界/引入重复的 MTP 模式。例如*-*-|*-*-把 4 个 GPT 等价 block 放到两个流水线段上。分隔符不计入层数。从 parse_hybrid_pattern 的实现可见统一模式的完整语法为main_pattern/mtp_pattern/mtp_pattern/...且所有 MTP 子模式必须完全一致例如M*M*/MM/MM表示主解码器M*M*加 2 层 MTP 深度。1.2 HybridModel 提供的能力不同的层家族可以组合进同一个模型无需强制每个 decoder block 结构一致attention 与稠密/专家 MLP 层可以独立放置与配置Mamba、GDN、标准或 DeepSeek attention、稠密 MLP 与 MoE 层共享同一个模式驱动的模型接口HybridStack通过hybrid_stack_spec统一装配见 hybrid_layer_specs.py包含 Mamba 的模式可以用亚二次subquadratic序列混合与固定大小的循环推理状态替代部分二次复杂度 attention 层流水线与虚拟流水线切分可以直接用模型模式表达而不再依赖独立的层布局参数同一个模型抽象既能描述纯 Transformer重复*-、纯 Mamba 模型也能描述异构架构。需要强调的是这些能力并不自动带来吞吐或质量的提升。保留架构的*-或*E迁移应当验证数值等价性而引入了新层家族的模式应视为新架构需要独立做基准测试这一点在 HybridModel 类文档 中同样没有承诺任何性能收益。2. 如何转换检查点将GPTModel权重带入HybridModel运行有两条路径两者都停留在 Megatron 的分布式检查点distributed-checkpoint格式内不需要以 Hugging Face 作为中间格式并且后续加载时都可以跨不同的 TP/PP/EP/FSDP 布局重新分片方案 A——加载时翻译无需单独步骤直接用 hybrid 运行去加载 GPT 检查点。hybrid 模型会在加载期间把自己的检查点 state dict 重定向到 GPT 检查点的 key 上磁盘上不会写出第二份拷贝。同时支持torch_dist检查点与 Megatron-FSDP 的fsdp_dtensor检查点含优化器状态。该路径也支持包含 GPT 没有对应物的层家族的模式——例如 MambaM位置会保留自己的全新初始化。方案 B——离线转换为新检查点使用 tools/checkpoint/gpt_hybrid_conversion.py 写出一份独立的HybridModel检查点其 key 已经符合 hybrid 布局。适用于需要持久化 hybrid 检查点、需要一份保留架构的*-或*E拷贝、或希望在训练前先检查目标结果的场景。该路径只支持*-与*E两种布局。2.1 方案 A加载时翻译加载时翻译由 megatron/core/dist_checkpointing/gpt_checkpoint_interop.py 处理当非 hybridGPT检查点被加载进HybridModel运行时自动触发对于torch_dist运行的模型与优化器分片 state dict 会被改写成 GPT 检查点的同构层homogeneous-layer格式——把decoder.layers.g.mlp...中的层索引从存储 key 中剥离、将匹配的 GPT 层索引变为前置的分片轴恰好镜像TransformerBlock.sharded_state_dict在non_homogeneous_layersFalse时保存的格式见 retarget_sharded_state_dict_to_gpt_checkpoint 与 模块头注释对于fsdp_dtensor在 Torch DCP 规划之前把显式的参数名映射改写为 GPT 的 key见 retarget_fsdp_state_dict_to_gpt_checkpoint。检查点被直接读取权重与优化器状态会重新分片到当前 TP/PP/EP/ETP/FSDP 布局不运行任何转换工具磁盘上的 GPT 检查点永不被修改。反向不匹配会直接报错把 hybrid 运行保存的检查点加载进非 hybrid 运行会抛出RuntimeError并引导你使用 hybrid 训练入口。选择检查点语义当 hybrid 运行加载 GPT 检查点时必须设置--hybrid-layer-pattern以便把检查点层与 hybrid 层位置配对。将--load或--pretrained-checkpoint指向 GPT 检查点根目录见第 3 节。--finetune可选并保留其常规的检查点加载含义GPT-to-Hybrid 翻译不会替用户选择它不带--finetune时直接--load会按常规检查点与并行布局兼容规则恢复迭代、优化器、调度器、RNG 与 rerun 状态带--finetune时迭代、调度器、RNG 与 rerun 状态全部从头开始。模型权重仍会加载翻译后的优化器状态也会加载除非设置--no-load-optim现有的--pretrained-checkpoint回退路径当--load目录中没有检查点时使用微调语义。需要完整恢复语义时请直接使用--load。默认情况下GPT 运行的优化器状态也会被翻译并加载——Adam 动量与 fp32 master 参数会随 attention 与 MLP 层一起带入从而支持保留架构的继续训练。传入--no-load-optim可跳过这一步让每一层的优化器状态从头开始。对于torch_dist加载优化器状态要求 GPT 检查点使用 model-space 的分布式优化器格式fully_reshardable或fully_sharded_model_space即用--dist-ckpt-optim-fully-reshardable保存。bucket-space 格式按扁平 buffer 布局组织优化器状态 key额外的 hybrid 层会打乱该布局因此运行会报错并提示你重新保存检查点或传入--no-load-optim。对于 Megatron FSDP请使用--ckpt-format fsdp_dtensor保存与加载。加载器会重定向显式的 DTensor 模型 key 以及分布式优化器使用的模型参数名因此模型权重与优化器状态可以在自动 GPT-to-Hybrid 加载期间跨不同的 FSDP、TP、EP、ETP 布局重新分片。该路径面向 Megatron FSDPTorch FSDP2 的torch_dcp格式不受自动翻译支持。没有 GPT 对应物的层例如 MambaM位置在检查点中没有优化器状态其动量从头开始运行会打印一条警告说明受影响层数。支持的模式与 key 映射主模式/MTP 后缀之前的部分忽略|流水线分隔符可包含符号权重来源*配对 GPT 层的self_attention子模块-或E配对 GPT 层的mlp子模块MoE 张量也位于mlp.*之下M无 GPT 来源Mamba 层保留全新初始化decoder.final_norm从 GPT 的decoder.final_layernorm加载embedding 与输出权重原样复制。由于每个 GPT 层恰好提供一个 attention 与一个 MLP 子模块加载器会拒绝以下模式包含 MTP 层/...后缀——它们没有 GPT 源权重使用无法翻译的层类型如 GDNG或 DeepSeek Sparse AttentionD——其权重布局与 GPT attention 不同在同一模式中混用稠密-与 MoEEMLP 位置*与 MLP 位置数量不相等或为零。检查点的num_layers必须等于模式中*位置的数量不匹配会被拒绝。警告优化器翻译覆盖 Adam 动量与 fp32 master 参数。仅加载权重请传--no-load-optim需要迭代、调度器、RNG 与 rerun 状态重新开始而非恢复时使用--finetune。2.2 方案 B用gpt_hybrid_conversion.py离线转换选择保留架构的模式对于含N个 GPT 层的源检查点稠密 GPT 模型使用*-重复N次每层 MLP 均为 MoE 的 GPT 模型使用*E重复N次。转换器按出现次序而非仅按数值层号映射参数源参数目标参数GPT 第i层的 attention第i个*层GPT 第i层的 MLP 或 MoE第i个-或E层embedding 与输出权重原样复制模型角色不变decoder.final_layernorm重命名为decoder.final_norm这一映射逻辑在 convert_gpt_to_hybrid 中实现decoder.layers.下的 key 先按is_attention_param/is_mlp_param归类再通过replace_layer_num把层号改写为映射后的 hybrid 层号final_layernorm统一替换为final_norm。MoE 张量因两侧都位于mlp.{router,experts,shared_experts}.*之下而可原样往返无需折叠专家或重初始化 router见转换工具头注释。检查前置条件源检查点必须使用以下分布式检查点格式之一torch_distfsdp_dtensor优先把包含latest_checkpointed_iteration.txt的顶层检查点根目录作为--load-dir。如果--load-dir直接指向含metadata.json的目录转换器会写出一份没有 tracker 文件的扁平目标而标准训练入口期望检查点根目录与 tracker。在仓库根目录、Megatron 环境下运行转换器即可。普通python进程就足够不需要torchrun也不需要 GPU。工具会在 CPU 上聚合完整的逻辑张量因此宿主机需要有足够内存容纳未分片的源与目标模型 state dict。目标目录必须与源目录不同。警告这是权重转换不是可恢复的完整训练状态转换。分片的优化器、RNG、rerun 以及 Transformer Engine_extra_state张量不会被转换。部分非张量条目可能残留在common.pt中但它们不构成转换后的优化器或 RNG 状态。转换后的模型请使用全新的优化器与 RNG 状态启动。运行转换以下示例将一个四层稠密 GPT 模型转换为等价的八层 HybridModel模式*-*-*-*-uv run python tools/checkpoint/gpt_hybrid_conversion.py \ --direction gpt-to-hybrid \ --load-dir /path/to/gpt-checkpoints \ --save-dir /path/to/hybrid-checkpoints \ --hybrid-layer-pattern *-*-*-*- \ --reset-iterations务必给模式加引号因为*和|对 shell 有特殊含义。--input-format auto与--output-format auto是默认值工具自动检测源后端并写出相同后端也可用--input-format/--output-format显式指定torch_dist或fsdp_dtensor。--reset-iterations会重置检查点迭代数、已消费样本计数以及缓存的train_iters与train_samples当新运行需要保留这些调度元数据时请省略它。模式中*位置数与-或E位置数都必须等于源 GPT 层数。模式校验器会拒绝 GDN、DSA 以及稠密/MoE 混合布局对应 validate_pattern_gpt_compatible 的实现。当检查点中缓存了训练参数时工具还会拒绝交错 MoE、实验性或线性 attention、异构 block 规格、Multi-Latent Attention 与 MTP 检查点对应 _GPT_COMPAT_REJECT_FIELDS 中的字段清单moe_layer_freq、experimental_attention_variant、linear_attention_freq、heterogeneous_block_specs、heterogeneous_layers_config_path、multi_latent_attention、mtp_num_layers等。当common.pt中没有缓存的args、或较旧检查点缺少某个字段时该源特性校验不完整需要手动核对这些特性。转换只识别标准的 attention 与 MLP/MoE state-dict key其他层局部张量会被省略。文档化的hybrid_stack_spec使用 Transformer Engine 的 fused layernorm/linear 布局本地或其他非 TE 的源布局需要兼容的自定义 Hybrid stack 与 key 转换。务必执行下文所述的严格加载检查。不要在转换期间追加 MTP/...后缀。转换器只映射第一个/之前的主模式因此不会创建 MTP 参数。当源路径是含latest_checkpointed_iteration.txt的检查点根目录时输出会包含一个迭代目录与匹配的 tracker 文件。保存的全形状张量可以由后续 Megatron 加载按不同的 tensor、pipeline、expert 或 FSDP 布局重新分片。此外该工具还支持反向hybrid-to-gpt方向从模块头注释可见完整用法此时 SSM 层被丢弃并打印警告旧版mp_rank_XX/model_optim_rng.pt布局不受支持需要先转成torch_dist。3. 如何训练模型3.1 更新训练命令从训练 GPT 模型的命令出发做以下改动将pretrain_gpt.py替换为pretrain_hybrid.py删除--num-layers改用与转换时相同的、有序的主层符号序列。可以添加或移动流水线|分隔符。命令行解析器会从模式中推导num_layers见 arguments.py当--hybrid-layer-pattern存在时num_layers自动取get_hybrid_total_layer_count的结果并忽略--num-layers用--spec megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec选择 HybridModel 栈规格。从 hybrid_builder 的实现可以看到该--spec通过import_module动态导入未提供时会直接报错把检查点输入指向预训练权重并把新训练的检查点写到独立目录方案 A加载时翻译将--load直接指向GPT检查点以获得恢复语义或用--pretrained-checkpoint获得微调语义。优化器状态默认加载只有想要全新优化器时才加--no-load-optim。无需离线转换翻译过程也不要求--finetune。方案 B离线转换将--pretrained-checkpoint指向转换后的hybrid检查点并将--ckpt-format设置为转换器的torch_dist或fsdp_dtensor输出格式。一个最小的方案 A 迁移——直接加载 GPT 检查点及其优化器状态进行保留架构的继续训练——如下- torchrun --nproc_per_node8 pretrain_gpt.py \ - --num-layers 4 \ - --load /path/to/gpt-checkpoints \ - --save /path/to/gpt-checkpoints torchrun --nproc_per_node8 pretrain_hybrid.py \ --hybrid-layer-pattern *-*-*-*- \ --spec megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec \ --load /path/to/gpt-checkpoints \ # 首次启动之后切换到 save 目录 --save /path/to/new-training-checkpoints方案 B 则改为把--pretrained-checkpoint指向转换后的 hybrid 检查点。保持既有的架构、优化器、精度、数据与基础 TP/DP/EP/CP 参数不变除非本指南指明需要改动。模式驱动的流水线布局与 GPT 专属的数据集特性需要另行评估。警告pretrain_hybrid.py在设置--fim-data时不会选择GPTFIMDataset。使用 fill-in-the-middle 数据的 GPT 训练流程在迁移前需要自定义数据集路径或等价的 Hybrid 入口支持。当--load目录为空时--pretrained-checkpoint以微调语义加载预训练权重迭代从零开始、RNG 状态不恢复。方案 A 下优化器状态仍会热启动除非设置--no-load-optim方案 B 总是以全新优化器开始。任务向--load写入检查点后后续启动会正常恢复新的 HybridModel 训练状态。3.2 从零开始训练若要从头初始化每一层使用相同的pretrain_hybrid.py、--hybrid-layer-pattern与--spec参数但省略--pretrained-checkpoint。把--load与--save指向新运行目录以便后续启动恢复。与检查点转换不同从零训练可以使用所选 HybridModel 栈规格支持的所有兼容层家族并且可以包含 MTP 后缀例如M*M*/MM/MM。模式约束仍然适用——例如标准 attention*与 DSAD不能出现在同一个模型中。3.3 注意层索引扩展任何按 decoder 层索引的下标、映射或回调都必须改用 HybridModel 索引。以*E*E为例源第 0 层的 attention 是 Hybrid 第 0 层、其 MoE 是 Hybrid 第 1 层源第 1 层的 attention 是 Hybrid 第 2 层、其 MoE 是 Hybrid 第 3 层。只关注 attention 的列表需要为中间的 MLP 或 MoE 位置补上非活动条目。这一每类层各自编号的对应关系正是get_layer_maps_from_layer_type_listhybrid_layer_allocation.py所建立的每个全局层索引都被映射到该符号第几个出现的局部索引。3.4 配置流水线并行对于流水线并行直接添加|分隔符而无需改变有序层符号。例如转换后的模式*-*-*-*-可以按两个流水线段训练为*-*-|*-*-。管道分隔的段数必须能被--pipeline-model-parallel-size整除该约束在 select_pipeline_segment 中强制校验。模式取代了传统的流水线布局控制参数。请删除--num-layers-per-virtual-pipeline-stage--num-virtual-stages-per-pipeline-rank--pipeline-model-parallel-layout--account-for-embedding-in-pipeline-split--account-for-loss-in-pipeline-split当模式包含|时还需删除--decoder-first-pipeline-num-layers--decoder-last-pipeline-num-layers虚拟流水线切分改用额外的管道分隔段来表达。注意声明式的HybridModelBuilder目前拒绝虚拟流水线并行pretrain_hybrid.py的 CLI builder 支持 pipe 定义的虚拟 stage但自定义 builder 用户必须避免 VPP或使用显式支持它的路径。从源码看无|模式 pp_size 1的情况仍保留旧的自动切层逻辑按num_layers % pp_size均匀切分或通过--num-layers-in-first/last-pipeline-stage不均匀切分但会打印弃用警告建议显式加|分隔符select_pipeline_segment。3.5 更新自定义 provider 与转换映射自定义 provider 与转换映射还需适配以下 API 与 state-dict 差异构建或注册HybridModel而非GPTModel提供hybrid_stack_spec而非 GPT transformer-layer spec将程序化num_layers设为主模式中的层符号数与 CLI 路径不同自定义 provider 可能不会自动推导它把 attention 与 MLP/MoE 参数映射到各自独立的 Hybrid 层索引在 HybridModel 映射中使用decoder.final_norm而非decoder.final_layernorm把 attention-window schedule 等逐层设置扩展到完整的 HybridModel 模式。3.6 在扩展规模前验证在启动长任务之前严格加载转换后的检查点确认没有模型 key 或张量形状缺失/意外即开启 strict-load对*-或*E迁移在固定 batch 上把 logits 与源 GPT 模型在预期精度容差内做对比运行若干训练迭代检查 loss、梯度范数与逐层参数计数保存并重新加载一个新检查点确认新优化器与 RNG 状态能正确恢复。4. 仓库中的验证与测试资产以上迁移路径并非孤立文档仓库提供了多层次的测试与可运行样例用于印证转换工具单元测试tests/unit_tests/tools/checkpoint/test_gpt_hybrid_conversion.py 验证模式解析含 MTP/pipe 分隔符剥离、GPT↔Hybrid 层索引映射、final_layernorm↔final_norm重命名、SSM 参数初始化形状与 dtype以及 GPT→Hybrid→GPT 往返后 attention/MLP 权重保持test_gpt_hybrid_conversion_parallelism.py 覆盖并行布局相关行为加载时互操作测试tests/unit_tests/dist_checkpointing/models/test_gpt_hybrid_interop.py 直接验证gpt_checkpoint_interop.py的加载时翻译逻辑功能测试样例tests/functional_tests/test_cases/hybrid/下包含大量 HybridModel 训练/推理配置例如hybrid_mr_mcore_te_tp1_pp4_cp1_dgx_a100_1N8G、hybrid_mr_mcore_te_tp1_pp2_vpp2_*演示 pipe 定义 VPP、hybrid_static_inference_tp1_pp1_2B_logitsmatch演示 logits 一致性校验等可作为迁移后训练配置的参考模板Mamba 示例脚本examples/mamba/train.sh 展示了--hybrid-layer-pattern与--spec在实际训练命令中的用法。5. 总结从GPTModel到HybridModel的迁移在 Megatron Core 中是零中间格式的要么在加载时由 gpt_checkpoint_interop.py 自动把 hybrid 的 sharded state dict 重定向到 GPT 检查点的同构层格式支持torch_dist与fsdp_dtensor含优化器状态要么用 gpt_hybrid_conversion.py 离线生成独立的 hybrid 检查点。无论走哪条路模式中的*与-/E都按出现次序与 GPT 层一一配对M位置保持全新初始化G/D与 MTP 后缀因缺乏 GPT 源权重而被拒绝。训练侧只需把入口换成pretrain_hybrid.py、用--hybrid-layer-pattern取代--num-layers与各类流水线布局参数并在规模化之前完成严格加载与数值等价性验证——这既是迁移的安全边界也是充分发挥 HybridModel 架构灵活性的起点。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
Shuffled Word Order:词序打乱预训练 RoBERTa 在 GLUE 与 PAWS 上微调的完整实践指南)
