# Megatron-LM

## 概述

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

### 核心特性

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

## 张量并行实现

### ColumnParallelLinear

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

```python
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第二个线性层：

```python
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
```

### 通信操作

```python
# 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 长度的统计量
```

### 解决方案

```python
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用于归一化统计
- 与计算重叠后开销可忽略
```

## 模型切分策略

### 整体切分方案

```python
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核心实现

### 并行注意力

```python
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
```

### 环形通信优化

```python
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训练配置

```python
# 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
```

### 分布式优化器

```python
# 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 | 需要网络优化 |

### 通信优化

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

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

### 调试技巧

```bash
# 查看张量并行等级
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
```

### 性能监控

```python
# 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 |