Notes/机器学习/深度学习模型/07-计算架构/5-分布式训练框架/2-Megatron-LM.md

Megatron-LM

概述

Megatron-LM是NVIDIA开发的大规模Transformer训练框架，专门针对超大规模语言模型的张量并行（Tensor Parallelism）和流水线并行（Pipeline Parallelism）优化。

核心特性

特性	说明
张量并行	实现Column切分和Row切分，支持超大规模单层切分
序列并行	解决序列维度上的显存问题
流水线并行	优化的1F1B调度
混合并行	支持TP+PP+DP组合
通信优化	NVLink优化，AllReduce融合

张量并行实现

ColumnParallelLinear

用于QKV投影、FFN第一个线性层：

class ColumnParallelLinear(torch.nn.Module):
    """
    将权重按列切分，每个GPU持有 (out_features / TP, in_features)
    """
    
    def __init__(self, in_features, out_features, tp_size):
        super().__init__()
        self.tp_size = tp_size
        self.output_size_per_partition = out_features // tp_size
        
        # 分片参数
        self.weight = torch.nn.Parameter(
            torch.randn(self.output_size_per_partition, in_features)
        )
        
    def forward(self, x):
        # x: (batch, seq, in_features) 完整输入
        # 在每个GPU上独立计算
        y = F.linear(x, self.weight)
        
        # AllReduce合并结果（因为输入相同，输出应求和）
        # 但Column切分时输出沿hidden维度拼接，不需要AllReduce
        # 实际是拼接操作
        return y

RowParallelLinear

用于FFN第二个线性层：

class RowParallelLinear(torch.nn.Module):
    """
    将权重按行切分，每个GPU持有 (out_features, in_features / TP)
    """
    
    def __init__(self, in_features, out_features, tp_size):
        super().__init__()
        self.tp_size = tp_size
        self.input_size_per_partition = in_features // tp_size
        
        self.weight = torch.nn.Parameter(
            torch.randn(out_features, self.input_size_per_partition)
        )
        
    def forward(self, x):
        # x被切分到各GPU (batch, seq, in_features / TP)
        # 每个GPU计算部分结果
        y_parallel = F.linear(x, self.weight)
        
        # AllReduce求和（因为结果需要汇总）
        y = tensor_model_parallel_all_reduce(y_parallel)
        
        return y

通信操作

# Tensor并行通信原语

def tensor_model_parallel_all_reduce(tensor):
    """TP域内的AllReduce"""
    return torch.distributed.all_reduce(
        tensor, op=torch.distributed.ReduceOp.SUM,
        group=tensor_model_parallel_group
    )

def broadcast_from_first_rank(tensor):
    """从TP组第一个GPU广播到所有GPU"""
    return torch.distributed.broadcast(
        tensor,
        src=first_rank_in_group,
        group=tensor_model_parallel_group
    )

def all_gather_coalesced(tensor):
    """收集TP组内所有GPU的tensor"""
    return torch.distributed.all_gather(
        tensor, group=tensor_model_parallel_group
    )

序列并行（Sequence Parallelism）

问题背景

在张量并行中，LayerNorm等操作需要对完整序列做归约，但各GPU只有部分数据：

TP=8时，每个GPU只有 seq_len/8 长度的数据
但LayerNorm需要 seq_len 长度的统计量

解决方案

class LayerNorm(nn.Module):
    """序列并行版本的LayerNorm"""
    
    def forward(self, x):
        # x: (batch, seq/tp, hidden)
        
        # 1. 收集完整序列用于计算均值和方差
        x_gather = all_gather_sequence_parallel(x)  # (batch, seq, hidden)
        
        # 2. 计算统计量
        mean = x_gather.mean(dim=-2, keepdim=True)  # (batch, 1, hidden)
        var = x_gather.var(dim=-2, keepdim=True)   # (batch, 1, hidden)
        
        # 3. 广播回各GPU
        mean = broadcast_to_sequence_parallel_region(mean)
        var = broadcast_to_sequence_parallel_region(var)
        
        # 4. 归一化
        x_norm = (x - mean) / sqrt(var + eps)
        
        return x_norm

收益分析

显存节省：
- LayerNorm输入: batch × seq/tp × hidden → 节省TP倍
- Attention score: batch × heads/tp × seq/tp × seq/tp → 节省TP²倍

通信开销：
- AllGather用于归一化统计
- 与计算重叠后开销可忽略

模型切分策略

整体切分方案

def get_tensor_model_parallel_world_size():
    return tp_size

def get_pipeline_model_parallel_world_size():
    return pp_size

def get_data_parallel_world_size():
    return dp_size

# 计算全局GPU数
total_gpus = tp_size * pp_size * dp_size

Transformer层内切分

Layer结构:
[Input] → [Input LayerNorm] → [Attention] → [+] → [Post Attention LayerNorm] → [MLP] → [+]
                    ↓                    ↓
            Column切分(QKV)        Row切分(Projection)
            Column切分(MLP1)       Row切分(MLP2)

各TP程度的切分方式

TP Size	切分层	说明
1	无	标准模型
2	QKV, MLP1	2路切分
4	QKV, MLP1	4路切分
8	QKV, MLP1	8路切分（需NVLink）

负载均衡

切分策略选择原则：
1. 计算量均衡：每GPU计算量接近
2. 通信均衡：避免通信成为瓶颈
3. 显存均衡：各GPU峰值显存接近

对于Transformer：
- Attention: 计算量 ∝ 4×batch×seq²×heads×head_dim
- MLP: 计算量 ∝ 6×batch×seq×hidden×intermediate

均衡切分需考虑注意力头的划分

与DeepSpeed对比

特性	Megatron-LM	DeepSpeed
张量并行	原生支持，优化完善	支持但非核心focus
流水线并行	支持，优化调度	支持
ZeRO	不支持（使用FSDP替代）	完整实现
序列并行	原生支持	有限支持
CPU Offload	有限	完善
3D并行	TP+PP原生，DP需配合	完整支持
适用场景	超大规模单节点/少数节点	大规模多节点

架构差异

Megatron-LM设计哲学：
- 张量并行是核心，专注于单层内的高效切分
- 认为通信是瓶颈，优化AllReduce模式
- 对NVLink等高速互联友好

DeepSpeed设计哲学：
- 数据并行是核心，ZeRO解决显存问题
- 通过分片减少通信需求
- 通过Offload支持超大规模

Megatron核心实现

并行注意力

class ParallelSelfAttention(nn.Module):
    """
    序列并行的自注意力
    """
    
    def __init__(self, hidden_size, num_heads, tp_size):
        super().__init__()
        self.tp_size = tp_size
        self.num_heads_per_partition = num_heads // tp_size
        
        # QKV投影（Column切分）
        self.query_key_value = ColumnParallelLinear(
            hidden_size,
            3 * hidden_size,  # Q, K, V 拼接
            tp_size
        )
        
        # 输出投影（Row切分）
        self.dense = RowParallelLinear(
            hidden_size,
            hidden_size,
            tp_size
        )
        
    def forward(self, x, attention_mask=None):
        # x: (batch, seq/tp, hidden)
        
        # QKV计算
        qkv = self.query_key_value(x)
        q, k, v = qkv.split(hidden_size, dim=-1)
        
        # 收集完整K, V（用于跨TP组计算注意力）
        # 需要AllGather获取其他GPU的K, V
        k = all_gather_kv_along_seq_dim(k)
        v = all_gather_kv_along_seq_dim(v)
        
        # 注意力计算
        attn_output = scaled_dot_product_attention(q, k, v)
        
        # 输出投影
        output = self.dense(attn_output)
        
        return output

环形通信优化

def ring_send_recv(k_chunks, v_chunks):
    """
    环形通信计算注意力
    用于减少AllGather通信量
    """
    num_chunks = len(k_chunks)
    local_k, local_v = k_chunks[0], v_chunks[0]
    
    outputs = []
    
    for i in range(num_chunks):
        # 计算当前chunk的注意力
        attn = scaled_dot_product_attention(q, local_k, local_v)
        outputs.append(attn)
        
        if i < num_chunks - 1:
            # 接收下一个chunk（环形）
            recv_idx = (rank + 1) % num_chunks
            local_k = k_chunks[recv_idx]
            local_v = v_chunks[recv_idx]
    
    return outputs

Megatron训练配置

# Megatron训练配置示例
megatron_config = {
    "tensor_model_parallel_size": 8,
    "pipeline_model_parallel_size": 4,
    "num_layers": 32,
    "hidden_size": 4096,
    "num_attention_heads": 32,
    "ffn_hidden_size": 16384,
    "micro_batch_size": 4,
    "global_batch_size": 1024,
    "lr": 1e-4,
    "seq_len": 2048,
    "bf16": True
}

# 启动训练
# python train.py --config megatron_config.yaml

分布式优化器

# Megatron的分布式优化器
# 每个GPU持有优化器状态的1/TP分片

class MegatronOptimizer:
    """
    分片优化器状态
    """
    
    def __init__(self, model, optimizer_class, tp_size):
        self.tp_size = tp_size
        
        # 优化器状态按TP分片
        for param in model.parameters():
            param.state_full = None  # 未分片的完整状态
            param.state_sharded = self.shard_tensor(param)  # 分片后状态
    
    def shard_tensor(self, tensor):
        # 沿第一个维度切分
        shard_size = tensor.shape[0] // self.tp_size
        shard = tensor[shard_size * rank : shard_size * (rank + 1)]
        return shard
    
    def step(self):
        # 更新参数
        for param in self.params:
            # 本地更新
            param_sharded = self.update_local(param.state_sharded)
            # AllReduce同步更新
            param_updated = all_reduce(param_sharded)
            # 写回完整参数（如需要）

最佳实践

硬件配置建议

GPU数	TP	PP	备注
8	8	1	最大TP单节点
16	8	2	TP+PP组合
32+	8	4	需要网络优化

通信优化

# 启用融合通信
model = parallel_model(model)

# 使用NVLink通信
# 确认 NCCL_ALGO=RING
# 确认 NCCL_IB_HCA=mlx5_0,mlx5_1

调试技巧

# 查看张量并行等级
echo $NV_TENSOR_MODEL_PARALLEL_SIZE

# 查看流水线等级
echo $PIPELINE_MODEL_PARALLEL_SIZE

# 检查通信性能
python -c "import torch.distributed as dist; dist.is_available()"

# 启用详细日志
export PYTHONFAULTHANDLER=1
export NCCL_DEBUG=INFO

性能监控

# Megatron性能指标
class PerformanceMonitor:
    """监控TP训练的性能指标"""
    
    def __init__(self):
        self.forward_time = []
        self.backward_time = []
        self.comm_time = []
    
    def record_forward(self, duration):
        self.forward_time.append(duration)
    
    def compute_stats(self):
        return {
            "avg_forward": mean(self.forward_time),
            "avg_backward": mean(self.backward_time),
            "throughput": batch_size * seq_len / mean(self.forward_time)
        }

框架生态

NVIDIA NGC Container: 包含优化好的Megatron镜像
├── Megatron-LM (主干)
├── NeMo-Megatron (企业级封装)
├── PyTorch Lightning (集成)
└── Triton (算子优化)

常见问题

问题	原因	解决方案
TP通信瓶颈	跨节点TP	限制TP在节点内
梯度不一致	AllReduce未同步	检查梯度分片逻辑
显存不均	参数未均匀切分	检查TP分片
速度慢	Occupancy低	调整batch size/seq len