Ulysses 序列并行：百万令牌上下文训练

来源：https://huggingface.co/blog/ulysses-sp 返回文章 (https://huggingface.co/blog)

Kashif Rasul的头像 (https://huggingface.co/kashif)

Stas Bekman的头像 (https://huggingface.co/stas)

• 内容 (https://huggingface.co/blog/ulysses-sp#contents)
• 长序列训练的挑战 (https://huggingface.co/blog/ulysses-sp#the-challenge-of-long-sequence-training)
• Ulysses 的工作原理 (https://huggingface.co/blog/ulysses-sp#how-ulysses-works)- 通信复杂度 (https://huggingface.co/blog/ulysses-sp#communication-complexity)
• 与 Accelerate 的集成 (https://huggingface.co/blog/ulysses-sp#integration-with-accelerate)- 配置 (https://huggingface.co/blog/ulysses-sp#configuration) - 关键参数 (https://huggingface.co/blog/ulysses-sp#key-parameters) - 使用 Accelerator (https://huggingface.co/blog/ulysses-sp#using-the-accelerator) - 损失聚合 (https://huggingface.co/blog/ulysses-sp#loss-aggregation)
• 与 Transformers Trainer 的集成 (https://huggingface.co/blog/ulysses-sp#integration-with-transformers-trainer)- 配置 (https://huggingface.co/blog/ulysses-sp#configuration-1) - Trainer 自动处理的内容 (https://huggingface.co/blog/ulysses-sp#what-the-trainer-handles-automatically) - 启动命令 (https://huggingface.co/blog/ulysses-sp#launch-command)
• 与 TRL SFTTrainer 的集成 (https://huggingface.co/blog/ulysses-sp#integration-with-trl-sfttrainer)- 配置 (https://huggingface.co/blog/ulysses-sp#configuration-2) - Ulysses 的关键 SFTConfig 参数 (https://huggingface.co/blog/ulysses-sp#key-sftconfig-parameters-for-ulysses) - Accelerate 配置文件 (https://huggingface.co/blog/ulysses-sp#accelerate-config-file) - 完整训练命令 (https://huggingface.co/blog/ulysses-sp#complete-training-command) - Shift Labels 处理 (https://huggingface.co/blog/ulysses-sp#shift-labels-handling)
• 比较 Ulysses 和 Ring Attention (https://huggingface.co/blog/ulysses-sp#comparing-ulysses-and-ring-attention)- 何时选择 Ulysses vs Ring Attention (https://huggingface.co/blog/ulysses-sp#when-to-choose-ulysses-vs-ring-attention)
• 最佳实践 (https://huggingface.co/blog/ulysses-sp#best-practices)- 1. 序列长度整除性 (https://huggingface.co/blog/ulysses-sp#1-sequence-length-divisibility) - 2. 使用 Flash Attention (https://huggingface.co/blog/ulysses-sp#2-use-flash-attention) - 3. 与 DeepSpeed ZeRO 结合 (https://huggingface.co/blog/ulysses-sp#3-combine-with-deepspeed-zero) - 5. 使用内存碎片友好的 PyTorch 分配器 (https://huggingface.co/blog/ulysses-sp#5-use-memory-fragmentation-friendly-pytorch-allocator) - 6. 二维并行配置 (https://huggingface.co/blog/ulysses-sp#6-2d-parallelism-configuration) - 7. Liger-Kernel (https://huggingface.co/blog/ulysses-sp#7-liger-kernel) - 8. 跨秩的令牌分布 (https://huggingface.co/blog/ulysses-sp#8-token-distribution-across-ranks)
• 基准测试 (https://huggingface.co/blog/ulysses-sp#benchmarks)- 设置 (https://huggingface.co/blog/ulysses-sp#setup) - 损失曲线匹配诊断 (4 GPU) (https://huggingface.co/blog/ulysses-sp#loss-curve-matching-diagnostics-4-gpu) - 内存降低 (https://huggingface.co/blog/ulysses-sp#memory-reduction) - 吞吐量 (https://huggingface.co/blog/ulysses-sp#throughput)
• 需求 (https://huggingface.co/blog/ulysses-sp#requirements)
• 资源 (https://huggingface.co/blog/ulysses-sp#resources)- 文档 (https://huggingface.co/blog/ulysses-sp#documentation) - 示例 (https://huggingface.co/blog/ulysses-sp#examples) - 论文 (https://huggingface.co/blog/ulysses-sp#papers) - 相关博客文章 (https://huggingface.co/blog/ulysses-sp#related-blog-posts)

在长序列上训练大型语言模型已成为构建先进 AI 系统的必需条件。随着模型越来越多地用于文档分析、代码理解、复杂推理和 RAG 工作负载等任务，处理数十万甚至数百万令牌序列的需求急剧增长。为了直观感受这一规模，一本平均书籍大约包含 25 万个令牌，因此在多文档上下文或书籍长度的输入上进行训练需要处理超出单个 GPU 容纳范围的序列。然而，用这样的长上下文进行训练带来了重大的内存挑战：注意力计算与序列长度呈二次方关系，对于超过数万令牌的上下文，很快就会超出 GPU 内存。

Ulysses 序列并行（来自 Snowflake AI Research 的 Arctic 长序列训练 (ALST) 协议 (https://huggingface.co/papers/2506.13996) 的一部分）通过通过注意力头并行来在多个 GPU 之间分布注意力计算，提供了一个优雅的解决方案。在本文中，我们将探讨 Ulysses 如何工作，以及它如何在整个 Hugging Face 生态系统中进行集成——从 Accelerate 到 Transformers Trainer 和 TRL 的 SFTTrainer。

内容

• 长序列训练的挑战 (https://huggingface.co/blog/ulysses-sp#the-challenge-of-long-sequence-training)
• Ulysses 的工作原理 (https://huggingface.co/blog/ulysses-sp#how-ulysses-works)
• 与 Accelerate 的集成 (https://huggingface.co/blog/ulysses-sp#integration-with-accelerate)
• 与 Transformers Trainer 的集成 (https://huggingface.co/blog/ulysses-sp#integration-with-transformers-trainer)
• 与 TRL 的 SFTTrainer 的集成 (https://huggingface.co/blog/ulysses-sp#integration-with-trl-sfttrainer)
• 比较 Ulysses 和 Ring Attention (https://huggingface.co/blog/ulysses-sp#comparing-ulysses-and-ring-attention)
• 最佳实践 (https://huggingface.co/blog/ulysses-sp#best-practices)
• 基准测试 (https://huggingface.co/blog/ulysses-sp#benchmarks)
• 资源 (https://huggingface.co/blog/ulysses-sp#resources)

长序列训练的挑战

Transformer 中的注意力机制与序列长度呈二次方关系进行扩展。对于长度为 n 的序列，标准注意力需要 O(n²) FLOPs 和 O(n²) 内存来计算和存储注意力分数矩阵。优化的实现（如 FlashAttention (https://huggingface.co/papers/2205.14135)）通过平铺计算和从不物化完整注意力矩阵，将内存减少到 O(n)——但 O(n²) 的计算仍然存在。对于非常长的序列（32k+ 令牌），即使使用 FlashAttention，训练仍然接近单 GPU 内存的极限。

考虑以下需要长上下文训练的场景：

• 文档理解：处理整本书籍、法律文件或研究论文
• 代码分析：理解具有多个相互连接文件的大型代码库
• 推理任务：模型在推理过程中可能生成数千个令牌
• 检索增强生成：将许多检索到的段落纳入上下文

传统的数据并行无法帮助解决这个问题——每个 GPU 仍然需要在注意力块内处理完整序列。我们需要一种方法来在多个设备之间分割序列本身。

Ulysses 的工作原理

Ulysses 序列并行（SP）（在 DeepSpeed Ulysses 论文 (https://huggingface.co/papers/2309.14509) 中引入）采取了一种聪明的方法：除了在序列维度上进行分割外，它还在 GPU 之间对注意力头进行分区。

Ulysses 序列并行概述 Ulysses 沿着序列维度在 P 个 GPU 之间分割输入序列，并使用 all-to-all 通信来交换键值对，使每个 GPU 能够计算注意力头的子集。（来源：Snowflake 工程博客 (https://www.snowflake.com/en/engineering-blog/ulysses-low-latency-llm-inference/)）

工作原理如下：

1. 序列分片：输入序列沿着序列维度在 P 个 GPU 之间进行分割。每个 GPU i 持有令牌 [i·n/P, (i+1)·n/P)。
2. QKV 投影：每个 GPU 为其本地序列块计算查询、键和值投影。
3. All-to-All 通信：一个 all-to-all 集体操作重新分配数据，使得每个 GPU 在投影后持有所有序列位置，但仅用于注意力头的子集。
4. 本地注意力：每个 GPU 使用标准注意力机制（FlashAttention 或 SDPA）为其分配的头计算注意力。
5. All-to-All 通信：另一个 all-to-all 操作反转重新分配，返回到序列分片格式。
6. 输出投影：每个 GPU 为其本地序列块计算输出投影。

关键洞察是注意力头是独立的——每个头可以单独计算。通过用头局部性交换序列局部性，Ulysses 以相对较低的通信开销启用了高效并行化。

通信复杂度

Ulysses 每个注意力层需要两个 all-to-all 操作，总通信量为每个 GPU O(n·d/P)，其中：

• n 是序列长度
• d 是隐藏维度
• P 是并行度

Ring Attention 每个 GPU 通信 O(n·d)——高 P 倍——通过 P-1 次顺序点对点转移绕环。Ulysses 也受益于更低的延迟，因为 all-to-all 可以在单个集体步骤中利用全二等分带宽，而 Ring Attention 在 P-1 跳上序列化。

与 Accelerate 的集成

Accelerate 通过其 ParallelismConfig (https://huggingface.co/docs/accelerate/concept_guides/sequence_parallelism#alstulysses-sp-backend-configuration) 类和 DeepSpeed 集成为 Ulysses 序列并行提供了基础。

配置

from accelerate import Accelerator
from accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig

parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,  # 跨 4 个 GPU 分割
    dp_shard_size=1,  # 必须满足：dp_replicate × dp_shard × sp_size = num_processes
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length=None,  # 无以处理可变长度序列
        sp_seq_length_is_variable=True,
        sp_attn_implementation="flash_attention_2",  # 或 "sdpa"
    ),
)

accelerator = Accelerator(parallelism_config=parallelism_config)

关键参数

使用 Accelerator

当调用 accelerator.prepare() 时，Ulysses 会自动设置：

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B")
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)

# 这会使用 Ulysses 注册模型并包装数据加载器
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

prepare() 调用会：

1. 使用 DeepSpeed 的 UlyssesSPAttentionHF 注册模型
2. 使用 UlyssesSPDataLoaderAdapter 包装数据加载器以处理序列分片
3. 自动注入 shift_labels 以进行正确的损失计算

损失聚合

使用 Ulysses，每个 GPU 计算序列不同部分的损失。损失必须正确聚合，并按每个秩的有效令牌数加权。如果使用 Transformers Trainer 或 TRL 的 SFTTrainer，这会自动处理——以下代码仅在编写自定义 Accelerate 训练循环时需要：

sp_size = parallelism_config.sp_size
if sp_size > 1:
    from deepspeed.utils import groups

    sp_group = groups._get_sequence_parallel_group()

    # 从所有 SP 秩收集损失和令牌计数
    losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=sp_group)
    good_tokens = (batch["shift_labels"] != -100).view(-1).sum()
    good_tokens_per_rank = torch.distributed.nn.functional.all_gather(good_tokens, group=sp_group)

    # 加权聚合
    total_loss = sum(
        losses_per_rank[i] * good_tokens_per_rank[i]
        for i in range(sp_size)
        if good_tokens_per_rank[i] > 0
    )
    loss = total_loss / max(sum(good_tokens_per_rank), 1)

accelerator.backward(loss)

加权损失聚合确保在令牌跨秩不均匀分布时（例如，当某些秩仅包含填充或被掩盖的提示令牌时）获得正确的梯度。

Ulysses 和 Ring Attention 在训练期间都使用 position_ids 而不是 attention_mask 来进行因果掩盖。在这些序列长度处的 4D 注意力掩码将同样令人望而生畏，如注意力分数本身——在 128k 令牌处，那是另一个 ~1TB 张量。位置 ID 以 O(n) 内存而不是 O(n²) 实现相同的因果行为。在评估/推理期间，DeepSpeed 的 SP 注意力层可以完全绕过 SP 操作（通过 disable_in_eval）并回退到模型的默认注意力实现。

与 Transformers Trainer 的集成

Transformers Trainer (https://huggingface.co/docs/transformers/main_classes/trainer) 通过 TrainingArguments.parallelism_config (https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments.parallelism_config) 提供了无缝的 Ulysses 集成。它自动处理所有 SP 特定的细节——数据加载器包装、序列分片和损失聚合——因此无需编写上面显示的任何自定义损失代码。

配置

只需将上面的相同 parallelism_config 传递到 TrainingArguments：

from transformers import TrainingArguments

training_args = TrainingArguments(
    parallelism_config=parallelism_config,  # 与上面相同的 ParallelismConfig
    per_device_train_batch_size=1,
)

Trainer 自动处理的内容

1. 数据加载器包装：在模型准备后，Trainer 使用 UlyssesSPDataLoaderAdapter 包装数据加载器
2. 损失计算：compute_loss 方法检测 SP 模式并路由到专门的 _deepspeed_sp_compute_loss，它处理：
   • 在 SP 秩之间收集损失
   • 计算每个秩的有效令牌计数
   • 加权损失聚合
3. 批大小计算：有效数据并行世界大小考虑了 SP：dp_world_size = world_size // sp_size
4. 数据加载器长度调整：训练步骤计算针对 SP 对迭代计数的影响进行了调整

启动命令

使用 accelerate 配置文件或命令行参数：

accelerate launch \
    --config_file deepspeed_ulysses.yaml \
    train.py \
    --per_device_train_batch_size 1

与 TRL SFTTrainer 的集成

TRL 的 SFTTrainer (https://huggingface.co/docs/trl/sft_trainer) 构建在 Transformers Trainer 之上并添加