446 lines
11 KiB
Markdown
446 lines
11 KiB
Markdown
# 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 | |