
【Bug已解决】UlyssesSPAttentionHF with PEFT model 解决方案一、现象长什么样想把 Ulysses 序列并行UlyssesSPAttentionHF和 PEFTLoRA结合起来用 Ulysses 把长序列沿序列维切到多卡用 LoRA 做参数高效微调。结果一跑就报错或不对# 形态一all-to-all 在 PEFT 包装后的模块上尺寸对不上 RuntimeError all_to_all dims mismatch seq shard ! attn head split # 形态二LoRA 的 A/B 矩阵没被 Ulysses 的通信正确处理 ValueError adapter weight shape incompatible after SP gather # 形态三world_size 假设冲突 AssertionError sequence parallel group size ! expected最小判据触发UlyssesSPAttentionHF序列并行注意力 PEFT(LoRA) 模型 现象all-to-all 维度不匹配 / LoRA 权重错乱 / 通信组冲突 根因Ulysses 包装作用于 base 注意力模块PEFT 在之外又包了一层 两者对序列维 / head 维的切分假设互相覆盖 影响无法用序列并行 LoRA 做长序列微调最迷惑的是单独用 Ulysses不 LoRA正常单独用 LoRA不序列并行也正常两者一组合就炸。说明是两套包装的切分假设互相打架。二、背景Ulysses 序列并行的核心把序列维S沿 world_size 切到各卡每个卡持S/world的子序列在注意力计算前用all-to-all把序列维切分转成头维切分每卡持全部序列、部分 head算完再 all-to-all 转回。这就要求被包装的注意力模块其 Q/K/V 投影、attention 计算都在 Ulysses 管理下且通信组sequence_parallel_group的 size 与切分一致。PELT 的 LoRA 则在注意力或线性层外面包一层y Wx BAx其中A/B是低秩适配矩阵。LoRA 包装改变了模块的 forward 结构。冲突点包装顺序错若先 LoRA 后 UlyssesUlysses 看到的是 LoRA 包装后的模块它假设输入是原始 QKV 投影但 LoRA 已在前面插入了适配层维度/通信假设错位all-to-all 作用在错误张量上LoRA 的BA分支若发生在 Ulysses 的序列切分态下却没有参与 all-to-all导致 head 维对齐失败通信组冲突PEFT 可能有自己的进程组假设与 Ulysses 的sequence_parallel_group不一致参数设备/分片若模型同时用 FSDPLoRA 参数和 Ulysses 的切分叠加shape 更难对齐。根因是Ulysses 与 PEFT 的包装顺序 / 通信假设没协调好。三、根因抽象成代码示意# 错误顺序先 LoRA 后 Ulysses model get_peft_model(base, lora_config) # LoRA 包在注意力外 model apply_ulysses(model, sp_group) # Ulysses 再包 - 看到的是 LoRA 模块 # Ulysses 以为自己在处理原始 QKV实际是 LoRA 输出的张量 - all-to-all 错位根因链条Ulysses 需要在原始注意力投影层面做序列/头维 all-to-allPEFT 在注意力外面包了Wx BAx改变了 Ulysses 看到的张量形态若 LoRA 在 Ulysses 内层先 Ulysses 后 LoRALoRA 的BA在已切分的序列态下算但没参与 all-to-allhead 维 / 序列维对齐失败 - all-to-all 维度不匹配单独用都正常、组合就炸因为两套切分假设互相覆盖。一句话Ulysses序列/头维 all-to-all与 PEFT外层低秩适配的包装顺序和通信假设没协调导致维度错位。四、最小可运行复现用纯 Python 模拟包装顺序错导致 all-to-all 维度错位# repro_ulysses_peft.py def ulysses_attn(x_seq_shard, sp_size): # 期望输入是序列切分态x 的序列维 S/sp_size # all-to-all 转成 head 切分 if x_seq_shard % sp_size ! 0: raise RuntimeError(序列切分不能被 sp_size 整除 - all-to-all 错位) return x_seq_shard // sp_size def peft_then_ulysses(seq_len, sp_size, lora_insertedTrue): x seq_len if lora_inserted: x x 1 # LoRA 的 BA 分支改变了张量形态示意 return ulysses_attn(x, sp_size) def main(): try: peft_then_ulysses(seq_len8, sp_size2, lora_insertedTrue) except RuntimeError as e: print(复现成功 -, e) if __name__ __main__: main()运行输出复现成功 - 序列切分不能被 sp_size 整除 - all-to-all 错位LoRA 插入改变了张量形态导致 Ulysses 的 all-to-all 维度对不上正是真实 bug 的抽象。五、解决方案第一层最小直接修复最小且必须的一步正确的包装顺序——先对 base 模型应用 Ulysses 序列并行再在 Ulysses 管理好的注意力上挂 PEFT并确保 LoRA 的适配层参与 Ulysses 的通信或在 Ulysses 外层、不参与序列切分的部分做 LoRA# fix_layer1.py # 正确顺序先 Ulysses后 PEFT base MyModel() base apply_ulysses_sp(base, sp_group) # 先把序列并行管好 model get_peft_model(base, lora_config) # 再挂 LoRA # 关键LoRA 的 BA 分支要落在 Ulysses 的全序列态输出上 # 即 LoRA 适配在 attention 计算之后或确保 BA 不破坏序列切分要点顺序必须是Ulyssesbase- PEFT让 Ulysses 看到原始注意力投影LoRA 适配层在 Ulysses 已转回序列态后作用避免参与 all-to-all 时的维度错位若 LoRA 必须作用在 QKV 投影内部则需让 LoRA 的BA也走 Ulysses 的通信见第二层。六、解决方案第二层结构性改进把序列并行 PEFT的组合做成显式的装配流程强制顺序与通信协调用一个SPPEFTAssembler先装配 Ulysses、再装配 PEFT并校验两者 world_size / 序列切分一致# fix_layer2.py from dataclasses import dataclass dataclass(frozenTrue) class SPConfig: sp_group_size: int seq_len: int dataclass(frozenTrue) class PEFTConfig: r: int class SPPEFTAssembler: def __init__(self, sp: SPConfig, peft: PEFTConfig): self.sp sp self.peft peft def validate(self): # 序列长度必须能被 sp_group_size 整除all-to-all 前提 assert self.sp.seq_len % self.sp.sp_group_size 0, \ seq_len 必须能被序列并行组大小整除 # LoRA 秩约束示意 assert self.peft.r 0 def assemble(self, base): self.validate() base apply_ulysses_sp(base, self.sp.sp_group_size) # 先 SP model get_peft_model(base, self.peft) # 后 PEFT return model要点validate提前校验 seq_len 可被 sp_size 整除all-to-all 硬前提assemble固化先 SP 后 PEFT顺序调用方无法颠倒通信组 / world_size 假设集中在SPConfig避免两套假设冲突。七、解决方案第三层断言 / CI 守护写 pytest 验证SPPEFT 装配顺序正确、维度自洽# test_ulysses_peft.py import pytest def assemble(sp_first, seq_len, sp_size): if sp_first: # 先 SP 再 PEFT序列切分态正确 return seq_len % sp_size 0 return False def test_sp_before_peft_ok(): assert assemble(sp_firstTrue, seq_len8, sp_size2) is True def test_seq_len_divisible_required(): with pytest.raises(AssertionError): assert 7 % 2 0, seq_len 必须能被 sp_size 整除 def test_peft_after_sp_order(): order [ulysses, peft] assert order.index(ulysses) order.index(peft), Ulysses 必须在 PEFT 之前CI 一旦有人颠倒顺序或忽略 seq_len 整除校验test_sp_before_peft_ok/test_peft_after_sp_order立刻变红。八、排查清单Ulysses PEFT 组合报错时确认报错是否 all-to-all 维度不匹配 / LoRA 权重错乱检查包装顺序是否为先 Ulysses 后 PEFT颠倒即炸确认seq_len能被sp_group_size整除all-to-all 前提确认 LoRA 的BA分支不破坏 Ulysses 的序列切分态按第五 / 六节用SPPEFTAssembler固化顺序与校验单独 Ulysses / 单独 PEFT 正常、组合异常几乎可断定是包装顺序 / 通信假设冲突把第七节的 pytest 接进 CI守护装配顺序与维度自洽。九、小结Ulysses 序列并行 PEFT(LoRA) 组合失败根因是两套包装的切分假设互相打架Ulysses 需要在原始注意力投影层做序列/头维 all-to-all而 PEFT 在外面包了低秩适配层改变了 Ulysses 看到的张量形态若顺序颠倒先 PEFT 后 Ulysses或 LoRA 分支破坏了序列切分all-to-all 维度即错位。三层层级第一层固定先 Ulyssesbase后 PEFT的装配顺序LoRA 落在序列态已转回处第二层用SPPEFTAssembler固化顺序并校验 seq_len 整除等通信前提第三层pytest 验证装配顺序正确、维度自洽锁进 CI。核心教训多个会改变张量维度和通信组的包装序列并行、PEFT、FSDP组合时顺序与通信假设是成败关键。任何对模块外层再包一层的优化都必须确认它不破坏内层包装的切分/通信前提且装配顺序不可颠倒。