🤗 accelerate 中的序列并行

本指南将介绍在 🤗accelerate 中使用序列并行（sequence parallelism）的基础知识。

另请参阅非常相关的上下文并行。

为什么需要序列并行？

随着大语言模型以及近期推理模型的兴起，序列长度正在快速增长。加之注意力机制具有二次方级的内存复杂度，这使得我们需要更高效的方式来训练长序列模型。 当序列长度为 128k 时，在采用原生注意力实现的情况下，注意力矩阵的内存需求为 128k * 128k * 2 bytes * num_heads = ~32 GB * num_heads（bf16 精度）。当然，如果使用 flash attention 或 SDPA 这类不会显式生成注意力权重的技术，内存需求会大幅降低，但其增长幅度依然相当可观。

Ulysses 序列并行允许我们沿序列维度对注意力计算的输入进行分片，并以常规方式计算注意力，但每个 GPU 仅使用注意力头的一个切片。借助此方法，配合少量额外工具，我们即可训练长序列模型，将序列长度扩展至 1500 万以上。要了解如何将 Ulysses SP 与 TiledMLP、Liger-Kernel、激活检查点卸载至 CPU 以及其他一些技巧相结合，请参阅论文：Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences。

Ulysses SP 与 FSDP CP 有何不同？

在文档 Context Parallelism 中，您可以了解另一种名为 Context Parallelism 的部署技术。该技术同样在序列维度上进行切分，但使用 Ring Attention 而非在头维度上进行切分。

以下文章对这两种技术之间的差异进行了非常详细的说明：

• https://insujang.github.io/2024-01-11/tensor-parallelism-and-sequence-parallelism-detailed-analysis/
• https://huggingface.co/blog/exploding-gradients/ulysses-ring-attention

以下是改编自其中一篇文章的快速总结：

• Ulysses SP 的通信开销相对较低，但受限于 Attention Heads 的数量，因此对网络拓扑有特定要求（单个副本的 attention heads 数量必须能被参与计算的 GPU 数量整除）。All-to-all 通信对延迟较为敏感，并且需要 Deepspeed。
• FSDP CP Ring-Attention 的 P2P 环形通信没有上述整除性要求，但通信量更大。

最终，应能够按照论文 USP: A Unified Sequence Parallelism Approach for Long Context Generative AI 中所述的方式将 SP + CP 组合起来，以支持更长的序列长度，尽管这尚未集成到 🤗accelerate 中。

支持的序列并行后端

目前唯一的序列并行后端是 deepspeed，它源自 Arctic Long Sequence Training 技术 中经过现代化改造的 Ulysses SP。如果您希望将其直接集成到自己的代码中，也可以参考 教程。

如何使用序列并行？

from accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig

+# Example: 4 GPUs with sp_size=4, dp_shard_size=1
+# Ensure: dp_replicate_size × dp_shard_size × sp_size = 1 × 1 × 4 = 4 GPUs
parallelism_config = ParallelismConfig(
+     sp_backend="deepspeed",
+     sp_size=4,
+     dp_shard_size=1,  # Explicit: no data parallelism
+     sp_handler=DeepSpeedSequenceParallelConfig(
+         sp_seq_length_is_variable: true,
+         sp_attn_implementation="sdpa",
+     ),
+ )

accelerator = Accelerator(
    ...,
    parallelism_config=parallelism_config,
)

与 🤗accelerate 中的任何其他功能一样，你可以通过向 accelerate launch 传递相应的标志来启用序列并行。在这种情况下，操作方式并无不同：

accelerate launch --parallelism-config-sp-size 8  ...

您也可以在 accelerate config 命令中设置 sp_size 和其他配置，它们会被保存在您的 accelerate 配置文件中，因此每次启动脚本时都无需重复传入这些参数。

序列并行可与数据并行结合使用，无需额外的 GPU。 因此，如果您有 8 个 GPU，可以执行：--parallelism-config-dp-shard-size 8 --parallelism-config-sp-size 8。或者，您可以使用 ParallelismConfig 类以编程方式设置它们。

重要：您必须确保 dp_replicate_size × dp_shard_size × sp_size = num_processes。例如，使用 8 个 GPU 且 sp_size=8 时，您需要 dp_shard_size=1（因为 1 × 1 × 8 = 8）。使用 4 个 GPU 且 sp_size=2 时，您可以使用 dp_shard_size=2（因为 1 × 2 × 2 = 4）来实现 2D 并行。

ALST/Ulysses SP 后端配置

ALST/UlyssesSP 通过注意力头并行来实现序列并行，如这篇论文中所述。为简化起见，我们沿用了序列并行的概念和配置，从用户角度来看，二者是相同的：使用多块 GPU 来处理单个批次。

为了让大家了解 ALST 所带来的能力——它使我们能够在单张 H100 GPU 上以 bf16 精度训练 500K tokens，在单个节点上训练 3.7M tokens，并仅使用四个节点就在 Llama-8B 上训练 15M tokens。HF Accelerate 的这一特性仅实现了 ALST 三个组件中的 1 个，因此可实现的序列长度会更小。若要充分发挥 ALST 的全部性能，您需要启用 TiledMLP、将 Activation checkpoint 卸载到 CPU 以及其他若干功能。有关详细信息，请参阅 本教程。

要配置 deepspeed 后端：

# Example: 4 GPUs with sp_size=4, dp_shard_size=1
# Ensure: dp_replicate_size × dp_shard_size × sp_size = 1 × 1 × 4 = 4 GPUs
parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,
    dp_shard_size=1,  # Explicit: no data parallelism
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length=256,
        sp_seq_length_is_variable=True,
        sp_attn_implementation="sdpa",
    ),
)
accelerator = Accelerator(
    ...,
    parallelism_config=parallelism_config,
)

• sp_backend：在此处设置为 deepspeed
• sp_size 是序列并行的程度——在上面的示例中为 4，因此将使用 4 个 GPU 来处理单个批次（同时在相同的 GPU 上执行 DP=4）。
• sp_seq_length 和 sp_seq_length_is_variable 用于处理序列长度。如果 sp_seq_length_is_variable=True，后端将处理可能在批次之间变化的序列长度，此时 sp_seq_length 的值可以设置为任何能被序列并行度整除的数值，或者完全不设置。在这种情况下，每次 forward 时，序列变量都将从输入中推导得出。如果为 False，则 seq_length 需要与批次的序列长度维度相匹配，并且该维度必须填充为始终相同。默认值为 True。
• sp_attn_implementation 是 sdpa、flash_attention_2 或 flash_attention_3 之一。该序列并行实现使用 position_ids 而非 attention_mask，因此，在 eager 支持使用 position_ids 之前，它在此处无法正常工作。此外，请注意 sdpa 无法正确处理合并为单个样本的多个样本；它会将整个样本视为一个整体进行注意力计算。如果样本未合并，sdpa 可以正常工作。因此，Flash Attention 应是理想选择，因为它始终有效。

除了将这些值设置在 DeepSpeedSequenceParallelConfig 对象中，您也可以通过环境变量来实现相同的效果——它们与上述列表末尾的项一一对应。

• PARALLELISM_CONFIG_SP_BACKEND
• PARALLELISM_CONFIG_SP_SEQ_LENGTH
• PARALLELISM_CONFIG_SP_SEQ_LENGTH_IS_VARIABLE
• PARALLELISM_CONFIG_SP_ATTN_IMPLEMENTATION

如果未在代码中传入，sp_size 可以通过 --parallelism_config_sp_size 命令行参数进行设置。其他参数同理。您也可以采用 accelerate 配置文件的方式进行配置，例如针对 2 个 GPU：

distributed_type: DEEPSPEED
deepspeed_config:
  deepspeed_config_file: path/to/ds_config.json
machine_rank: 0
num_machines: 1
num_processes: 2
parallelism_config:
  parallelism_config_dp_replicate_size: 1
  parallelism_config_dp_shard_size: 1  # Must satisfy: 1 × 1 × 2 = 2 num_processes
  parallelism_config_sp_size: 2
  parallelism_config_sp_backend: deepspeed
  parallelism_config_sp_seq_length_is_variable: true
  parallelism_config_sp_attn_implementation: sdpa

如前所述，Ulysses 序列并行通常与数据并行叠加使用——相同的 rank 既用于输入独立的数据流，也用于执行 Ulysses 序列并行。但您也可以创建副本，如下所示：

# Example: 4 GPUs with 2D parallelism (SP=2, DP=2)
# Ensure: dp_replicate_size × dp_shard_size × sp_size = 2 × 1 × 2 = 4 GPUs
parallelism_config = ParallelismConfig(
    dp_replicate_size=2,
    dp_shard_size=1,  # Explicit: no sharding within replicas
    sp_size=2,
    sp_backend="deepspeed",
    sp_handler=DeepSpeedSequenceParallelConfig(...),
)

这里我们使用 4 个 GPU，并配置了 2 个序列并行副本。数据并行由 Deepspeed-ZeRO 驱动。

请注意，UlyssesSPDataLoaderAdapter 内部隐藏了许多机制。它在幕后使用，用于包装您原始的 DataLoader 对象，但如果您遇到任何问题，应当了解它的存在。此外，在批次被分片到各个参与 rank 之前，它还会自动向批次字典中注入正确的 shift_labels。

现在，开始使用 ALST/UlyssesSP 的最后一个步骤是使用可微分的 all_gather 在多个 rank 之间聚合损失，以正确获取梯度。以下代码实现了这一操作，同时排除了被 -100 屏蔽的 token，以获得正确的平均值：

sp_size = parallelism_config.sp_size if parallelism_config is not None else 1
if sp_size > 1:
    sp_group = accelerator.torch_device_mesh["sp"].get_group()
    sp_world_size = parallelism_config.sp_size

# Normal training loop
for iter, batch in enumerate(dl):
    optimizer.zero_grad()

    batch = move_to_device(batch, model.device)

    # The model automatically receives shift_labels via **kwargs and uses it for loss computation.
    # Both standard transformers models and Liger-patched models handle this correctly.
    outputs = model(**batch)
    loss = outputs.loss
    shift_labels = batch["shift_labels"]

    if sp_size > 1:
        # differentiable weighted per-shard-loss aggregation across ranks
        losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=sp_group)
        # special dealing with SFT that has prompt tokens that aren't used in loss computation
        good_tokens = (shift_labels != -100).view(-1).sum()
        good_tokens_per_rank = torch.distributed.nn.functional.all_gather(
            good_tokens, group=sp_group
        )
        # Skip ranks with zero valid tokens to avoid NaN contamination (NaN * 0 = NaN)
        total_loss = sum(
            losses_per_rank[rank] * good_tokens_per_rank[rank]
            for rank in range(sp_world_size)
            if good_tokens_per_rank[rank] > 0
        )
        total_good_tokens = sum(good_tokens_per_rank)
        loss = total_loss / max(total_good_tokens, 1)

    if rank == 0: accelerator.print(f"{iter}: {loss=}")
    accelerator.log(dict(train_loss=loss, step=iter))

    accelerator.backward(loss)
    optimizer.step()

请注意，当批次中存在 shift_labels 时，模型会自动处理它。模型的前向传播通过 **kwargs 接收 shift_labels，并将其传递给损失函数，从而正确计算序列并行下的损失。如果您使用 Liger Kernel，它也能无缝处理 shift_labels，并以非常节省内存的方式计算损失。对于长序列长度，强烈建议使用 Liger，因为它通过融合操作（例如，融合 logit-loss 计算，该计算永远不会在内存中实例化完整的 logits 张量）来释放 GPU 内存。

如果您想了解 HF Accelerate 在幕后做了什么，请阅读此完整集成教程。

有关启用 ALST/UlyssesSP 的 Accelerate 训练循环示例，请参阅 examples/alst_ulysses_sequence_parallelism。

[!警告]

此 API 相当新，仍处于实验阶段。虽然我们致力于提供稳定的 API，但公共 API 的某些小部分可能会在未来发生变化。

由于这是一个 Deepspeed 后端，因此适用常规的 Deepspeed 配置，因此您也可以将序列并行与优化器状态和/或权重卸载相结合，以释放更多 GPU 内存并支持更长的序列长度。该技术已测试可与 DeepSpeed ZeRO stage 2 和 3 配合使用。