序列并行

序列并行：解决超长序列的显存瓶颈 张量并行（TP）将模型参数切分到多个GPU，流水线并行（PP）将模型层切分到多个GPU，但两者都没有解决序列长度带来的显存问题。 当序列长度达到128K甚至1M时，即使使用TP和PP，每个GPU上存储的激活值仍然是 O(seq_len × hidden_dim)。序列并行（Sequence Parallelism, SP）正是为了解决这个问题——将序列维度也切分到多个GPU上。 序列并行的核心思想 基本切分 将序列维度切分到N个GPU上，每个GPU只处理 seq_len/N 个token： 原始序列: [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}] 通信需求 序列切分后，LayerNorm和Dropout等逐token操作可以独立进行（无需通信），但注意力计算需要完整的序列信息——每个token的Query需要与所有token的Key/Value计算注意力。 这是序列并行的核心挑战：如何在不重新聚合完整序列的前提下完成注意力计算？ DeepSpeed-Ulysses 原理 DeepSpeed-Ulysses使用All-to-All通信来实现序列并行注意力： 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 通信分析 Ulysses需要2次All-to-All通信（前向）+ 2次All-to-All（反向）。 ...