序列并行：超长序列训练

原始序列: [t0, t1, t2, ..., t_{L-1}]
GPU 0:     [t0, t1, ..., t_{L/N-1}]
GPU 1:     [t_{L/N}, ..., t_{2L/N-1}]
...
GPU N-1:   [t_{L-L/N}, ..., t_{L-1}]

class UlyssesAttention(nn.Module):
    def __init__(self, d_model, n_heads, n_gpus):
        super().__init__()
        self.n_gpus = n_gpus
        self.local_heads = n_heads // n_gpus
        # 每个GPU只存储部分注意力头的参数
        self.W_q = nn.Linear(d_model, self.local_heads * (d_model // n_heads))
        self.W_k = nn.Linear(d_model, self.local_heads * (d_model // n_heads))
        self.W_v = nn.Linear(d_model, self.local_heads * (d_model // n_heads))
        self.W_o = nn.Linear(self.local_heads * (d_model // n_heads), d_model)
    
    def forward(self, x):
        # x的形状: [batch, local_seq_len, d_model]
        # 每个GPU持有序列的一部分
        
        # 1. 本地计算QKV
        Q_local = self.W_q(x)  # [batch, local_seq, local_heads * head_dim]
        K_local = self.W_k(x)
        V_local = self.W_v(x)
        
        # 2. All-to-All通信：重新排列序列和头维度
        # 从 [batch, local_seq, local_heads, head_dim]
        # 到 [batch, global_seq, local_heads, head_dim]
        Q_global = all_to_all(Q_local, split_dim=1, concat_dim=2)
        K_global = all_to_all(K_local, split_dim=1, concat_dim=2)
        V_global = all_to_all(V_local, split_dim=1, concat_dim=2)
        
        # 3. 本地注意力计算（现在每个GPU有完整的序列，但只有部分头）
        attn_out = flash_attention(Q_global, K_global, V_global)
        
        # 4. All-to-All通信：将结果按序列切分回去
        attn_local = all_to_all(attn_out, split_dim=2, concat_dim=1)
        
        # 5. 输出投影
        output = self.W_o(attn_local)
        return output

class MegatronSP(nn.Module):
    def forward(self, x):
        # x按序列维度切分：[batch, local_seq, d_model]
        
        # 1. LayerNorm：本地独立计算（无需通信！）
        x_norm = layer_norm(x)  # 逐token操作，天然支持序列并行
        
        # 2. QKV投影：本地计算
        Q = self.W_q(x_norm)
        K = self.W_k(x_norm)
        V = self.W_v(x_norm)
        
        # 3. 注意力：需要All-Gather恢复完整序列
        Q_full = all_gather(Q, dim=1)  # [batch, seq, local_heads * head_dim]
        K_full = all_gather(K, dim=1)
        V_full = all_gather(V, dim=1)
        
        # 4. 注意力计算
        attn_out = attention(Q_full, K_full, V_full)
        
        # 5. Reduce-Scatter：按序列切分回去
        attn_local = reduce_scatter(attn_out, dim=1)
        
        # 6. 输出投影
        output = self.W_o(attn_local)
        
        # 7. Dropout：本地独立计算
        output = dropout(output)
        
        return output + x  # 残差连接

# 组合配置
parallel_config = {
    "tensor_model_parallel_size": 4,    # TP=4
    "pipeline_model_parallel_size": 2,  # PP=2
    "sequence_parallel_size": 4,        # SP=4 (Megatron方式)
    "ulysses_parallel_size": 2,         # UP=2 (Ulysses方式)
}
# 总GPU数 = 4 × 2 × 4 × 2 = 64

def dynamic_sp(model, input_ids, seq_parallel_threshold=8192):
    """动态序列并行"""
    seq_len = input_ids.shape[1]
    
    if seq_len > seq_parallel_threshold:
        # 长序列：启用序列并行
        return model.forward_sp(input_ids)
    else:
        # 短序列：标准计算（避免通信开销）
        return model.forward(input_ids)