长序列训练的挑战

Ulysses Sequence Parallelism（来自 Snowflake AI Research 的 Arctic Long Sequence Training (ALST) 协议的一部分）通过跨多个 GPU 分布式注意力计算，利用注意力头并行化的方式提供了一个优雅的解决方案。在本文中，我们将探讨 Ulysses 如何工作，以及它如何在 Hugging Face 生态系统中集成——从 Accelerate 到 Transformers Trainer 和 TRL 的 SFTTrainer。

长序列训练的挑战

与 Accelerate 的集成

与 Transformers Trainer 的集成

与 TRL SFTTrainer 的集成

Ulysses 与 Ring Attention 的对比

长序列训练的挑战

Transformer 中的注意力机制随着序列长度呈二次方增长。对于长度为 n 的序列，标准注意力需要 O(n²) 的 FLOPs 和 O(n²) 的内存来计算和存储注意力分数矩阵。FlashAttention 等优化实现通过分块计算和避免物化完整注意力矩阵，将内存降低到 O(n)——但 O(n²) 的计算量仍然存在。对于极长序列（32k+ tokens），即使使用 FlashAttention，训练仍然会推至单 GPU 内存的极限。

考虑这些长上下文训练至关重要的场景：

• 文档理解：处理整本书籍、法律文件或研究论文
• 代码分析：理解具有多个相互关联文件的大型代码库
• 推理任务：模型"逐步思考"可能在推理过程中生成数千个 tokens
• 检索增强生成：将多个检索到的段落合并到上下文中

传统的数据并行化在这里没有帮助——每个 GPU 仍然需要在注意力块内处理完整序列。我们需要一种方法将序列本身分割到多个设备上。

Ulysses Sequence Parallelism (SP)（在 DeepSpeed Ulysses 论文中介绍）采取了一种聪明的方法：除了在序列维度上进行分割，它还将注意力头分区分散到各个 GPU 上。

序列分片：输入序列沿序列维度分割到 P 个 GPU 上。每个 GPU i 持有 tokens [i·n/P,(i+1)·n/P)。

序列分片：输入序列沿序列维度分割到 P 个 GPU 上。每个 GPU i 持有 tokens [i·n/P,(i+1)·n/P)。

QKV 投影：每个 GPU 计算其本地序列块的查询、键和值投影。

QKV 投影：每个 GPU 计算其本地序列块的查询、键和值投影。

全到全通信：一个全到全集合操作重新分布数据，使每个 GPU 在投影后持有所有序列位置，但仅针对注意力头的一个子集。

全到全通信：一个全到全集合操作重新分布数据，使每个 GPU 在投影后持有所有序列位置，但仅针对注意力头的一个子集。

本地注意力：每个 GPU 使用标准注意力机制（FlashAttention 或 SDPA）计算其分配的头部的注意力。

本地注意力：每个 GPU 使用标准注意力机制（FlashAttention 或 SDPA）计算其分配的头部的注意力。

全到全通信：另一个全到全操作反转重新分布，返回到序列分片格式。

全到全通信：另一个全到全操作反转重新分布，返回到序列分片格式。

输出投影：每个 GPU 计算其本地序列块的输出投影。

输出投影：每个 GPU 计算其本地序列块的输出投影。

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

通信复杂度

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

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

Ring Attention 每个 GPU 通信 O(n·d)——多出 P 倍——通过在环周围进行 P-1 次顺序的点对点传输。Ulysses 还受益于更低的延迟，因为全到全可以在单个集合步骤中利用完整的对分带宽，而 Ring Attention 则在 P-1 跳上序列化。

与 Accelerate 的集成

Accelerate 通过其 ParallelismConfig 类和 DeepSpeed 集成为 Ulysses 序列并行化提供了基础。

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

parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,  # Split across 4 GPUs
    dp_shard_size=1,  # Must satisfy: dp_replicate × dp_shard × sp_size = num_processes
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length=None,  # None for variable-length sequences
        sp_seq_length_is_variable=True,
        sp_attn_implementation="flash_attention_2",  # or "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)

# This registers the model with Ulysses and wraps the dataloader
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

• 使用 DeepSpeed 的 UlyssesSPAttentionHF 注册模型
• 使用 UlyssesSPDataLoaderAdapter 包装 dataloader 以处理序列分片
• 自动注入 shift_labels 以确保正确的损失计算

使用 Ulysses，每个 GPU 计算序列不同部分的损失。必须正确聚合损失，按每个秩的有效 tokens 数加权。如果你使用 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()

# Gather losses and token counts from all SP ranks
    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)

# Weighted aggregation
    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)

加权损失聚合确保在 tokens 不均匀分布在各秩上时获得正确的梯度（例如，当某些秩仅包含填充或屏蔽的提示 tokens 时）。

Ulysses 和 Ring Attention 都在训练期间使用 position_ids 而不是 attention_mask 来进行因果掩蔽。在这些序列长度下的 4D 注意力掩蔽会和注意力分数本身一样有问题——在 128k tokens 时，这又是另一个 ~1TB 的张量。Position IDs 以 O(n) 的内存实现相同的因果行为，而不是 O(n²)。在评估/推理期间，DeepSpeed 的 SP 注意力层可以完全绕过 SP 操作（通过 disable_in_eval），回退到模型的默认注意力实现。

与 Transformers Trainer 的集成

Transformers Trainer 通过 TrainingArguments.parallelism_config 提供了无缝的 Ulysses 集成。它自动处理所有 SP 特定的细节——dataloader 包装、序列分片和损失聚合——所以你无需编写上面显示的任何自定义损失代码。

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

from transformers import TrainingArguments

training_args = TrainingArguments(
    parallelism_config=parallelism_config,  # same ParallelismConfig as above
    per_device_train_batch_size=1,
)

Trainer 自动处理的内容

• Dataloader 包装：在模型准备后，Trainer 使用 UlyssesSPDataLoaderAdapter 包装 dataloader
• 损失计算：compute_loss 方法检测 SP 模式并路由到专门的 _deepspeed_sp_compute_loss，该方法处理：在 SP 秩上收集损失、计算每个秩的有效 token 计数、加权损失聚合

跨 SP 排序的损失值汇聚

计算每个排序的有效 token 计数

加权损失聚合

批大小计算：有效数据并行世界大小计入 SP

dp_world_size = world_size // sp_size

数据加载器长度调整：训练步骤计算已针对 SP 对迭代计数的影响进行调整

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

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

TRL SFTTrainer 集成

TRL 的 SFTTrainer 构建于 Transformers Trainer 之上，并为长序列的监督微调添加了特定优化。

from trl import SFTConfig, SFTTrainer
from accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig

parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=2,
    dp_shard_size=2,  # 2D 并行：SP × DP = 4 GPUs
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length_is_variable=True,
        sp_attn_implementation="flash_attention_2",
    ),
)

training_args = SFTConfig(
    ...,
    parallelism_config=parallelism_config,
    max_length=32768,
    pad_to_multiple_of=2,  # 必须等于 sp_size
    per_device_train_batch_size=1,
)

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
)
trainer.train()

Ulysses 的关键 SFTConfig 参数

Accelerate 配置文件

创建 alst_ulysses_4gpu.yaml：

compute_environment: LOCAL_MACHINE
distributed_type: DEEPSPEED
mixed_precision: bf16
num_processes: 4
deepspeed_config:
  zero_stage: 3
  seq_parallel_communication_data_type: bf16
parallelism_config:
  parallelism_config_sp_size: 2
  parallelism_config_sp_backend: deepspeed
  parallelism_config_dp_shard_size: 2
  parallelism_config_sp_seq_length_is_variable: true
  parallelism_config_sp_attn_implementation: flash_attention_2

完整训练命令

accelerate launch --config_file alst_ulysses_4gpu.yaml \
    trl/scripts/sft.py \
    --model_name_or_path meta-llama/Llama-3.1-8B \
    --dataset_name trl-lib/Capybara \
    --max_length 32768 \
    --packing \
    --pad_to_multiple_of 2 \
    --per_device_train_batch_size 1

Shift Labels 处理

当启用 Ulysses 时，SFTTrainer 会自动处理预移位的 labels：

# 使用 SP 时，数据加载器适配器会预移位 labels
# 训练器检测到这一点并直接使用 shift_labels
labels = inputs["labels"] if "shift_labels" not in inputs else None

# 损失计算使用预移位的 labels
if "shift_labels" in inputs:
    shift_logits = outputs.logits.contiguous()
    shift_labels = inputs["shift_labels"]
else:
    shift_logits = outputs.logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()

比较 Ulysses 和 Ring Attention

Ulysses 和 Ring Attention 都支持长上下文训练，但特性不同：

何时选择 Ulysses vs Ring Attention

由于在两者之间切换只需改变 accelerate 配置，我们建议在你的具体设置上尝试两者并比较性能和内存使用情况。主要限制是 Ulysses 要求 num_heads >= sp_size，而 Ring Attention 没有这样的限制。

最佳实践

1. 序列长度可整除性

始终确保序列长度能被 sp_size 整除：

training_args = SFTConfig(
    pad_to_multiple_of=4,  # 当 sp_size=4 时
    max_length=32768,  # 必须能被 4 整除
)

2. 使用 Flash Attention

Flash Attention 2 相比 SDPA 提供了更清晰的输出和更好的性能：

parallelism_config = ParallelismConfig(
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_attn_implementation="flash_attention_2",
    ),
)

对 Hopper 使用 Flash Attention 3，并留意 Flash Attention 4 针对 Blackwell 的发布（FA2 在 Blackwell 上速度很慢）。

3. 与 DeepSpeed ZeRO 结合

对于超大模型，将 Ulysses 与 ZeRO Stage 3 结合：

deepspeed_config:
  zero_stage: 3
  offload_optimizer:
    device: cpu

如果模型非常庞大，也可以卸载参数，通过在上面添加以下内容：

  offload_param:
    device: cpu

5. 使用内存碎片友好的 PyTorch 分配器

此环境变量将允许更长的序列长度：

export PYTORCH_ALLOC_CONF=expandable_segments:True

6. 2D 并行配置

为 GPU 数量平衡 SP 和 DP：

记住：dp_replicate_size × dp_shard_size × sp_size = num_processes

如果你期望的模型架构被 Liger-Kernel 支持，它与 Ulysses SP 完全兼容，可通过单一标志启用：

training_args = SFTConfig(
    use_liger_kernel=True,
)

主要内存节省来自 FusedLinearCrossEntropy，它避免了在损失计算过程中实现完整的 logits 张量。节省量随着更长序列而增长，此时 logits 张量更大。

此外，你可以启用 TiledMLP 进一步延长序列长度——像 FusedLinearCrossEntropy 一样，它通过平铺大型矩阵操作来节省工作内存。

8. 跨排序的 Token 分布

无需手动平衡跨 SP 排序的 token——损失聚合代码可优雅地处理不均匀分布（包括有效 token 为零的排序）。通过在合理大小的数据集上随机批处理，分布在训练过程中统计上趋于均匀。

基准测试结果

为量化 Ulysses SP 的优势，我们在 Gutenberg 英文流式数据集上使用 TRL 的 SFTTrainer 训练了 Qwen3-4B。所有实验在 H100 80GB GPU 上运行，配备 DeepSpeed ZeRO-3、CPU 优化器卸载、梯度检查点和 flash-attn2 作为注意力后端。

上表中的基准运行采用相同的全局批大小（8 微批）、cosine 学习率调度和种子，因此这些基准损失曲线直接可比。

损失曲线匹配诊断（4 GPU）

为验证 SP 与 DP 的损失等价性，我们进行了受控的 4-GPU A/B 实验，使用相同的种子、模型、优化器、学习率调度和数据顺序。

DP=4, SP=1, GAS=1（基线）

DP=1, SP=4, GAS=4（Ulysses SP）

为公平比较，GAS 必须与 SP 缩放：

Ulysses SP 将序列分割到 SP 排序，所以每个 SP 排序在每个微步中看到大约 1/SP 的序列 token。

如果 GAS 不变，SP 中的每个优化器步骤聚合的总 token 数比 DP 基线少。

设置 GAS=SP 保证每个优化器步骤的有效 token 匹配：DP tokens/step: dp_world_size * micro_batch * seq_len * GAS = 4 * B * L * 1 SP tokens/step: dp_world_size * micro_batch * (L/SP) * GAS * SP_ranks = 1 * B * (L/4) * 4 * 4 = 4 * B * L

DP tokens/step: dp_world_size * micro_batch * seq_len * GAS = 4 * B * L * 1

SP tokens/step: dp_world_size * micro_batch * (L/SP) * GAS * SP_ranks = 1 * B * (L/4) * 4 * 4 = 4 * B * L

在受控等价性工具中 4 GPU 上的 20 步测量：

要点：在匹配的 token 预算下，SP 和非 SP 在规范化 token 的损失上匹配。剩余差异在于训练器报告的日志（损失），而非底层交叉熵目标。

在 8K token 处，DP=4 和 SP=4 每 GPU 使用几乎相同的内存（约 22 GB，配备 ZeRO-3）。SP 的优势在于它可以扩展到更长的序列：在 96K token 处（长 12 倍），峰值内存为 66 GB——仍在 H100 的 80 GB 容量内。在 128K，模型内存溢出，确定了这个配置的实际限制。不带 SP 的 DP=4 无法超过 8K。

在相同序列长度（8K）处，SP=4 与单 GPU 基线的吞吐量相当——all-to-all 通信开销在 NVLink 连接的 GPU 上最小。真正的优势来自更长的序列：随着序列长度增长，二次注意力计算主导通信和其他开销，使每个训练步骤越来越高效。每一步也处理更多成比例的 token，所以吞吐量随序列长度缩放。在 64K，SP=4 处理 13,396 token/秒——比基线快 3.7 倍。

这些结果仅使用 4 GPU 配备 SP=4。用 8 GPU（SP=8），你可以推进到更长的序列——高达 256K+ token——或使用 2D 并行（SP=4, DP=2）来结合长上下文训练与数据并行吞吐量。

依赖要求

HF Accelerate: deepspeed>=0.18.1 accelerate>=1.12

HF Trainer: deepspeed>=0.18.1 accelerate>=1.12 transformers>=5.0

HF TRL: deepspeed>=0.18.1 accelerate>=1.12 transformers>=5.0 trl>=0.18.0

对 Ampere GPU 使用 flash_attention_2，或对 Hopper GPU 使用 flash_attention_3。等待 Blackwell 上的 flash_attention_4 🕰。

参考资源

Accelerate: Context Parallelism Guide

TRL: Distributing Training

DeepSpeed Sequence Parallelism

Accelerate ALST 示例

TRL Accelerate 配置

Arctic 长序列训练：面向数百万 Token 序列的可扩展、高效训练

DeepSpeed Ulysses：支持超长序列 Transformer 模型训练的系统优化

Accelerate ND-Parallel：高效多 GPU 训练指南

理解 Ulysses 与 Ring Attention

在 Axolotl 中使用序列并行实现长上下文训练

本文提及的模型 1

本文提及的数据集 1

本文提及的论文 3

博客中的更多文章

超越 LoRA：你能击败最流行的微调技术吗？

你 🫵 可以在 transformers 中使用的 OpenAI gpt-oss 技巧

· 注册或登录后发表评论

本文提及的模型 1

本文提及的数据集 1

本文提及的论文 3