图 1 序列并行实现方案
图 1 展示了序列并行策略的实现方案。由于 Transformer 结构较为规整，除 attention 计算外，其他计算过程中 token 之间不会互相影响（即每个 token 的计算是独立的），这一条件为序列并行提供了有利条件。上图展示了序列并行的核心设计。设由 P 个 GPUs 共同计算一个长度为 N 的长序列，在 Attention 计算的第一阶段，长度为 N / P 的子序列会通过线性层投影为 Query、Key、Value。接下来， QKV Tensor 会在参与序列并行计算的多个 GPUs 之间通过高度优化的 all-to-all 通信算子汇聚，得到序列长度为 N ，但更少注意力头的子序列。注意力计算后，通过另一个 all-to-all 通信算子将其转换为长度为 N / P 的子序列，进行后续计算。伪代码如下所示。 .. code-block:: python # Pseudo code for an Attention Layer # Input: hidden_states with shape (bs, seq_len, dim) # Output: attn_out with shape (bs, seq_len, dim) def attn_forward(hidden_states): q, k, v = qkv_proj(hidden_states) q, k, v = reshape(q, k, v) # (bs, q_len, dim) -> (bs, q_len, nhead, hdim) q, k = apply_rotary_pos_emb(q, k, cos, sin) sp_size = get_sequence_parallel_world_size() # (bs, q_len, nhead, hdim) -> (bs, q_len * sp_size, nhead / sp_size, hdim) q, k, v = all_to_all(q, k, v, sp_size) attn_out = local_attn(q, k, v) # (bs, q_len * sp_size, nhead / sp_size, hdim) -> (bs, q_len, nhead, hdim) attn_out = all_to_all(attn_out) attn_out = reshape(attn_out) # (bs, q_len, nhead, hdim) -> (bs, q_len, dim) attn_out = o_proj(attn_out) return attn_out 序列并行 API ============= 为了方便在其他 repo 中使用序列并行策略，XTuner 中抽象出了序列并行所必须的五个 API 接口： - 序列并行分布式环境初始化 (init_sequence_parallel) - 适配序列并行的 Data Sampler (SequenceParallelSampler) - 数据 Pad 与切分 (pad_for_sequence_parallel, split_for_sequence_parallel) - 适配序列并行的 Attention (dispatch_modules) - reduce loss 以正确打印训练损失 (reduce_sequence_parallel_loss) 分布式环境初始化 ------------------- 由于序列并行算法会将长序列切分为 `sequence_parallel_world_size` 块，并将每个子序列分发给对应的 GPU 独立进行计算。因此需要在训练开始前初始化序列并行分布式环境，以指定哪几块 GPU 共同负责一个长序列输入的计算。 一个 `sequence_parallel_world_size = 4` 的示例如下： .. code-block:: python # We have to initialize the distributed training environment first. # Here is an example when training on slurm scheduler # from xtuner.parallel.sequence import init_dist # init_dist('slurm', 'nccl', init_backend='deepspeed') from xtuner.parallel.sequence import init_sequence_parallel sequence_parallel_world_size = 4 init_sequence_parallel(sequence_parallel_world_size) .. tip:: 上述过程在 ``xtuner/engine/_strategy/deepspeed.py`` 中实现。 Data Sampler -------------- 在使用序列并行后，Dataloader 的采样策略需要进一步调整。例如当 `sequence_parallel_world_size = 4` 时，4 块 GPU 从 Dataloader 拿到的数据需要是完全一样的。 在构建 Dataloader 时搭配 XTuner 中提供的 `SequenceParallelSampler` 使用即可： .. code-block:: python from xtuner.parallel.sequence import SequenceParallelSampler dataloader = DataLoader( train_dataset, sampler=SequenceParallelSampler(train_dataset), **other_dataloader_params) 数据 Pad 与切分 --------------- 由于每条训练数据的长度可能不尽相同，我们需要将数据进行 Pad 以使得序列长度可以被 `sequence_parallel_world_size` 整除，这样一条长数据才能被均等地分发给不同的 GPU 上。 训练过程中需要被 Pad 的 Tensor 往往有 input_ids, labels, position_ids, attention_mask 四个，pad 的过程可以通过以下方式实现： .. code-block:: python from xtuner.parallel.sequence import pad_for_sequence_parallel input_ids, labels, position_ids, attention_mask = pad_for_sequence_parallel( input_ids, labels, position_ids, attention_mask) 如果训练过程用不到 attention_mask，那么可以： .. code-block:: python input_ids, labels, position_ids, _ = pad_for_sequence_parallel( input_ids, labels, position_ids) Pad 后，我们需要对长序列均等切分： .. code-block:: python from xtuner.parallel.sequence import split_for_sequence_parallel # attention mask should not be split input_ids, labels, position_ids = split_for_sequence_parallel( input_ids, labels, position_ids) .. tip:: 以上两步在 ``xtuner/dataset/collate_fns/default_collate_fn.py`` 中实现。 Attention ----------- 在 Attention 的计算过程中，序列中的不同 token 是不能独立运算的，但不同的 attention head 之间的计算却是独立的。因此，如第一节所述，需要在计算 Attention 前后（即 qkv_proj 后和 o_proj 前）分别插入一个 all-to-all 操作。 XTuner 提供了 dispatch_modules 接口以支持修改模型 Attention 的计算方式： .. code-block:: python from xtuner.model.modules import dispatch_modules model: AutoModelForCausalLM dispatch_modules(model) .. tip:: 上述过程在 ``xtuner/model/sft.py`` 中实现。 Reduce Loss ------------- 这个 API 对于保证训练的正确性不是必须的，但对于观测模型训练状态，打印训练 loss 是非常有用的。 .. code-block:: python from xtuner.parallel.sequence import reduce_sequence_parallel_loss outputs = llm(input_ids=input_ids, labels=labels, **kwargs) num_tokens_per_rank = (labels != -100).sum() # Suppose sequence parallel world size equals to 4, # losses on rank0, rank1, rank2, rank3 are different. loss = reduce_sequence_parallel_loss(outputs.loss, num_tokens_per_rank) # After loss reduction, losses on rank0, rank1, rank2, rank3 are the same. .. tip:: 上述过程在 ``xtuner/model/sft.py`` 中实现。