SWIFT大模型超长序列训练实操→极速搞定显存直降75%

各位跨境同行，咱们今天来聊一个大模型训练中的“硬骨头”——超长序列并行训练。对于咱们这些身处跨境实战一线的从业者来说，模型的性能，尤其是它处理复杂、超长文本的能力，直接关系到咱们在Agent工作流中应用的可靠性，更是决定了模型在实际业务场景中的“抗压”表现。但长序列训练有个特别让人头疼的地方，那就是随着序列长度的增加，Attention计算的内存消耗会呈现几何级数的增长，这对咱们有限的显卡资源来说，无疑是个巨大的挑战。

在这种情况下，序列并行（Sequence Parallel，简称SP）技术就成了咱们手中的一把利器。简单来说，序列并行就是把一个超长的输入序列，巧妙地拆分成几个子序列，让它们在不同的显卡上并行计算。这样一来，每张卡需要处理的数据量就大大减少了，自然也就降低了单卡的内存需求。

目前主流的序列并行方法有几种，包括Ulysses、Ring-Attention，以及Megatron-SP/CP。在今天的分享中，咱们主要聚焦于Ulysses和Ring-Attention这两种基于Transformers生态的解决方案。它们在实际应用中更灵活，也更容易上手。至于Megatron-CP和Ring-Attention技术原理上有些相似，而Megatron-SP主要是对激活值进行切分，通常与Megatron-TP（张量并行）结合使用，咱们这里就先不展开细说了。

咱们新媒网跨境获悉，先来看看一份实战数据。下图展示了在Qwen2.5-3B模型上，集成Ulysses和Ring-Attention技术后，长序列训练所能达到的内存优化效果：

大家可以看到，当我们将序列拆分成8份时，显存占用从将近80GiB直接降到了不到18GiB，这个效果非常显著！这意味着咱们用普通的商用级显卡，也能跑得动原本需要“高端配置”才能训练的长序列大模型了。当然，凡事有利有弊，由于序列拆分会增加卡间通信量，而且GPU的负载分配也会有所不同，所以训练时间会相应延长。这里也提一句，当拆分成2个子序列时，咱们主要使用了Ulysses；当拆分成4个或8个子序列时，则是Ulysses（2个子序列）结合Ring-Attention（2或4个子序列）的方式。

接下来，咱们就深入浅出地聊聊这两种技术的原理和具体实现方案。

Ulysses：巧借“分头行动”降低内存

Ulysses是外媒DeepSpeed团队开发的一种序列并行算法。它的核心思想，其实用一句话就能概括：把序列切分成几段，每张卡先拿到一小段。在注意力计算前，这些卡之间先进行一轮数据交换，让每张卡都能“看到”完整序列的信息。但请注意，它并不是把整个序列都放在一张卡上，而是巧妙地把不同的注意力头（Attention Heads）分发到不同的卡上进行计算，这样每张卡虽然处理的是完整序列的信息，但因为它只负责一部分注意力头的计算，所以内存占用就大大降低了。计算完成后，再把结果交换回来。

咱们来看个简单的例子，假设把序列分成两段，每段有两个注意力头：

这张图来自Ulysses的论文，更直观地展示了整个流程：

这里面的N代表序列长度，d代表隐藏层维度（hidden_size），你可以把它理解为注意力头数量乘以每个头的实际维度。

经过Ulysses切分后的状态大概是这样的：

Ulysses的技术原理非常清晰，关键点就在于它实现了从N/P到d/P的all-to-all通信。由于在Attention计算时，QKV（Query, Key, Value）是完整的，而不同的注意力头分布在不同的卡上，所以它对GQA（分组查询注意力）和MHA（多头注意力）这类场景都普遍适用。更棒的是，它还能和flash-attn、SDPA等高效的Attention实现，以及padding_free（无填充）等技术完全兼容。

当然，由于存在跨卡通信，反向传播（backward）时需要额外的处理。但Ulysses也有其局限性，那就是它受限于注意力头的数量。尤其是在GQA模型中，KV头的数量远少于Q头，这可能会限制Ulysses在更多卡上进行切分。所以啊，Ulysses虽然好用，但也不是万能的，咱们在实际选择时得根据模型特点来。

Ring-Attention：化整为零的循环接力

在聊Ring-Attention之前，咱们得先简单回顾一下Flash-Attention。这可是个内存优化的大杀器，它的核心思想也很简单：把Attention计算中的QKV和softmax操作，分解成一个个小块（block-wise）进行并行计算和更新，最大限度地利用了SRAM，同时还能显著减少内存占用。

这是Flash-Attention前向传播的伪代码：

大家注意看上面算法1中第10到12行的伪代码，这部分其实是在对分块计算的LSE（log-sum-exp，对数和指数）进行合并更新。LSE可以理解为softmax归一化分母的对数形式，它在Attention计算中扮演着数值稳定性的关键角色。同时，Attention-Out的结果也是通过类似的方式合并更新的。

有了Flash-Attention这个基础，Ring-Attention的思路就水到渠成了。它的想法是：如果每张卡只负责序列中的一部分长度，然后计算结果在卡之间传递，那不就能实现跨卡的Flash-Attention了吗？

所以，Ring-Attention的核心在于：利用Attention计算可以分块进行的原理，将序列块分散到多张卡上分别计算，然后通过循环传递和合并计算结果，最终得到完整的结果。

假设咱们把序列分成了N个块，每张卡处理一个块。Q、K、V这些信息可以在不同块之间进行通信和流动。咱们先从Softmax这部分来看，它涉及到概率归一化。为了避免数值溢出，咱们通常会采用数值稳定的计算方法，这就引出了LSE的更新。

大家可能看着这些公式有点头疼，但别担心，咱们不是来当数学家的，而是要理解它的核心逻辑。对于LSE的更新，可以想象成一个累积的过程：我们已经有了之前累积的结果，现在又处理了一个新的数据块，那么新的总LSE就应该把两者结合起来。通过一系列数学推导（主要是利用对数和指数的性质，以及PyTorch中像softplus这样的操作符来保证数值稳定性），我们就能得到一个递归更新LSE的公式。

同样地，对于Attention-Out（最终的注意力输出）的更新也是如此。我们手头有当前块的LSE和Attention-Out，也有之前累积的LSE和Attention-Out。利用这些信息，咱们同样可以通过递归公式，把它们有效地合并起来，得到最新的总Attention-Out。本质上，这些公式都遵循了Flash-Attention论文中提到的分块更新和在线Softmax的思想。

明白了前向计算的原理，那反向传播（backward）又该如何处理呢？反向传播时，我们需要从最终结果一步步地逆推回去。这里涉及到LSE梯度（lse_grad）的计算，由于篇幅有限，咱们就不详细推导公式了，但大家知道，它也是一套精密的递归计算过程。

好啦，理论准备得差不多了，接下来咱们看看实现。Ring-Attention有很多变体，比如strip-ring-attention。但在这些实现中，"zigzag"（锯齿形）实现方案在负载均衡方面表现最出色。

为什么需要zigzag呢？因为原始的Ring-Attention有个小毛病，那就是负载不均衡。大家看这张图：

如果GPU 0只处理句子的前半部分，而因“因果注意力”（causal=True）的限制，它不能计算句子的后半部分。但GPU 3却可以计算从0到2的完整序列。这样一来，每张卡的计算任务就不一致了，有的卡忙，有的卡闲，整体效率自然就上不去。

为了解决这个问题，像Megatron-CP和一些优秀的实现方案就采用了zigzag分区：

假设咱们要在4张卡上进行分区，并且序列可以均匀地分成8个小块。那么，咱们会把第0块和第7块组合到一起，第1块和第6块组合到一起，以此类推，第2块和第5块、第3块和第4块也分别组合。这种巧妙的组合方式，能确保每张卡的计算负载都尽可能均衡。

这种计算方式还有几个特别之处：

• 当计算本地QKV（比如序列编号为0的块）时，会直接按因果注意力（causal=True）模式计算。
• 当流动的序列编号小于或等于当前卡的编号时，咱们只需要计算KV的前半部分。
• 当流动的序列编号大于当前卡的编号时，咱们只需要计算Q的后半部分。
  

这些精细化的处理，进一步降低了每张卡的计算负担，提高了整体效率。

具体的代码实现，大家可以参考这个链接：

Ulysses与Ring-Attention的融合：强强联合

不难看出，Ulysses和Ring-Attention这两种序列并行方案各有千秋。Ulysses的通信开销相对较低，但受限于注意力头的数量，而且all-to-all通信对网络延迟和拓扑结构有一定要求。Ring-Attention采用P2P（点对点）循环通信，对网络要求相对较低，但通信量较大，而且不受注意力头数量的限制。

既然各有优势，那为啥不把它们结合起来用呢？

从原理上讲，Ulysses和Ring-Attention确实可以协同作战。咱们可以先用Ulysses这种通信开销相对较小的方案来做初步分区。如果遇到注意力头数量不足（比如GQA模型），或者需要分区的序列太多，那就再叠加Ring-Attention来补强。SWIFT框架就实现了这样一套融合计算技术，它能很好地适用于纯文本、多模态、SFT、DPO、GRPO等各种训练场景。

在底层代码实现上，咱们借鉴了一些外媒社区的优秀开源工作，并在此基础上进行了部分重写和优化。

项目地址：

使用起来也非常简单，只需要在命令行中额外添加一个参数：

--sequence_parallel_size N

框架就会自动为你计算分区方案。更厉害的是，它甚至能支持GPU数量不是偶数（比如3、5、7张卡）的场景，这在实战中可是大大方便了咱们的部署，让资源利用更加灵活高效。

分区方法揭秘

最自然的思路是，先用Ulysses进行局部聚合，再用Ring-Attention进行全局计算，从而得出全局的LSE和Attention-Out。

假设咱们要把序列分成4个子序列（其中Ulysses的world_size=2，Ring-Attention的world_size=2），模型的注意力头数为4：

经过Ulysses的all-to-all通信后，GPU0和GPU1属于同一个Ulysses组，它们都持有序列0/3的信息，但分别负责不同的注意力头（前半部分和后半部分）。GPU2和GPU3的情况也类似。

在Ring-Attention计算阶段，GPU0和GPU2会组成一个Ring-Attention组进行循环通信，而GPU1和GPU3则组成另一个组。

值得注意的是，在进行分区之前，咱们需要对序列进行填充（padding），确保它的长度能够被world_size * 2整除，因为zigzag模式需要对子块进行重新组合。

多模态模型的适配：巧解“黑盒”难题

多模态模型的序列分区可不是个轻松活儿，这里面坑不少，主要有两点：

（1）多模态模型的序列长度在实际前向传播之前是无法确定的。有些模型可能会用一个简单的<image>标签来代表图像部分，但经过ViT（视觉转换器）编码后，这个标签会被替换成一个非常长的图像特征序列。

（2）有些模型的输入序列会包含闭合标签，比如<image></image>。如果咱们在这些标签被实际图像编码替换之前就进行分区，很可能直接报错。

通常，多模态大模型包含内层和外层模型。外层负责ViT处理和lm_head计算逻辑，而内层则计算decode_layers，咱们通常称之为骨干网络（backbone）。

为了适配多模态模型的分区，SWIFT在实现时采用了一个非常巧妙的工程技巧：它不是在数据准备阶段（data_collator）就进行分区，而是等到模型骨干网络（backbone）的前向钩子（forward hook）里才做分区。这样做的好处显而易见：首先，这时候图像编码已经和文本部分融合了，序列长度是准确的；其次，这种做法对纯文本模型也完全兼容；最关键的是，它避免了咱们为了适配多模态而修改原始模型代码，大大降低了维护成本，这一点对咱们开发者来说非常重要。

Padding-free（无填充）适配：精打细算每一寸内存

Padding-free（无填充）可以理解为Flash-Attention中一种特殊的输入格式：把多个序列拼接成一个超长的序列，尽可能地减少填充带来的计算浪费。

这种输入方式虽然高效，但给实际工程实现带来了不少挑战。但别急，SWIFT也给出了应对方案：

（1）它会先分解原始的padding-free输入，对每个序列单独进行填充（同样要保证能被world_size * 2整除），然后再分别分区。
（2）在计算Attention之前，SWIFT会根据填充的位置，将Q和V中对应填充的部分设置为0，K中对应填充的部分设置为一个极小值，这样能有效防止填充部分对Attention计算产生负面影响。
（3）这里有个特别需要注意的地方，对于GRPO和DPO这类训练方法，最终的损失计算需要完整的序列。如果在Logits（模型输出的原始分数）阶段就进行聚合（gather），会大幅增加通信量；如果聚合得太晚，又可能导致损失计算异常。所以，SWIFT需要针对每种训练方法的损失计算逻辑进行全面的重写，才能保证结果的准确性。
（4）在反向梯度更新时，由于通信序列号大于当前卡时Q只有半长，所以需要将其恢复到全长。这就要求咱们对每个序列的梯度进行单独填充，并且将LSE填充为极小值。

反向传播：内存与速度的权衡

根据咱们前面公式的推导，LSE和Attention-Out分块更新的反向传播需要按顺序进行，并且依赖一些前向计算的信息，比如每块的LSE、每块的Attention-Out等。虽然这些信息可以在前向的flash_attn_forward中获取并保存在上下文中，但这可能会额外消耗内存。因此，SWIFT选择了一种更为务实的方案：在反向传播时，重新计算一次flash_attn_forward，然后根据中间结果计算LSE的梯度，最后再对QKV进行实际的反向传播。虽然这个方案可能不是速度最优解，但它在内存效率和工程实现上找到了一个很好的平衡点，对咱们资源有限的实战场景来说，是一个非常务实的考量。

内存优化实战成果

咱们再回头看看文章开头那张表，用数据说话最有说服力。咱们用一个3B模型，在8块A100 GPU上进行了内存优化测试：

NPROC_PER_NODE=8 \
swift sft \
--model Qwen/Qwen2.5-3B-Instruct \
--dataset 'test.jsonl' \
# 9000 tokens per sequence
--train_type lora \
--torch_dtype bfloat16 \
--per_device_train_batch_size 4 \
--target_modules all-linear \
--gradient_accumulation_steps 8 \
--save_total_limit 2 \
--save_only_model true \
--save_steps 50 \
--max_length 65536 \
--warmup_ratio 0.05 \
--attn_impl flash_attn \
--sequence_parallel_size 8 \
--logging_steps 1 \
--use_logits_to_keep false \
--padding_free true

可以看到，当拆分成8份时，训练内存占用从接近80GiB骤降到不足20GiB。这意味着，咱们平时用的普通商用级显卡，也能跑得动原本需要“高端配置”才能训练的长序列大模型了，这大大降低了咱们进入长序列训练的门槛，非常有价值！

展望未来：持续探索，共筑辉煌

技术的发展永无止境，虽然咱们在Ulysses + Ring-Attention的融合训练上取得了不错的进展，但未来的优化空间依然广阔。

例如，在反向传播过程中，重新计算flash_attention_forward是否就是速度最优的实现方案呢？P2P通信量和异步执行方面，是不是还有继续优化的潜力？这些都是值得咱们深入思考和探索的方向。

咱们新媒网跨境也期待各位对这方面感兴趣的开发者能够积极参与进来，贡献您的宝贵意见和智慧，共同推动SWIFT在长序列大模型训练能力上的进步，为咱们国家的AI技术发展贡献一份力量！

新媒网（公号: 新媒网跨境发布），是一个专业的跨境电商、游戏、支付、贸易和广告社区平台，为百万跨境人传递最新的海外淘金精准资讯情报。

本文来源：新媒网 https://nmedialink.com/posts/swift-llm-long-seq-sp-train-mem-down-75pc.html