4. 张量并行实现
任务目标
在单机多GPU环境下，基于Megatron-LM的核心思想实现Transformer层的张量并行：将多头注意力的QKV投影矩阵按头维度切分到不同GPU，每个GPU只计算自己负责的注意力头，最后通过All-Gather合并输出。训练一个参数约50万的小型模型，验证张量并行训练与单卡训练在损失下降趋势上一致。
环境与依赖
- Python 3.8+，PyTorch 1.13+（需支持 NCCL）
- 2张可用GPU（可通过
torch.cuda.device_count() 确认）
- 安装：
pip install torch numpy
核心实现
以下代码实现了一个张量并行的单层Transformer块。关键点：QKV 权重按输出维度切分为 tp_size 份，每份包含 num_heads / tp_size 个注意力头；前向传播时各GPU独立计算自己的注意力头，输出通过 AllGather 合并；MLP的中间层也按列切分，输出层按行切分并All-Reduce求和。
```python
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributed import init_process_group, all_gather, all_reduce, get_world_size, get_rank
# ---------- 初始化分布式环境 ----------
def setup(backend="nccl"):
init_process_group(backend=backend)
torch.cuda.set_device(get_rank())
# ---------- 张量并行线性层 ----------
class ColumnParallelLinear(nn.Module):
"""按输出维度切分：Y = X @ W^T，W的shape为[out_dim/tp, in_dim]"""
def __init__(self, in_dim, out_dim, tp_size):
super().__init__()
self.tp_size = tp_size
self.out_dim_per_tp = out_dim // tp_size
self.weight = nn.Parameter(torch.randn(self.out_dim_per_tp, in_dim) * 0.02)
def forward(self, x):
return F.linear(x, self.weight)
class RowParallelLinear(nn.Module):
"""按输入维度切分，输出All-Reduce求和：Y = X @ W^T，W的shape为[out_dim, in_dim/tp]"""
def __init__(self, in_dim, out_dim, tp_size):
super().__init__()
self.tp_size = tp_size
self.in_dim_per_tp = in_dim // tp_size
self.weight = nn.Parameter(torch.randn(out_dim, self.in_dim_per_tp) * 0.02)
def forward(self, x):
out = F.linear(x, self.weight)
all_reduce(out)
return out
# ---------- 张量并行注意力 ----------
class TPAttention(nn.Module):
def __init__(self, hidden_dim, num_heads, tp_size):
super().__init__()
assert hidden_dim % num_heads == 0
assert num_heads % tp_size == 0
self.hidden_dim = hidden_dim
self.num_heads = num_heads
self.tp_size = tp_size
self.head_dim = hidden_dim // num_heads
self.heads_per_tp = num_heads // tp_size
# 每个GPU只保存自己负责的头的QKV投影
self.qkv = ColumnParallelLinear(hidden_dim, 3 * hidden_dim, tp_size)
self.out_proj = RowParallelLinear(hidden_dim, hidden_dim, tp_size)
def forward(self, x):
bsz, seq_len, _ = x.shape
# qkv: [bsz, seq_len, 3 * hidden_dim / tp]
qkv = self.qkv(x)
qkv = qkv.reshape(bsz, seq_len, 3, self.heads_per_tp, self.head_dim)
q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] # 各 [bsz, seq, h_per_tp, head_dim]
q = q.transpose(1, 2) # [bsz, h_per_tp, seq, head_dim]
k = k.transpose(1, 2)
v = v.transpose(1, 2)
attn_weights = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
attn_weights = F.softmax(attn_weights, dim=-1)
attn_out = torch.matmul(attn_weights, v) # [bsz, h_per_tp, seq, head_dim]
attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, self.heads_per_tp * self.head_dim)
# All-Gather合并所有GPU的注意力头
gathered = [torch.empty_like(attn_out) for _ in range(self.tp_size)]
all_gather(gathered, attn_out)
attn_out = torch.cat(gathered, dim=-1) # [bsz, seq, hidden_dim]
return self.out_proj(attn_out)
# ---------- 张量并行MLP ----------
class TPMLP(nn.Module):
def __init__(self, hidden_dim, ffn_dim, tp_size):
super().__init__()
self.fc1 = ColumnParallelLinear(hidden_dim, ffn_dim, tp_size)
self.fc2 = RowParallelLinear(ffn_dim, hidden_dim, tp_size)
def forward(self, x):
return self.fc2(F.gelu(self.fc1(x)))
# ---------- 张量并行Transformer层 ----------
class TPTransformerLayer(nn.Module):
def __init__(self, hidden_dim, num_heads, ffn_dim, tp_size):
super().__init__()
self.attn = TPAttention(hidden_dim, num_heads, tp_size)
self.mlp = TPMLP(hidden_dim, ffn_dim, tp_size)
self.ln1 = nn.LayerNorm(hidden_dim)
self.ln2 = nn.LayerNorm(hidden_dim)
def forward(self, x):
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
return x
# ---------- 小型模型 ----------
class SmallGPT(nn.Module):
def __init__(self, vocab_size=256, hidden_dim=128, num_heads=8, ffn_dim=512, num_layers=2, tp_size=2):
super().__init__()
self.tp_size = tp_size
self.embed = nn.Embedding(vocab_size, hidden_dim)
self.layers = nn.ModuleList([
TPTransformerLayer(hidden_dim, num_heads, ffn_dim, tp_size) for _ in range(num_layers)
])
self.ln_f = nn.LayerNorm(hidden_dim)
self.lm_head = nn.Linear(hidden_dim, vocab_size, bias=False)
self.embed.weight = self.lm_head.weight # 权重绑定
def forward(self, input_ids):
x = self.embed(input_ids)
for layer in self.layers:
x = layer(x)
x = self.ln_f(x)
logits = self.lm_head(x)
return logits
# ---------- 训练入口 ----------
def train():
setup()
tp_size = get_world_size()
rank = get_rank()
torch.manual_seed(42 + rank)
model = SmallGPT(tp_size=tp_size).cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
# 构造随机训练数据（batch=4, seq_len=32）
input_ids = torch.randint(0, 256, (4, 32)).cuda()
targets = torch.randint(0, 256, (4, 32)).cuda()
model.train()
for step in range(50):
optimizer.zero_grad()
logits = model(input_ids) # [4, 32, 256]
loss = F.cross_entropy(logits.view(-1, 256), targets.view(-1))
loss.backward()
optimizer.step()
if rank == 0 and step % 10 == 0:
print(f"step {step}, loss = {loss.item():.4f}")
# 验证：保存模型输出，供单卡对比
if rank == 0:
torch.save(model.state_dict(), "tp_model.pt")
print("训练完成，模型已保存至 tp_model.pt")
if __name__ == "__main__":
train()
```
运行方式
- 将上述代码保存为
tp_。
- 使用
torchrun 启动2卡训练：
torchrun --nproc_per_node=2 tp_
验证方法
- 训练收敛验证：观察rank 0打印的loss，应随step增加而下降（例如从约5.5降至约4.0以下），说明张量并行训练正常更新参数。
- 并行正确性验证：将
tp_size 改为1（即单卡模式），用相同随机种子和相同数据训练相同步数，对比两种模式下每步的loss值。若张量并行实现正确，两种模式的loss应完全一致（误差<1e-4）。
- GPU利用率验证：训练时运行
nvidia-smi，确认2张GPU均有显存占用和计算活动。
常见问题
- All-Gather后维度不匹配：确保
num_heads % tp_size == 0，且每个GPU的 heads_per_tp * head_dim * tp_size == hidden_dim。
- 权重初始化不一致：不同rank需使用不同随机种子初始化，否则各GPU权重完全相同，All-Reduce/All-Gather结果虽正确但无法体现并行切分效果。
- OOM问题：若显存不足，可减小
batch_size 或 seq_len，本案例默认配置约需2×2GB显存。