# ai阅读megatron # Megatron-LM 详尽文档 ## 目录 1. 项目概述 2. 架构设计 3. 核心组件 4. 训练脚本 5. 并行策略 6. 数据处理 7. 优化器与调度器 8. 模型实现 9. 示例与工具 10. 性能优化 --- ## 项目概述 ### 什么是 Megatron-LM？ **Megatron-LM** 是 NVIDIA 开发的用于大规模训练 Transformer 模型的 GPU 优化库。它包含两个主要部分： - **Megatron Core**: 可组合的生产级库，提供 GPU 优化的构建模块 - **Megatron-LM**: 参考实现，包含完整的训练脚本和示例 ### 版本信息 - **当前版本**: 0.14.0 (2025) - **许可证**: Apache 2.0 - **Python 版本**: 3.12 推荐 - **PyTorch**: 最新稳定版本 ### 主要特性 ✅ **大规模训练**: 支持从 2B 到 462B+ 参数的模型 ✅ **高性能**: 在 H100 集群上达到 47% 的 Model FLOP Utilization (MFU) ✅ **多种并行策略**: TP, PP, DP, CP, EP 及其组合 ✅ **混合精度**: FP16, BF16, FP8 支持 ✅ **模型支持**: GPT, LLaMA, Mixtral, Mamba, BERT, T5, 多模态等 ✅ **容错训练**: 集成 NVRx 进行故障检测和恢复 --- ## 架构设计 ### 项目结构 Megatron-LM/ ├── megatron/ # 核心代码库 │ ├── core/ # Megatron Core (生产级库) │ │ ├── models/ # 模型实现 (GPT, BERT, T5, Mamba, 多模态等) │ │ ├── transformer/ # Transformer 构建块 │ │ ├── tensor_parallel/ # 张量并行 │ │ ├── pipeline_parallel/ # 流水线并行 │ │ ├── distributed/ # 分布式训练 (FSDP, DDP) │ │ ├── optimizer/ # 优化器 │ │ ├── datasets/ # 数据集加载器 │ │ ├── inference/ # 推理引擎 │ │ └── export/ # 模型导出 (TensorRT-LLM) │ ├── training/ # 训练基础设施 │ ├── inference/ # 推理服务器 │ ├── legacy/ # 遗留组件 │ └── post_training/ # 后训练 (RLHF 等) ├── examples/ # 训练示例 ├── tools/ # 实用工具 ├── tests/ # 测试套件 └── docs/ # 文档 ### 设计理念 #### Megatron Core: 可组合库 **目标用户**: - 框架开发者构建自定义训练框架 - 研究团队需要自定义训练循环 - ML 工程师构建容错训练流水线 **提供内容**: - 可组合的 Transformer 构建块 (attention, MLP 等) - 高级并行策略 (TP, PP, DP, EP, CP) - 流水线调度和分布式优化器 - GPU 优化内核和内存管理 - 高性能数据加载器 #### Megatron-LM: 参考实现 **目标用户**: - 训练最先进的基础模型 - 研究团队探索新架构 - 学习分布式训练概念 - 快速实验 **提供内容**: - 预配置的训练脚本 (GPT, LLaMA, DeepSeek, Qwen 等) - 端到端示例 (从数据准备到评估) - 研究工具和实用程序 --- ## 核心组件 ### 1. Parallel State ( megatron/core/parallel_state.py) **功能**: 管理所有并行组和全局状态 **关键组件**: # 张量并行组 _TENSOR_MODEL_PARALLEL_GROUP _TENSOR_MODEL_PARALLEL_WORLD_SIZE _TENSOR_MODEL_PARALLEL_RANK # 流水线并行组 _PIPELINE_MODEL_PARALLEL_GROUP _PIPELINE_MODEL_PARALLEL_WORLD_SIZE _PIPELINE_MODEL_PARALLEL_RANK # 数据并行组 _DATA_PARALLEL_GROUP _DATA_PARALLEL_WORLD_SIZE _DATA_PARALLEL_RANK # 专家并行组 (MoE) _EXPERT_MODEL_PARALLEL_GROUP _EXPERT_TENSOR_PARALLEL_GROUP # 上下文并行组 _CONTEXT_PARALLEL_GROUP _HIERARCHICAL_CONTEXT_PARALLEL_GROUPS **主要函数**: - `initialize_model_parallel()`: 初始化所有并行组 - `get_tensor_model_parallel_group()`: 获取 TP 组 - `get_pipeline_model_parallel_group()`: 获取 PP 组 - `get_data_parallel_group()`: 获取 DP 组 - RankGenerator: 生成不同并行模式的 rank 组 ### 2. Transformer Config ( megatron/core/transformer/transformer_config.py) **功能**: 统一的 Transformer 配置对象 **核心配置类**: #### TransformerConfig 包含所有 Transformer 和并行配置: @dataclass class TransformerConfig: # 模型架构 num_layers: int = 0 hidden_size: int = 0 num_attention_heads: int = 0 ffn_hidden_size: int = None # 注意力配置 num_query_groups: int = None # GQA kv_channels: int = None attention_dropout: float = 0.1 # 激活函数 gated_linear_unit: bool = False activation_func: Callable = F.gelu # 并行配置 tensor_model_parallel_size: int = 1 pipeline_model_parallel_size: int = 1 context_parallel_size: int = 1 expert_model_parallel_size: int = 1 # 混合精度 fp16: bool = False bf16: bool = False fp8: str = None # MoE 配置 num_moe_experts: int = None moe_router_topk: int = 2 moe_grouped_gemm: bool = False # 优化 recompute_granularity: str = None recompute_method: str = None sequence_parallel: bool = False #### MLATransformerConfig 用于 Multi-Latent Attention (DeepSeek-V3): @dataclass class MLATransformerConfig(TransformerConfig): q_lora_rank: int = 1536 kv_lora_rank: int = 512 qk_nope_head_dim: int = 128 qk_rope_head_dim: int = 64 v_head_dim: int = 128 ### 3. Transformer Block ( megatron/core/transformer/transformer_block.py) **功能**: 核心 Transformer 层实现 **主要类**: - `TransformerBlock`: 主 Transformer 块 - `TransformerLayer`: 单个 Transformer 层 - `Attention`: 注意力机制 - `MLP`: 前馈网络 **特性**: - 支持多种注意力机制 (self-attention, cross-attention, MLA) - 激活检查点 - 序列并行 - CUDA Graph 支持 ### 4. 模型实现 (`megatron/core/models/`) #### GPT Model (`gpt/gpt_model.py`) class GPTModel(MegatronModule): """GPT Transformer 语言模型 Args: config: TransformerConfig transformer_layer_spec: 层规范 vocab_size: 词汇表大小 max_sequence_length: 最大序列长度 pre_process: 是否包含嵌入层 post_process: 是否包含输出层 share_embeddings_and_output_weights: 共享嵌入和输出权重 """ def __init__(self, config, transformer_layer_spec, vocab_size, ...): # 嵌入层 self.embedding = LanguageModelEmbedding(...) # Transformer 解码器 self.decoder = TransformerBlock(...) # 输出层 self.output_layer = ColumnParallelLinear(...) def forward(self, input_ids, position_ids, attention_mask, labels=None): # 嵌入 hidden_states = self.embedding(input_ids, position_ids) # Transformer 层 hidden_states = self.decoder(hidden_states, attention_mask) # 输出 logits logits = self.output_layer(hidden_states) # 计算损失 if labels is not None: loss = self.compute_language_model_loss(labels, logits) return loss return logits #### 其他模型 - **BERT** (`bert/bert_model.py`): 双向编码器 - **T5** (`T5/t5_model.py`): 编码器-解码器 - **Mamba** (`mamba/mamba_model.py`): 状态空间模型 - **Mixtral** (MoE): 专家混合模型 - **多模态** (`multimodal/`): 视觉-语言模型 (LLaVA, NVLM) --- ## 训练脚本 ### 预训练脚本 #### 1. GPT 预训练 (`pretrain_gpt.py`) def model_provider(pre_process=True, post_process=True): """构建 GPT 模型""" config = get_language_model_config(args) model = GPTModel( config=config, transformer_layer_spec=get_gpt_layer_with_transformer_engine_spec(), vocab_size=args.padded_vocab_size, max_sequence_length=args.max_position_embeddings, pre_process=pre_process, post_process=post_process ) return model def train_valid_test_datasets_provider(train_val_test_num_samples): """构建训练/验证/测试数据集""" config = GPTDatasetConfig( random_seed=args.seed, sequence_length=args.seq_length, blend=args.data_path, split=args.split ) return BlendedMegatronDatasetBuilder( GPTDataset, train_val_test_num_samples, config ).build() def forward_step(data_iterator, model): """前向传播步骤""" tokens, labels, loss_mask, attention_mask, position_ids = get_batch(data_iterator) output_tensor = model(tokens, position_ids, attention_mask, labels=labels) return output_tensor, partial(loss_func, loss_mask) if __name__ == "__main__": pretrain( train_valid_test_datasets_provider, model_provider, ModelType.encoder_or_decoder, forward_step ) #### 2. Mamba 预训练 (`pretrain_mamba.py`) from mamba_builders import mamba_builder def forward_step(data_iterator, model: MambaModel): """Mamba 前向传播""" tokens, labels, loss_mask, attention_mask, position_ids = get_batch(data_iterator) output_tensor = model(tokens, position_ids, attention_mask, labels=labels) return output_tensor, partial(loss_func, loss_mask, model=model) if __name__ == "__main__": pretrain( train_valid_test_datasets_provider, partial(model_provider, mamba_builder), ModelType.encoder_or_decoder, forward_step, args_defaults={'tokenizer_type': 'GPT2BPETokenizer'} ) #### 3. 其他预训练脚本 - `pretrain_bert.py`: BERT 预训练 - `pretrain_t5.py`: T5 预训练 - `pretrain_retro.py`: RETRO 预训练 - `pretrain_vlm.py`: 视觉-语言模型预训练 - `train_rl.py`: 强化学习训练 (RLHF) ### 训练主循环 (`megatron/training/training.py`) **核心函数**: def pretrain( train_valid_test_datasets_provider, model_provider, model_type, forward_step_func, ... ): """主训练函数 流程: 1. 初始化 Megatron 2. 设置模型、优化器和学习率调度器 3. 加载检查点 4. 训练循环 5. 保存检查点 """ # 初始化 initialize_megatron(...) # 构建模型和优化器 model, optimizer, opt_param_scheduler = setup_model_and_optimizer( model_provider, model_type ) # 加载检查点 iteration = load_checkpoint(model, optimizer, opt_param_scheduler) # 训练循环 while iteration < args.train_iters: train_step(forward_step_func, data_iterator, model, optimizer, ...) if iteration % args.save_interval == 0: save_checkpoint(iteration, model, optimizer, opt_param_scheduler) **关键训练步骤**: def train_step(forward_step_func, data_iterator, model, optimizer, ...): """单个训练步骤 1. 前向-反向传播 2. 梯度裁剪 3. 优化器步骤 4. 学习率调度 """ # 前向-反向传播 losses_reduced = forward_backward_func( forward_step_func, data_iterator, model, ... ) # 梯度裁剪 grad_norm = clip_gradients(model, optimizer, ...) # 优化器步骤 optimizer.step() # 学习率调度 opt_param_scheduler.step(increment=1) return losses_reduced --- ## 并行策略 ### 1. 数据并行 (Data Parallelism, DP) **原理**: 在多个 GPU 上复制模型，每个 GPU 处理不同的数据批次 **实现方式**: #### 标准 DDP --data-parallel-sharding-strategy no_shard #### Fully Sharded Data Parallel (FSDP) # Megatron 优化的 FSDP (~15% 比 PyTorch FSDP2 快) --use-custom-fsdp # PyTorch FSDP2 --use-torch-fsdp2 # 分片策略 --data-parallel-sharding-strategy optim # ZeRO-1: 分片优化器状态 --data-parallel-sharding-strategy optim_grads # ZeRO-2: 分片梯度 + 优化器 --data-parallel-sharding-strategy optim_grads_params # ZeRO-3: 分片参数 + 梯度 + 优化器 **代码实现** (`megatron/core/distributed/`): # 自定义 FSDP from megatron.core.distributed.fsdp import MegatronFSDP model = MegatronFSDP( module=model, sharding_strategy=ShardingStrategy.FULL_SHARD, # ZeRO-3 mixed_precision=mixed_precision_policy, ) ### 2. 张量并行 (Tensor Parallelism, TP) **原理**: 在多个 GPU 上分割单个层的张量 **使用场景**: - 单个层太大无法放入单个 GPU - 减少激活内存占用 **配置**: --tensor-model-parallel-size 4 # 4-way TP --sequence-parallel # 启用序列并行 (推荐与 TP 一起使用) **实现** (`megatron/core/tensor_parallel/`): #### 列并行线性层 class ColumnParallelLinear(torch.nn.Module): """列并行线性层 将权重矩阵按列分割: Y = XA, A = [A1 | A2 | ... | Ap] 每个 GPU 计算 Yi = X * Ai """ def forward(self, input_): # 输入在 TP 组内全收集 (如果需要) input_parallel = copy_to_tensor_model_parallel_region(input_) # 本地矩阵乘法 output_parallel = F.linear(input_parallel, self.weight) # 输出保持分割状态 return output_parallel #### 行并行线性层 class RowParallelLinear(torch.nn.Module): """行并行线性层 将权重矩阵按行分割: Y = XA, A = [A1; A2; ... ; Ap] 每个 GPU 计算 Yi = X * Ai，然后 all-reduce """ def forward(self, input_): # 输入已经分割 output_parallel = F.linear(input_, self.weight) # All-reduce 跨 TP 组 output = reduce_from_tensor_model_parallel_region(output_parallel) return output #### 序列并行 # 在序列维度上分割激活 # LayerNorm 输入: [seq_len, batch, hidden] -> [seq_len/tp, batch, hidden] class SequenceParallelLayerNorm(torch.nn.Module): def forward(self, input_): # All-gather 跨序列维度 input_gathered = gather_from_sequence_parallel_region(input_) # LayerNorm output = F.layer_norm(input_gathered, ...) # Scatter 回序列维度 output_scattered = scatter_to_sequence_parallel_region(output) return output_scattered ### 3. 流水线并行 (Pipeline Parallelism, PP) **原理**: 按深度分割模型，不同 GPU 处理不同层 **配置**: --pipeline-model-parallel-size 8 # 8 个流水线阶段 --virtual-pipeline-model-parallel-size 4 # 虚拟流水线以改善负载平衡 **流水线调度** (`megatron/core/pipeline_parallel/schedules.py`): #### 1F1B (One-Forward-One-Backward) def forward_backward_pipelining_with_interleaving( forward_step_func, data_iterator, model, num_microbatches, ... ): """交错 1F1B 流水线调度 阶段: 1. Warmup: 填充流水线 2. 1F1B: 一个前向，一个反向 3. Cooldown: 清空流水线 """ # Warmup 阶段 for i in range(num_warmup_microbatches): output_tensor = forward_step(...) # 1F1B 阶段 for i in range(num_microbatches - num_warmup_microbatches): output_tensor = forward_step(...) input_tensor_grad = backward_step(...) # Cooldown 阶段 for i in range(num_warmup_microbatches): input_tensor_grad = backward_step(...) #### P2P 通信 def send_forward(output_tensor, config): """发送激活到下一个流水线阶段""" if not is_pipeline_last_stage(): send(output_tensor, get_pipeline_model_parallel_next_rank()) def recv_forward(config): """从上一个流水线阶段接收激活""" if not is_pipeline_first_stage(): return recv(get_pipeline_model_parallel_prev_rank()) ### 4. 上下文并行 (Context Parallelism, CP) **原理**: 在多个 GPU 上分割长序列 **使用场景**: - 处理超长上下文 (>8K tokens) - 减少注意力计算的内存占用 **配置**: --context-parallel-size 2 # 2-way CP --cp-comm-type p2p # 通信类型: p2p, a2a, allgather --hierarchical-context-parallel-sizes 2 4 # 分层 CP **实现**: # Ring Attention with CP def ring_attention_forward(query, key, value, cp_group): """Ring Attention 实现 每个 GPU 持有序列的一部分: GPU 0: Q0, K0, V0 GPU 1: Q1, K1, V1 通过环形通信计算完整注意力 """ for step in range(cp_size): # 本地注意力计算 attn_output = scaled_dot_product_attention(query, key, value) # 发送 KV 到下一个 rank send_recv_kv(key, value, cp_group) ### 5. 专家并行 (Expert Parallelism, EP) **原理**: 在 MoE 模型中分割专家 **配置**: --expert-model-parallel-size 4 # 4-way EP --num-experts 8 # 每个 MoE 层 8 个专家 --moe-grouped-gemm # 优化专家计算 **实现** (`megatron/core/transformer/moe/`): class MoELayer(torch.nn.Module): """MoE 层实现 组件: 1. Router: 将 tokens 路由到专家 2. Experts: 专家网络 3. Token Dispatcher: 管理 token 分发 """ def forward(self, hidden_states): # 路由 router_logits = self.router(hidden_states) topk_weights, topk_indices = torch.topk(router_logits, k=self.top_k) # Token 分发 (All-to-All) dispatched_input = self.token_dispatcher.dispatch( hidden_states, topk_indices ) # 专家计算 (每个 GPU 上的本地专家) expert_output = self.experts(dispatched_input) # Token 收集 (All-to-All) output = self.token_dispatcher.combine( expert_output, topk_weights, topk_indices ) return output ### 并行策略组合 #### 推荐配置 (基于 NVIDIA NeMo) |模型|大小|GPUs|TP|PP|CP|EP|备注| |---|---|---|---|---|---|---|---| |LLaMA-3|8B|8|1|1|2|1|CP 用于长序列 (8K)| |LLaMA-3|70B|64|4|4|2|1|TP+PP| |LLaMA-3.1|405B|1024|8|8|2|1|3D 并行| |GPT-3|175B|128-512|4|8|1|1|大模型配置| |Mixtral|8x7B|64|1|4|1|8|EP 用于 MoE| |Mixtral|8x22B|256|4|4|8|8|TP+EP 用于大 MoE| |DeepSeek-V3|671B|1024|2|16|1|64|大 MoE 配置| **重要**: 当组合 EP 和 TP 时，**必须启用序列并行 (SP)** --- ## 数据处理 ### 数据集实现 (`megatron/core/datasets/`) #### 1. GPT Dataset (`gpt_dataset.py`) class GPTDataset(MegatronDataset): """GPT 数据集 功能: - 从二进制索引文件加载 - 动态采样和混合 - 支持文档级别的采样 """ def __getitem__(self, idx): # 获取文档索引 doc_idx = self.doc_idx[idx] # 采样序列 tokens = self.indexed_dataset[doc_idx] # 创建训练样本 tokens = tokens[:self.seq_length + 1] input_ids = tokens[:-1] labels = tokens[1:] return { 'input_ids': input_ids, 'labels': labels, 'loss_mask': torch.ones_like(labels) } #### 2. Blended Dataset (`blended_megatron_dataset_builder.py`) class BlendedMegatronDatasetBuilder: """混合数据集构建器 支持: - 多个数据集混合 - 自定义混合权重 - 训练/验证/测试分割 """ def build(self): # 构建单个数据集 datasets = [] for path, weight in zip(self.paths, self.weights): dataset = self.dataset_cls(path, ...) datasets.append((dataset, weight)) # 混合数据集 blended_dataset = BlendedDataset(datasets, self.weights) # 分割 train_ds, valid_ds, test_ds = split_dataset( blended_dataset, self.split ) return train_ds, valid_ds, test_ds ### 数据预处理 #### 工具: `tools/preprocess_data.py` # 基本预处理 python tools/preprocess_data.py \ --input data.jsonl \ --output-prefix processed_data \ --tokenizer-type HuggingFaceTokenizer \ --tokenizer-model /path/to/tokenizer.model \ --workers 8 \ --append-eod # 输出文件: # - processed_data_text_document.bin # 二进制数据 # - processed_data_text_document.idx # 索引文件 **JSONL 格式**: {"text": "Your training text here..."} {"text": "Another training sample..."} #### 索引数据集 (`indexed_dataset.py`) class MMapIndexedDataset: """内存映射索引数据集 特性: - 零拷贝加载 - 快速随机访问 - 支持大规模数据集 """ def __init__(self, path): # 加载索引 self.index = MMapIndexedDatasetIndex(path + '.idx') # 内存映射二进制文件 self.bin_buffer = np.memmap( path + '.bin', dtype=self.index.dtype, mode='r' ) def __getitem__(self, idx): # 从索引获取偏移量和大小 ptr, size = self.index[idx] # 从内存映射缓冲区读取 return self.bin_buffer[ptr:ptr+size] --- ## 优化器与调度器 ### 优化器 (`megatron/core/optimizer/`) #### 1. 分布式优化器 (`distrib_optimizer.py`) class DistributedOptimizer: """分布式优化器 特性: - 在数据并行组内分片优化器状态 - 减少内存占用 - 加速检查点保存/加载 """ def __init__(self, optimizer, config): self.optimizer = optimizer # 分片参数 self.param_to_rank = self._shard_parameters() def step(self): # 本地优化器步骤 self.optimizer.step() # All-gather 更新的参数 self._all_gather_parameters() **使用**: --use-distributed-optimizer #### 2. CPU 卸载 (`cpu_offloading/`) # 将优化器状态卸载到 CPU 以节省 GPU 内存 --cpu-optimizer #### 3. 支持的优化器 - **Adam**: 标准 Adam - **AdamW**: 带权重衰减的 Adam - **SGD**: 随机梯度下降 - **FusedAdam**: Apex 融合 Adam (更快) ### 学习率调度器 (`optimizer_param_scheduler.py`) class OptimizerParamScheduler: """优化器参数调度器 支持: - 学习率调度 - 权重衰减调度 - Warmup """ def step(self, increment): # 更新迭代计数 self.iteration += increment # 计算学习率 lr = self._get_learning_rate() # 更新优化器 for param_group in self.optimizer.param_groups: param_group['lr'] = lr **调度策略**: # 余弦衰减 --lr-decay-style cosine \ --lr 1e-4 \ --min-lr 1e-5 \ --lr-warmup-iters 1000 \ --lr-decay-iters 100000 # 线性衰减 --lr-decay-style linear # 常数学习率 --lr-decay-style constant ### 梯度裁剪 (`clip_grads.py`) def clip_grad_norm(parameters, max_norm, norm_type=2): """梯度范数裁剪 在分布式设置中: 1. 计算本地梯度范数 2. All-reduce 跨数据并行组 3. 裁剪梯度 """ # 计算本地范数 total_norm = 0.0 for p in parameters: if p.grad is not None: param_norm = p.grad.data.norm(norm_type) total_norm += param_norm ** norm_type total_norm = total_norm ** (1.0 / norm_type) # All-reduce torch.distributed.all_reduce( total_norm, group=get_data_parallel_group() ) # 裁剪 clip_coef = max_norm / (total_norm + 1e-6) if clip_coef < 1.0: for p in parameters: if p.grad is not None: p.grad.data.mul_(clip_coef) return total_norm --- ## 模型实现 ### 1. GPT 模型族 #### GPT-3 / GPT-NeoX # 175B GPT-3 --num-layers 96 \ --hidden-size 12288 \ --num-attention-heads 96 \ --seq-length 2048 \ --max-position-embeddings 2048 #### LLaMA 系列 **LLaMA-3 8B**: --num-layers 32 \ --hidden-size 4096 \ --ffn-hidden-size 14336 \ --num-attention-heads 32 \ --num-query-groups 8 \ # GQA --seq-length 8192 \ --max-position-embeddings 8192 \ --use-rotary-position-embeddings \ --swiglu \ --normalization RMSNorm **LLaMA-3 70B**: --num-layers 80 \ --hidden-size 8192 \ --ffn-hidden-size 28672 \ --num-attention-heads 64 \ --num-query-groups 8 ### 2. MoE 模型 #### Mixtral 8x7B --num-layers 32 \ --hidden-size 4096 \ --num-attention-heads 32 \ --num-query-groups 8 \ --num-experts 8 \ --moe-router-topk 2 \ --expert-model-parallel-size 8 \ --moe-grouped-gemm #### DeepSeek-V3 671B --num-layers 61 \ --hidden-size 7168 \ --num-attention-heads 128 \ --num-experts 256 \ --moe-router-topk 8 \ --expert-model-parallel-size 64 \ --use-mla \ # Multi-Latent Attention --q-lora-rank 1536 \ --kv-lora-rank 512 ### 3. Mamba (状态空间模型) # Mamba 混合模型 --num-layers 24 \ --hidden-size 2048 \ --is-hybrid-model \ --mamba-state-dim 16 \ --mamba-head-dim 64 \ --mamba-num-groups 1 ### 4. 多模态模型 #### LLaVA # 视觉编码器 + LLM --vision-model-type clip \ --vision-encoder-hidden-size 1024 \ --language-model-type llama \ --add-cross-attention --- ## 示例与工具 ### 训练示例 (`examples/`) #### 1. 简单训练循环 # 2 GPU, 模拟数据 torchrun --nproc_per_node=2 examples/run_simple_mcore_train_loop.py #### 2. LLaMA-3 训练 # 8 GPU, FP8 精度 ./examples/llama/train_llama3_8b_fp8.sh #### 3. Mixtral MoE 训练 ./examples/mixtral/train_mixtral_8x7b.sh #### 4. 多模态训练 # LLaVA 训练 ./examples/multimodal/train_llava.sh ### 工具 (`tools/`) #### 1. 数据预处理 # 文本数据 python tools/preprocess_data.py --input data.jsonl ... # 多模态数据 python tools/preprocess_mmdata.py --input mm_data.jsonl ... #### 2. 检查点转换 # Megatron -> Hugging Face python tools/checkpoint/convert_checkpoint.py \ --model-type GPT \ --load-dir /path/to/megatron/checkpoint \ --save-dir /path/to/hf/checkpoint \ --target-tensor-parallel-size 1 \ --target-pipeline-parallel-size 1 #### 3. 推理服务器 # GPT 推理 python tools/run_text_generation_server.py \ --load /path/to/checkpoint \ --tensor-model-parallel-size 4 # Mamba 推理 python tools/run_mamba_text_generation_server.py \ --load /path/to/checkpoint #### 4. 理论内存计算 # 估算模型内存需求 python tools/report_theoretical_memory.py \ --num-layers 32 \ --hidden-size 4096 \ --num-attention-heads 32 \ --seq-length 2048 \ --batch-size 8 \ --tensor-parallel-size 4 \ --pipeline-parallel-size 2 --- ## 性能优化 ### 1. 混合精度训练 #### FP16 --fp16 #### BF16 (推荐用于大模型) --bf16 #### FP8 (Hopper, Ada, Blackwell GPUs) --fp8-hybrid \ --fp8-recipe hybrid # 或 'delayed_scaling' ### 2. FlashAttention # 默认: cuDNN attention via Transformer Engine # 提供高达 50% 的前向加速和 84% 的反向加速 # 使用 flash-attn 包 --use-flash-attn ### 3. 激活检查点 # 基本激活检查点 --recompute-activations # 完全重计算 (极端内存约束) --recompute-granularity full \ --recompute-method uniform # 选择性重计算 --recompute-granularity selective \ --recompute-method block ### 4. 通信优化 #### 数据并行通信重叠 --overlap-grad-reduce # 重叠梯度 reduce --overlap-param-gather # 重叠参数 gather #### 张量并行通信重叠 --tp-comm-overlap #### 流水线并行通信重叠 # 默认启用 ### 5. CUDA Graph --cuda-graph ### 6. 融合内核 --use-fused-swiglu # 融合 SwiGLU 激活 --use-fused-rotary-pos-emb # 融合 RoPE ### 7. 分布式优化器 --use-distributed-optimizer ### 8. MoE 优化 --moe-grouped-gemm # GroupedGEMM --moe-permute-fusion # 融合排列内核 --moe-router-fusion # 融合路由器 --moe-expert-parallel-overlap # EP All-to-All 重叠 ### 性能基准 **配置**: - 词汇表大小: 131,072 - 序列长度: 4096 - H100 GPUs **结果**: |模型大小|GPUs|MFU|吞吐量 (tokens/s/GPU)| |---|---|---|---| |22B|128|41%|~3200| |175B|1024|47%|~2800| |462B|6144|48%|~2600| **超线性扩展**: MFU 从 41% (22B) 增加到 47-48% (175B+)，因为更大的 GEMM 具有更高的算术强度。 --- ## 高级特性 ### 1. 容错训练 (NVRx) # 启用故障检测和恢复 --enable-fault-tolerance \ --check-for-nan-in-loss-and-grad \ --check-for-spiky-loss ### 2. 多数据中心训练 # N/S 连接支持 --multi-datacenter-training ### 3. 检查点 #### 保存 --save /path/to/checkpoints \ --save-interval 1000 #### 加载 --load /path/to/checkpoints #### 分布式检查点 --use-dist-ckpt ### 4. 日志和监控 #### TensorBoard --tensorboard-dir /path/to/logs #### Weights & Biases --wandb-project my-project \ --wandb-exp-name my-experiment ### 5. 评估 # 验证间隔 --eval-interval 1000 \ --eval-iters 100 --- ## 生态系统 ### 使用 Megatron Core 的库 1. **[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)** - Hugging Face ↔ Megatron 检查点转换 - 灵活的训练循环 - 生产就绪的配方 2. **[NeMo RL](https://github.com/NVIDIA-NeMo/RL)** - RLHF, DPO 等后训练方法 - 可扩展的强化学习 3. **[NeMo Framework](https://docs.nvidia.com/nemo-framework/)** - 企业框架 - 云原生支持 - 端到端示例 4. **[TensorRT Model Optimizer](https://github.com/NVIDIA/TensorRT-Model-Optimizer)** - 量化、剪枝、蒸馏 ### Megatron Core 使用的库 1. **[Megatron Energon](https://github.com/NVIDIA/Megatron-Energon)** - 多模态数据加载器 (文本、图像、视频、音频) 2. **[Transformer Engine](https://github.com/NVIDIA/TransformerEngine)** - 优化内核 - FP8 混合精度 3. **[NVRx](https://github.com/NVIDIA/nvidia-resiliency-ext)** - 容错训练 --- ## 常见问题 ### Q1: 如何选择并行策略？ **答**: - **小模型 (<10B)**: DP only - **中等模型 (10B-100B)**: TP + PP + DP - **大模型 (>100B)**: TP + PP + DP + CP - **MoE 模型**: 添加 EP ### Q2: 内存不足怎么办？ **答**: 1. 启用激活检查点: `--recompute-activations` 2. 使用 FSDP: `--use-custom-fsdp --data-parallel-sharding-strategy optim_grads_params` 3. 增加 PP: `--pipeline-model-parallel-size` 4. 减少批次大小: `--micro-batch-size` 5. CPU 卸载: `--cpu-optimizer` ### Q3: 如何加速训练？ **答**: 1. 使用 FP8: `--fp8-hybrid` 2. 启用通信重叠: `--overlap-grad-reduce --overlap-param-gather --tp-comm-overlap` 3. 使用 FlashAttention: 默认启用 4. 启用 CUDA Graph: `--cuda-graph` 5. 使用融合内核: `--use-fused-swiglu --use-fused-rotary-pos-emb` ### Q4: 如何调试？ **答**: # 启用详细日志 --log-level DEBUG # 检查 NaN/Inf --check-for-nan-in-loss-and-grad # 使用模拟数据 --mock-data # 单 GPU 测试 --tensor-model-parallel-size 1 \ --pipeline-model-parallel-size 1 --- ## 参考资源 ### 官方文档 - [Megatron Core 文档](https://docs.nvidia.com/megatron-core/) - [API 指南](https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/) - [开发者指南](https://docs.nvidia.com/megatron-core/developer-guide/latest/) ### 论文 - [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://arxiv.org/abs/1909.08053) - [Efficient Large-Scale Language Model Training on GPU Clusters](https://arxiv.org/abs/2104.04473) ### GitHub - [Megatron-LM 仓库](https://github.com/NVIDIA/Megatron-LM) - [Issues](https://github.com/NVIDIA/Megatron-LM/issues) ### 社区 - [NVIDIA 开发者论坛](https://forums.developer.nvidia.com/) --- ## 更新日志 ### v0.14.0 (2025) - MoE 性能优化 (路由器融合、EP A2A 重叠) - FP8 优化 (填充优化、重计算支持) - 推理改进 (多批次 CUDA Graph、FP8 权重) - MiMo 视频 VLM 支持 ### v0.13.0 - 灵活的非对称虚拟流水线并行 - 混合分片数据并行 (MoE) - FP8 支持 Multi-Token-Prediction - 细粒度重计算 ### v0.12.0 - FP8 配方选择 - 上下文并行损失缩放修复 - 混合模型 CUDA Graph 支持 - MoE CUDA Graph 支持 --- ## 总结 Megatron-LM 是一个功能强大、高度优化的大规模 Transformer 模型训练框架。它提供: ✅ **完整的并行策略**: TP, PP, DP, CP, EP ✅ **最先进的性能**: 47% MFU on H100 ✅ **丰富的模型支持**: GPT, LLaMA, Mixtral, Mamba, 多模态 ✅ **生产就绪**: 容错、检查点、监控 ✅ **可扩展性**: 从 1 GPU 到 6000+ GPUs 无论您是研究人员、工程师还是框架开发者，Megatron-LM 都提供了构建和训练大规模语言模型所需的所有工具。