3. Transformer块中的张量并行
任务目标
本案例将Transformer块中的注意力与MLP层按张量并行方式切分到2个GPU设备上，并实现前向传播所需的通信原语（All-Reduce）。核心目标是理解列并行（Column Parallel）与行并行（Row Parallel）的切分方式，以及它们如何组合成完整的Transformer块。
环境与依赖
- Python 3.8+，仅依赖NumPy（用于张量运算模拟）
- 无需真实多GPU环境，通过NumPy模拟2个设备上的分片计算
实现原理
注意力层与MLP层均采用“列并行→行并行”的组合切分策略。以MLP为例：第一个线性层按列切分权重（每设备持有完整输入、部分输出），第二个线性层按行切分权重（每设备持有部分输入、部分输出）。两个设备各自完成局部计算后，通过All-Reduce求和得到与单设备等价的结果。注意力层的QKV投影采用列并行，输出投影采用行并行，通信模式与MLP完全一致。
关键完整代码
import numpy as np
# ========== 通信原语 ==========
def all_reduce_sum(local_tensor, world_size=2):
"""模拟All-Reduce SUM：各设备本地结果求和后广播"""
# 真实场景中此处为NCCL/MPI调用；这里用全量求和模拟
global_sum = np.sum(local_tensor, axis=0) # 假设沿设备维度求和
return global_sum
# ========== 张量并行线性层 ==========
class ColumnParallelLinear:
"""列并行：输入完整，输出按列切分到各设备"""
def __init__(self, in_features, out_features, rank, world_size=2):
self.rank = rank
self.world_size = world_size
self.out_per_device = out_features // world_size
# 每设备持有部分行（对应输出特征的部分列）
self.weight = np.random.randn(self.out_per_device, in_features) * 0.1
self.bias = np.random.randn(self.out_per_device) * 0.1
def forward(self, x):
# x: [batch, in_features]
return x @ self.weight.T + self.bias # [batch, out_per_device]
class RowParallelLinear:
"""行并行：输入按列切分，输出完整"""
def __init__(self, in_features, out_features, rank, world_size=2):
self.rank = rank
self.world_size = world_size
self.in_per_device = in_features // world_size
# 每设备持有部分列（对应输入特征的部分行）
self.weight = np.random.randn(out_features, self.in_per_device) * 0.1
self.bias = np.random.randn(out_features) * 0.1
def forward(self, x_partial):
# x_partial: [batch, in_per_device]
return x_partial @ self.weight.T # [batch, out_features]，bias在All-Reduce后加
# ========== Transformer块（张量并行版） ==========
class TensorParallelTransformerBlock:
def __init__(self, hidden_size, num_heads, rank, world_size=2):
self.rank = rank
self.world_size = world_size
self.hidden_size = hidden_size
self.head_dim = hidden_size // num_heads
self.num_heads = num_heads
# 注意力：QKV列并行，输出行并行
self.qkv_proj = ColumnParallelLinear(hidden_size, 3 * hidden_size, rank, world_size)
self.attn_out_proj = RowParallelLinear(hidden_size, hidden_size, rank, world_size)
# MLP：第一个线性层列并行（4x扩展），第二个行并行
self.mlp_up = ColumnParallelLinear(hidden_size, 4 * hidden_size, rank, world_size)
self.mlp_down = RowParallelLinear(4 * hidden_size, hidden_size, rank, world_size)
# LayerNorm参数（每设备完整副本）
self.ln1_gamma = np.ones(hidden_size)
self.ln1_beta = np.zeros(hidden_size)
self.ln2_gamma = np.ones(hidden_size)
self.ln2_beta = np.zeros(hidden_size)
def layer_norm(self, x, gamma, beta, eps=1e-5):
mean = x.mean(axis=-1, keepdims=True)
var = x.var(axis=-1, keepdims=True)
return (x - mean) / np.sqrt(var + eps) * gamma + beta
def attention(self, x):
batch = x.shape[0]
# 列并行QKV投影
qkv_local = self.qkv_proj.forward(x) # [batch, 3*hidden/2]（每设备）
# 本地切分Q、K、V（每设备处理部分头）
local_dim = self.hidden_size // self.world_size
q_local = qkv_local[:, :local_dim].reshape(batch, -1, self.head_dim)
k_local = qkv_local[:, local_dim:2*local_dim].reshape(batch, -1, self.head_dim)
v_local = qkv_local[:, 2*local_dim:3*local_dim].reshape(batch, -1, self.head_dim)
# 局部注意力计算
scores = np.einsum('bhd,bhd->bh', q_local, k_local) / np.sqrt(self.head_dim)
attn_weights = np.exp(scores - scores.max(axis=-1, keepdims=True))
attn_weights /= attn_weights.sum(axis=-1, keepdims=True)
context_local = attn_weights[..., None] * v_local # [batch, heads_local, head_dim]
context_local = context_local.reshape(batch, -1)
# 行并行输出投影（输入为部分上下文）
return self.attn_out_proj.forward(context_local)
def forward(self, x):
# x: [batch, hidden_size] 完整输入
# 注意力子层
residual = x
x = self.layer_norm(x, self.ln1_gamma, self.ln1_beta)
attn_out = self.attention(x) # 每设备输出完整hidden_size（行并行）
# All-Reduce：行并行输出需要跨设备求和
attn_out = all_reduce_sum(np.stack([attn_out, np.zeros_like(attn_out)])) # 模拟2设备
x = residual + attn_out + self.attn_out_proj.bias
# MLP子层
residual = x
x = self.layer_norm(x, self.ln2_gamma, self.ln2_beta)
# 列并行：每设备输出部分4x隐藏单元
mlp_hidden_local = self.mlp_up.forward(x) # [batch, 2*hidden]（每设备）
mlp_hidden_local = np.maximum(mlp_hidden_local, 0) # GELU近似用ReLU
# 行并行：输入为部分隐藏单元
mlp_out = self.mlp_down.forward(mlp_hidden_local) # [batch, hidden]
# All-Reduce
mlp_out = all_reduce_sum(np.stack([mlp_out, np.zeros_like(mlp_out)]))
x = residual + mlp_out + self.mlp_down.bias
return x
# ========== 单设备参考实现（用于验证） ==========
class ReferenceTransformerBlock:
def __init__(self, hidden_size, num_heads):
self.hidden_size = hidden_size
self.num_heads = num_heads
self.head_dim = hidden_size // num_heads
self.qkv_weight = np.random.randn(3 * hidden_size, hidden_size) * 0.1
self.qkv_bias = np.random.randn(3 * hidden_size) * 0.1
self.attn_out_weight = np.random.randn(hidden_size, hidden_size) * 0.1
self.attn_out_bias = np.random.randn(hidden_size) * 0.1
self.mlp_up_weight = np.random.randn(4 * hidden_size, hidden_size) * 0.1
self.mlp_up_bias = np.random.randn(4 * hidden_size) * 0.1
self.mlp_down_weight = np.random.randn(hidden_size, 4 * hidden_size) * 0.1
self.mlp_down_bias = np.random.randn(hidden_size) * 0.1
self.ln1_gamma = np.ones(hidden_size)
self.ln1_beta = np.zeros(hidden_size)
self.ln2_gamma = np.ones(hidden_size)
self.ln2_beta = np.zeros(hidden_size)
def layer_norm(self, x, gamma, beta, eps=1e-5):
mean = x.mean(axis=-1, keepdims=True)
var = x.var(axis=-1, keepdims=True)
return (x - mean) / np.sqrt(var + eps) * gamma + beta
def attention(self, x):
batch = x.shape[0]
qkv = x @ self.qkv_weight.T + self.qkv_bias
q, k, v = np.split(qkv, 3, axis=-1)
q = q.reshape(batch, self.num_heads, self.head_dim)
k = k.reshape(batch, self.num_heads, self.head_dim)
v = v.reshape(batch, self.num_heads, self.head_dim)
scores = np.einsum('bhd,bhd->bh', q, k) / np.sqrt(self.head_dim)
attn_weights = np.exp(scores - scores.max(axis=-1, keepdims=True))
attn_weights /= attn_weights.sum(axis=-1, keepdims=True)
context = attn_weights[..., None] * v
context = context.reshape(batch, -1)
return context @ self.attn_out_weight.T + self.attn_out_bias
def forward(self, x):
residual = x
x = self.layer_norm(x, self.ln1_gamma, self.ln1_beta)
x = residual + self.attention(x)
residual = x
x = self.layer_norm(x, self.ln2_gamma, self.ln2_beta)
mlp_hidden = np.maximum(x @ self.mlp_up_weight.T + self.mlp_up_bias, 0)
x = residual + mlp_hidden @ self.mlp_down_weight.T + self.mlp_down_bias
return x
# ========== 验证 ==========
if __name__ == "__main__":
np.random.seed(42)
hidden_size, num_heads = 8, 2
batch = 2
# 初始化并行模型（2设备）与参考模型
tp_block_0 = TensorParallelTransformerBlock(hidden_size, num_heads, rank=0)
tp_block_1 = TensorParallelTransformerBlock(hidden_size, num_heads, rank=1)
ref_block = ReferenceTransformerBlock(hidden_size, num_heads)
# 将参考权重同步到并行模型（确保一致性）
# 列并行：QKV权重按输出维切分
tp_block_0.qkv_proj.weight = ref_block.qkv_weight[:hidden_size//2]
tp_block_1.qkv_proj.weight = ref_block.qkv_weight[hidden_size//2:]
tp_block_0.qkv_proj.bias = ref_block.qkv_bias[:hidden_size//2]
tp_block_1.qkv_proj.bias = ref_block.qkv_bias[hidden_size//2:]
# 行并行：输出权重按输入维切分
tp_block_0.attn_out_proj.weight = ref_block.attn_out_weight[:, :hidden_size//2]
tp_block_1.attn_out_proj.weight = ref_block.attn_out_weight[:, hidden_size//2:]
tp_block_0.attn_out_proj.bias = ref_block.attn_out_bias
tp_block_1.attn_out_proj.bias = ref_block.attn_out_bias
# MLP同理
tp_block_0.mlp_up.weight = ref_block.mlp_up_weight[:2*hidden_size]
tp_block_1.mlp_up.weight = ref_block.mlp_up_weight[2*hidden_size:]
tp_block_0.mlp_up.bias = ref_block.mlp_up_bias[:2*hidden_size]
tp_block_1.mlp_up.bias = ref_block.mlp_up_bias[2*hidden_size:]
tp_block_0.mlp_down.weight = ref_block.mlp_down_weight[:, :2*hidden_size]
tp_block_1.mlp_down.weight = ref_block.mlp_down_weight[:, 2*hidden_size:]
tp_block_0.mlp_down.bias = ref_block.mlp_down_bias
tp_block_1.mlp_down.bias = ref_block.mlp_down_bias
# 输入
x = np.random.randn(batch, hidden_size)
# 并行前向（模拟2设备）
out_0 = tp_block_0.forward(x)
out_1 = tp_block_1.forward(x)
# 由于All-Reduce已返回完整结果，两个设备输出应一致
parallel_out = (out_0 + out_1) / 2 # 实际每设备都持有完整结果
# 参考前向
ref_out = ref_block.forward(x)
# 验证
diff = np.abs(parallel_out - ref_out).max()
print(f"最大误差: {diff:.6f}")
assert diff < 1e-5, "并行实现与参考实现不一致！"
print("✓ 张量并行Transformer块前向结果与单设备参考一致")运行与验证
- 直接运行脚本：
python tensor_parallel_transformer.py - 观察输出：若打印
最大误差: 0.000000且断言通过，说明2设备张量并行结果与单设备完全一致 - 验证覆盖：注意力QKV列并行、注意力输出行并行、MLP两层并行切分、All-Reduce通信原语均被实际执行
关键说明
- 代码中
all_reduce_sum为模拟实现，真实场景替换为NCCL的all_reduce即可 - LayerNorm参数每设备持有完整副本，不参与切分
- 注意力头数需能被设备数整除（本案例2头/2设备=每设备1头）
- All-Reduce后每设备均获得完整结果，因此两设备输出一致