4. 序列并行

任务目标

在张量并行（Tensor Parallelism）已对注意力与MLP的线性层进行列/行切分的基础上，本案例针对LayerNorm与Dropout这两个逐元素算子，实现序列维度切分（Sequence Parallelism）。核心目标是：让LayerNorm与Dropout的输入激活张量按序列长度切分到各设备，从而避免在每个设备上保留完整的序列激活副本，降低激活内存占用。

核心原理与切分策略

在标准张量并行中，Transformer块的LayerNorm输出需完整复制到每个设备（因为后续QKV投影需要完整序列）。本案例采用序列并行改造：将LayerNorm与Dropout放在序列切分域中执行，而将需要完整序列的注意力/MLP部分保留在张量并行域。具体流程：

• 前一层输出（已序列切分）→ 当前设备上的LayerNorm（仅处理本设备序列分片）→ Dropout → 通过All-to-All通信转换为张量并行所需的完整序列分布 → 注意力/MLP → 再通过All-to-All转回序列切分。
• 本案例聚焦LayerNorm与Dropout的序列切分实现，通信转换部分用简化的索引映射模拟，便于独立验证。

环境依赖

Python 3.8+，仅需NumPy（用于张量运算模拟）。无需真实多卡环境，通过矩阵分块模拟多设备行为。

关键完整代码

import numpy as np

def sequence_parallel_ln_dropout(x, gamma, beta, drop_rate, seq_world_size, rank):
    """
    序列并行版 LayerNorm + Dropout（训练模式）
    x: [seq_len, hidden] 当前rank的完整hidden，但seq已按world切分
    gamma, beta: [hidden]  LayerNorm参数（每设备完整副本）
    seq_world_size: 序列并行设备数
    rank: 当前设备编号
    """
    # 模拟：x已经是本rank的序列分片（seq_len = total_seq / seq_world_size）
    local_seq_len = x.shape[0]
    
    # 1. LayerNorm（逐元素，序列维度独立）
    mean = x.mean(axis=-1, keepdims=True)
    var = x.var(axis=-1, keepdims=True)
    x_norm = (x - mean) / np.sqrt(var + 1e-5)
    x_ln = gamma * x_norm + beta
    
    # 2. Dropout（训练模式，逐元素）
    mask = (np.random.rand(*x_ln.shape) > drop_rate) / (1.0 - drop_rate)
    x_drop = x_ln * mask
    
    # 3. 验证：序列切分正确性（对比未切分结果）
    # 仅当所有rank的x拼接后与完整序列LayerNorm结果一致
    return x_drop, mask

# ========== 验证闭环 ==========
def verify_sequence_parallel():
    np.random.seed(42)
    hidden = 8
    total_seq = 12
    seq_world_size = 3  # 模拟3个设备
    local_seq = total_seq // seq_world_size
    
    # 完整序列输入（用于对比）
    x_full = np.random.randn(total_seq, hidden).astype(np.float32)
    gamma = np.random.randn(hidden).astype(np.float32)
    beta = np.random.randn(hidden).astype(np.float32)
    drop_rate = 0.1
    
    # 未切分参考实现
    mean_full = x_full.mean(axis=-1, keepdims=True)
    var_full = x_full.var(axis=-1, keepdims=True)
    ln_full = (x_full - mean_full) / np.sqrt(var_full + 1e-5)
    ln_full = gamma * ln_full + beta
    
    # 序列并行模拟：每个rank处理自己的分片
    results = []
    for rank in range(seq_world_size):
        x_local = x_full[rank*local_seq:(rank+1)*local_seq].copy()
        x_drop_local, _ = sequence_parallel_ln_dropout(
            x_local, gamma, beta, drop_rate, seq_world_size, rank)
        results.append(x_drop_local)
    
    # 拼接各rank输出（验证LayerNorm部分，Dropout因随机性仅验证形状）
    cat_out = np.concatenate(results, axis=0)
    assert cat_out.shape == (total_seq, hidden), "输出形状错误"
    
    # 验证LayerNorm数值一致性（关闭Dropout后对比）
    ln_results = []
    for rank in range(seq_world_size):
        x_local = x_full[rank*local_seq:(rank+1)*local_seq].copy()
        # 用drop_rate=0.0调用，使Dropout退化为恒等
        x_ln_only, _ = sequence_parallel_ln_dropout(
            x_local, gamma, beta, 0.0, seq_world_size, rank)
        ln_results.append(x_ln_only)
    cat_ln = np.concatenate(ln_results, axis=0)
    
    # 数值一致性检查
    np.testing.assert_allclose(cat_ln, ln_full, rtol=1e-5, atol=1e-6)
    print("✓ 序列并行LayerNorm与完整序列结果一致")
    print(f"✓ 激活内存：每设备保存 {local_seq}/{total_seq} = {local_seq/total_seq:.0%} 序列")
    print(f"  原始完整序列需 {total_seq}x{hidden} 元素，序列并行每设备仅 {local_seq}x{hidden}")

if __name__ == "__main__":
    verify_sequence_parallel()

运行与验证

1. 将代码保存为 seq_parallel_ln_dropout.py。
2. 运行：python seq_parallel_ln_dropout.py
3. 通过条件：程序无异常退出，并输出两条以✓开头的验证信息。

内存收益说明

本实现中，每个设备仅需保存local_seq长度的激活（而非完整total_seq）。当序列长度远大于隐藏维度且设备数增加时，激活内存近似线性减少。上述验证代码通过形状断言与数值对比，确认了序列切分后的LayerNorm结果与未切分完全一致，而Dropout因随机性仅验证输出形状正确。