用途:使用 NVIDIA Megatron-Core 训练 2B 至 462B 参数的大语言模型,通过多种并行策略在 H100 上实现最高约 47% 的模型 FLOP 利用率(MFU)。
主要能力:支持张量并行(TP)、流水线并行(PP)、数据并行(DP)、上下文并行(CP)与专家并行(EP);支持 LLaMA 风格模型与 Mixtral 等稀疏 MoE 模型;提供 FP8 混合精度、Flash Attention、序列并行、梯度检查点与 CPU 优化器卸载等优化手段。
使用方式:推荐用 Docker 镜像 nvcr.io/nvidia/pytorch:25.04-py3,或 pip install megatron-core。简单分布式训练可运行 torchrun --nproc_per_node=2 examples/run_simple_mcore_train_loop.py,LLaMA-3 8B 可运行 ./examples/llama/train_llama3_8b_fp8.sh。典型流程为:按模型规模选择并行配置(如 7B 用 8 卡 DP=8;70B 用 64 卡 TP=4/PP=4/DP=4;405B 用 128 卡 TP=8/PP=8/CP=2),配置超参数,用 torchrun 或 SLURM 启动,并监控 MFU、吞吐、显存与 loss。MoE 训练需设置 --expert-model-parallel-size、--num-experts、--moe-router-topk 等参数。
适用场景:训练超过 10B 参数、追求高 MFU、使用 NVIDIA A100/H100/B200 的生产级大规模训练,以及需要细粒度并行控制的场景。
限制:需 NVIDIA Ampere 及以上 GPU,FP8 需 Hopper/Ada/Blackwell;多节点需 InfiniBand 或 400Gb+ 以太网;70B 模型每卡约需 80GB 显存,405B 需 TP=8、PP=8;检查点存储建议 1TB+ NVMe。若模型小于 70B 或追求更简单 API,可考虑 PyTorch FSDP、DeepSpeed、HuggingFace Accelerate 或 LitGPT。常见问题包括 MFU 偏低、显存不足、训练变慢与 loss 发散,可通过调整 micro-batch、并行度、重计算、warmup 与梯度裁剪缓解。