Megatron-LM 分布式训练实战2 卡跑通 GPT 训练的最小路径【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM单卡 80GB 显存装不下 7B 模型的激活第一反应是把一层 Transformer 拆到多张卡上——这正是张量并行的用武之地。Megatron-LM核心库为 Megatron Core是 NVIDIA 开源的 Transformer 分布式训练框架支持 TP/PP/CP/EP/FSDP 并行组合。下面走一遍仓库里的 2 卡路径环境 → loss 下降 → 跨并行配置迁移检查点。它凭什么高效拆解三种主力并行机制Megatron Core 不是给你一个 DDP 黑盒而是把模型装不下单卡拆成三个正交的切分维度。官方基准H100 集群上 MFU 最高 47%覆盖 2B 到 462B 参数模型张量并行把一层横向切成多份注意力 QKV 投影与 MLP 权重矩阵按行/列切分每张卡只持有一份分片层输出靠 allreduce 拼回一层计算只需两次卡间集合通信。单卡参数与激活占用随 TP 度数线性下降直接解决单层放不下的问题再配合 sequence parallelLayerNorm/Dropout 的激活也按序列维切分。流水线并行按深度切模型、用虚拟流水线摊平气泡把 N 层切成若干 stage每张卡负责连续几层微批数据流式通过前向反向交替使流水线保持填充。stage 间只传层输入输出P2P通信量远低于 allreduce。纯 PP 存在气泡虚拟流水线 stage交错调度把气泡摊到每个微批上适合 50 层的深模型与 TP 一个管宽、一个管深。上下文并行长序列不怕 8K token把注意力沿序列维的 K/V 分布到多卡Ring Attention 思路单卡显存不再与序列长度成正比。Dynamic CP 还能按变长序列自适应调整并行度数官方报告变长训练最高 1.48 倍加速适合长上下文场景。把环境跑起来pip 与 NGC 容器两条路最轻pip 装pip install uv uv pip install megatron-core需要 WB、tokenizer 等训练依赖就装megatron-core[training]。要改源码或调试git clone https://gitcode.com/GitHub_Trending/me/Megatron-LM cd Megatron-LM uv pip install -e .最稳NGC 容器docker run --ipchost --shm-size512m --gpus 2 -it nvcr.io/nvidia/pytorch:24.02-py3两个参数都有讲究--ipchost让 NCCL 走 host 共享内存做节点内 P2P 传输而不是退回 /tmp直接影响多卡带宽--shm-size512m把容器默认 64MB 的 /dev/shm 扩上去否则多进程 DataLoader 的 worker 在搬运 batch 时会挂。容器内再执行uv pip install --no-build-isolation megatron-core[training,dev]torch 已预装、TE 的 CUDA 内核也预编译好了省掉 20 分钟以上的编译。安装方式全集见 docs/get-started/install.md。验证一行带过import megatron.core; print(megatron.core.__version__)能打印版本号就 OK。用 2 卡跑通第一个 step从 TransformerConfig 到 loss 下降完整示例在 examples/run_simple_mcore_train_loop.py直接跑torchrun --nproc-per-node 2 examples/run_simple_mcore_train_loop.py。第一步初始化分布式 建配置。这段做两件事拉起 NCCL 进程组再告诉 Megatron 卡怎么分组TP2PP1。# 依赖: torch, megatron.core torch.distributed.init_process_group(backendnccl) parallel_state.initialize_model_parallel(2, 1) config TransformerConfig(num_layers2, hidden_size12, num_attention_heads4)无输出进程组与配置就位。第二步组装模型。GPTModel 搭一个 2 层的玩具模型再包一层 DDP 获得梯度同步能力。# 依赖: megatron.core.models.gpt, megatron.core.distributed model GPTModel(configconfig, transformer_layer_specget_gpt_layer_local_spec(), vocab_size100, max_sequence_length64) model DistributedDataParallel(configconfig, ddp_configddp_cfg, modulemodel)此时模型已分布在两张卡上尚无输出等待开跑。第三步造数据。用不到真实语料MockGPTDataset 生成 1000 条随机 token 序列足够验证链路。# 依赖: torch.utils.data, megatron.core.datasets cfg GPTDatasetConfig(sequence_length64, tokenizertok) ds BlendedMegatronDatasetBuilder(MockGPTDataset, [1000, None, None], lambda: True, cfg).build() data_iter iter(DataLoader(ds[0], batch_size8, shuffleTrue))数据就绪依旧无输出。第四步进入 forward/backward 循环。get_forward_backward_func()按当前并行组合自动选调度这里等价于每轮一次前向反向fwd_step是返回 (输出, loss 函数) 的前向回调。# 依赖: torch.optim, megatron.core.pipeline_parallel fwd_bwd get_forward_backward_func() for _ in range(5): optim.zero_grad() fwd_bwd(forward_step_funcfwd_step, data_iteratordata_iter, modelmodel, micro_batch_size8) finalize_model_grads([model]); optim.step()跑完控制台打印 5 行Iteration i: Losses reduced: {lm loss: ...}。数据是随机的loss 不会收敛但只要看到有限数值链路就通了。换并行配置不用重新训练分布式 checkpoint 保存与加载生产里的真实痛点TP2 训好的模型想换 TP4 续训普通 state_dict 的分片布局没法直接用。Megatron 的解法是分布式 checkpoint——sharded_state_dict()记录每个参数在各 rank 上的分片映射save/load 都是集合通信操作加载时框架自动按新并行组 reshard# 依赖: megatron.core.dist_checkpointing sd model.sharded_state_dict(prefix) dist_checkpointing.save(sharded_state_dictsd, checkpoint_dird) model.load_state_dict(dist_checkpointing.load(sharded_state_dictsd, checkpoint_dird))示例脚本的最后两步就是先存再读回打印Successfully loaded the model。正式训练里 pretrain_gpt.py 用--ckpt-format dist_ckpt打开同一套机制接口细节见 docs/api-guide/core/dist_checkpointing.md。分布式优化器同理把优化器状态沿 DP 维分片单卡优化器显存省一个 DP 倍数并行策略调优清单TP 放节点内避免跨机 allreduceTP 配 sequence parallel 省激活显存PP 用虚拟流水线 stage 摊平气泡长序列优先 CP 扩规模Hopper/Blackwell 上开 FP8 提精度吞吐扩规模前可对照图中弱扩展效率曲线预估收益【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
Megatron-LM 分布式训练实战：2 卡跑通 GPT 训练的最小路径
Megatron-LM 分布式训练实战2 卡跑通 GPT 训练的最小路径【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM单卡 80GB 显存装不下 7B 模型的激活第一反应是把一层 Transformer 拆到多张卡上——这正是张量并行的用武之地。Megatron-LM核心库为 Megatron Core是 NVIDIA 开源的 Transformer 分布式训练框架支持 TP/PP/CP/EP/FSDP 并行组合。下面走一遍仓库里的 2 卡路径环境 → loss 下降 → 跨并行配置迁移检查点。它凭什么高效拆解三种主力并行机制Megatron Core 不是给你一个 DDP 黑盒而是把模型装不下单卡拆成三个正交的切分维度。官方基准H100 集群上 MFU 最高 47%覆盖 2B 到 462B 参数模型张量并行把一层横向切成多份注意力 QKV 投影与 MLP 权重矩阵按行/列切分每张卡只持有一份分片层输出靠 allreduce 拼回一层计算只需两次卡间集合通信。单卡参数与激活占用随 TP 度数线性下降直接解决单层放不下的问题再配合 sequence parallelLayerNorm/Dropout 的激活也按序列维切分。流水线并行按深度切模型、用虚拟流水线摊平气泡把 N 层切成若干 stage每张卡负责连续几层微批数据流式通过前向反向交替使流水线保持填充。stage 间只传层输入输出P2P通信量远低于 allreduce。纯 PP 存在气泡虚拟流水线 stage交错调度把气泡摊到每个微批上适合 50 层的深模型与 TP 一个管宽、一个管深。上下文并行长序列不怕 8K token把注意力沿序列维的 K/V 分布到多卡Ring Attention 思路单卡显存不再与序列长度成正比。Dynamic CP 还能按变长序列自适应调整并行度数官方报告变长训练最高 1.48 倍加速适合长上下文场景。把环境跑起来pip 与 NGC 容器两条路最轻pip 装pip install uv uv pip install megatron-core需要 WB、tokenizer 等训练依赖就装megatron-core[training]。要改源码或调试git clone https://gitcode.com/GitHub_Trending/me/Megatron-LM cd Megatron-LM uv pip install -e .最稳NGC 容器docker run --ipchost --shm-size512m --gpus 2 -it nvcr.io/nvidia/pytorch:24.02-py3两个参数都有讲究--ipchost让 NCCL 走 host 共享内存做节点内 P2P 传输而不是退回 /tmp直接影响多卡带宽--shm-size512m把容器默认 64MB 的 /dev/shm 扩上去否则多进程 DataLoader 的 worker 在搬运 batch 时会挂。容器内再执行uv pip install --no-build-isolation megatron-core[training,dev]torch 已预装、TE 的 CUDA 内核也预编译好了省掉 20 分钟以上的编译。安装方式全集见 docs/get-started/install.md。验证一行带过import megatron.core; print(megatron.core.__version__)能打印版本号就 OK。用 2 卡跑通第一个 step从 TransformerConfig 到 loss 下降完整示例在 examples/run_simple_mcore_train_loop.py直接跑torchrun --nproc-per-node 2 examples/run_simple_mcore_train_loop.py。第一步初始化分布式 建配置。这段做两件事拉起 NCCL 进程组再告诉 Megatron 卡怎么分组TP2PP1。# 依赖: torch, megatron.core torch.distributed.init_process_group(backendnccl) parallel_state.initialize_model_parallel(2, 1) config TransformerConfig(num_layers2, hidden_size12, num_attention_heads4)无输出进程组与配置就位。第二步组装模型。GPTModel 搭一个 2 层的玩具模型再包一层 DDP 获得梯度同步能力。# 依赖: megatron.core.models.gpt, megatron.core.distributed model GPTModel(configconfig, transformer_layer_specget_gpt_layer_local_spec(), vocab_size100, max_sequence_length64) model DistributedDataParallel(configconfig, ddp_configddp_cfg, modulemodel)此时模型已分布在两张卡上尚无输出等待开跑。第三步造数据。用不到真实语料MockGPTDataset 生成 1000 条随机 token 序列足够验证链路。# 依赖: torch.utils.data, megatron.core.datasets cfg GPTDatasetConfig(sequence_length64, tokenizertok) ds BlendedMegatronDatasetBuilder(MockGPTDataset, [1000, None, None], lambda: True, cfg).build() data_iter iter(DataLoader(ds[0], batch_size8, shuffleTrue))数据就绪依旧无输出。第四步进入 forward/backward 循环。get_forward_backward_func()按当前并行组合自动选调度这里等价于每轮一次前向反向fwd_step是返回 (输出, loss 函数) 的前向回调。# 依赖: torch.optim, megatron.core.pipeline_parallel fwd_bwd get_forward_backward_func() for _ in range(5): optim.zero_grad() fwd_bwd(forward_step_funcfwd_step, data_iteratordata_iter, modelmodel, micro_batch_size8) finalize_model_grads([model]); optim.step()跑完控制台打印 5 行Iteration i: Losses reduced: {lm loss: ...}。数据是随机的loss 不会收敛但只要看到有限数值链路就通了。换并行配置不用重新训练分布式 checkpoint 保存与加载生产里的真实痛点TP2 训好的模型想换 TP4 续训普通 state_dict 的分片布局没法直接用。Megatron 的解法是分布式 checkpoint——sharded_state_dict()记录每个参数在各 rank 上的分片映射save/load 都是集合通信操作加载时框架自动按新并行组 reshard# 依赖: megatron.core.dist_checkpointing sd model.sharded_state_dict(prefix) dist_checkpointing.save(sharded_state_dictsd, checkpoint_dird) model.load_state_dict(dist_checkpointing.load(sharded_state_dictsd, checkpoint_dird))示例脚本的最后两步就是先存再读回打印Successfully loaded the model。正式训练里 pretrain_gpt.py 用--ckpt-format dist_ckpt打开同一套机制接口细节见 docs/api-guide/core/dist_checkpointing.md。分布式优化器同理把优化器状态沿 DP 维分片单卡优化器显存省一个 DP 倍数并行策略调优清单TP 放节点内避免跨机 allreduceTP 配 sequence parallel 省激活显存PP 用虚拟流水线 stage 摊平气泡长序列优先 CP 扩规模Hopper/Blackwell 上开 FP8 提精度吞吐扩规模前可对照图中弱扩展效率曲线预估收益【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
