Megatron-LM

活跃
GitHub Python NOASSERTION

简介

NVIDIA 开源的 GPU 优化型 Transformer 分布式训练库,提供张量/流水线/序列并行策略,是训练大模型的核心底层库。

核心特性

  • 张量并行(TP) — 将单个 Transformer 层分片到多个 GPU,大幅降低单卡显存需求
  • 流水线并行(PP) — 将模型层级分配到不同 GPU,支持超大规模模型
  • 上下文并行(CP) — 有效处理超长序列,支持百万 token 训练
  • 混合精度训练 — FP16、BF16、FP8、FP4 全链路低精度训练支持
  • 混合专家 MoE — 原生支持 DeepSeek-V3、Mixtral 等 MoE 架构
  • Megatron Bridge — 与 HuggingFace 模型格式双向转换

适用场景

💡 大语言模型预训练:在数千 GPU 上训练千亿参数模型
💡 MoE 模型训练:高效训练 DeepSeek-V3、Mixtral 等混合专家模型
💡 超长上下文训练:百万 token 级别的长序列预训练
💡 自定义训练框架:基于 Megatron Core 模块化组件构建定制化训练系统

标签

transformer distributed-training gpu llm parallel python

分类

🤖 Agent 框架

快速开始

# 安装 Megatron Core
uv pip install megatron-core

# 或从源码安装
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM && uv pip install -e .

# 运行快速开始训练脚本
cd examples
python pretrain_gpt.py ...

# 查看文档
# https://docs.nvidia.com/megatron-core/developer-guide/latest/get-started/quickstart.html
访问 GitHub 访问官网 查看文档

相关项目