ai阅读megatron

vllbc 收录于 Pai-Megatron Coding

2025-11-20 约 7094 字预计阅读 15 分钟次阅读

Megatron-LM 详尽文档

项目概述

什么是 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

并行策略组合

模型	大小	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 配置

数据处理

数据集实现 (`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 的库

Megatron Bridge
- Hugging Face ↔︎ Megatron 检查点转换
- 灵活的训练循环
- 生产就绪的配方
NeMo RL
- RLHF, DPO 等后训练方法
- 可扩展的强化学习
NeMo Framework
- 企业框架
- 云原生支持
- 端到端示例
TensorRT Model Optimizer
- 量化、剪枝、蒸馏

Megatron Core 使用的库

Megatron Energon
- 多模态数据加载器 (文本、图像、视频、音频)
Transformer Engine
- 优化内核
- FP8 混合精度
NVRx
- 容错训练

常见问题

Q1: 如何选择并行策略？

答:

小模型 (<10B): DP only
中等模型 (10B-100B): TP + PP + DP
大模型 (>100B): TP + PP + DP + CP
MoE 模型: 添加 EP

Q2: 内存不足怎么办？

答:

启用激活检查点: --recompute-activations
使用 FSDP: --use-custom-fsdp --data-parallel-sharding-strategy optim_grads_params
增加 PP: --pipeline-model-parallel-size
减少批次大小: --micro-batch-size
CPU 卸载: --cpu-optimizer

Q3: 如何加速训练？

答:

使用 FP8: --fp8-hybrid
启用通信重叠: --overlap-grad-reduce --overlap-param-gather --tp-comm-overlap
使用 FlashAttention: 默认启用
启用 CUDA Graph: --cuda-graph
使用融合内核: --use-fused-swiglu --use-fused-rotary-pos-emb

Q4: 如何调试？

答:

启用详细日志

–log-level DEBUG

检查 NaN/Inf