【Bug已解决】Question about multiple LoRAs and FSDP 解决方案

【Bug已解决】Question about multiple LoRAs and FSDP 解决方案一、现象长什么样很多人想在同一个基座模型上挂多个 LoRA adapter比如一个负责“代码风格”、一个负责“中文语气”、一个负责“安全对齐”并且用FSDP做多卡分片训练。组合时会冒出一堆疑问和坑多个 adapter 都add_adapter之后print_trainable_parameters()显示可训练参数 所有 adapter 之和FSDP 把它们全部分片了显存没省以为只激活一个就只占一份用set_adapter([a, b])同时激活两个 adapter 做加权融合推理但 FSDP 下报RuntimeError: ... parameter ... not sharded或激活的 adapter 权重不在本 rank切换set_adapter(b)后FSDP 的summon_full_params/ 梯度通信仍然按“全部 adapter 都参与”处理导致某些 rank 上b的lora_A/lora_B是空分片训练时只想训其中一个 adapter其余冻结但 FSDP 在no_sync/ 梯度分片时把冻结 adapter 的参数也算进通信组浪费带宽保存时save_pretrained只存了“当前激活”的 adapter其它 adapter 丢了或load_adapter回来后 FSDP 包裹状态错乱。一句话多 LoRA FSDP 的核心矛盾是——FSDP 在初始化时就对“所有存在的参数”做分片而多 adapter 意味着多份 LoRA 参数同时存在且 adapter 的激活/切换是运行期行为FSDP 的静态分片不知道“此刻谁活跃”。二、背景先厘清两个事实FSDP 分片是静态的FullyShardedDataParallel(model)在包裹时遍历所有参数把每个参数切成world_size份每 rank 持一份分片。它不关心参数当前requires_grad是 True 还是 False也不关心它属于哪个 adapter——只要参数是模型的一部分就被分片。多 adapter 的参数是“并存”的add_adapter(a)给每个目标层加一份lora_A.a/lora_B.aadd_adapter(b)再加一份lora_A.b/lora_B.b。它们都挂在模型树上FSDP 全部照单分片。于是显存多 adapter 的 LoRA 参数总和被分片省的是“相比全参数”的部分但多份 LoRA 之间不互相省——挂 N 个 adapter就有 N 份 LoRA 参数分片。这是预期不是 bug但常被误解。激活切换set_adapter(b)只是把前向里“用哪份 lora 权重”切到b参数本身仍在、仍被分片。FSDP 不需要重新分片这本应工作。出错往往是因为切换后某些 rank 没同步active_adapter状态或推理代码直接读了lora_A.default这种固定名字而非按active_adapter索引。同时激活多个set_adapter([a,b])在前向里对多份 lora 做加权求和。FSDP 下每份都被分片求和需要每份都 all-gather 到完整通信量随 adapter 数线性增长但不会报错——除非代码写错成只 gather 了一份。下面用最小可运行代码演示“多 adapter 并存 FSDP 分片语义 正确切换”。三、根因根因一句话FSDP 的参数是静态分片的初始化时对所有存在的 LoRA 参数做分片而多 adapter 的参数在模型树上是“并存”的、adapter 的激活是运行期行为混淆“参数并存”与“参数激活”会引发显存/通信/切换三方面的误用。展开显存误解多 adapter 多份 LoRA 参数FSDP 全部分片不互相省。切换状态不同步set_adapter后各 rank 的active_adapter必须一致否则前向用的权重分片对不上。固定名字读取代码若写死读lora_A.default而非按active_adapter索引切换后读错副本。冻结 adapter 仍占通信训练只训一个但 FSDP 对所有参数做梯度分片冻结的也参与通信组可优化但非错误。修复方向明确“参数并存、激活切换”的语义切换时保证所有 rankactive_adapter一致训练时用requires_grad_(False)冻结非目标 adapter 以减少无用通信保存时确认存了哪些 adapter。四、最小可运行复现下面用单进程演示“多 adapter 并存 切换激活 训练只训一个”FSDP 多卡部分用注释说明单卡可跑逻辑。import torch import torch.nn as nn from peft import get_peft_model, LoraConfig, PeftModel class Tiny(nn.Module): def __init__(self): super().__init__() self.fc1 nn.Linear(16, 16) self.fc2 nn.Linear(16, 4) def forward(self, x): return self.fc2(torch.relu(self.fc1(x))) base Tiny() # 加两个 adaptera代码风格、b中文语气 m get_peft_model(base, LoraConfig(r4, lora_alpha8, target_modules[fc1, fc2], lora_dropout0.0)) m.add_adapter(b, LoraConfig(r4, lora_alpha8, target_modules[fc1, fc2], lora_dropout0.0)) # 两个 adapter 的参数都并存于模型树上 names [n for n, _ in m.named_parameters() if lora_ in n] print(并存参数含 a 和 b:, any(default in n for n in names), any(b in n for n in names)) # 切换激活前向只走 b m.set_adapter(b) out_b m(torch.randn(2, 16)) # 切换回 a m.set_adapter(default) out_a m(torch.randn(2, 16)) print(a / b 两个 adapter 都能前向:, out_a.shape out_b.shape) # 训练时只训 b冻结 a减少 FSDP 下无用梯度通信 m.set_adapter(b) for n, p in m.named_parameters(): p.requires_grad_( lora_ in n and .b. in n ) trainable sum(p.numel() for p in m.parameters() if p.requires_grad) print(只训 b 时可训练参数:, trainable)运行后能看到 a、b 两份参数并存切换激活都能前向且训练时只b可训练。FSDP 多卡下这份“并存参数”会被全部分片切换激活不改变分片因为参数一直都在。五、解决方案第一层最小直接修复修复 1切换 adapter 时所有 rank 保持一致# 分布式下每个 rank 都必须 set 成同一个 adapter m.set_adapter(b) # 或在切换前 barrier确保状态同步 # dist.barrier()修复 2训练只训目标 adapter冻结其余如上文遍历参数只有.目标名.的lora_参数requires_grad_(True)其余False。FSDP 仍会对全部参数分片但冻结的不产生梯度、不参与优化器更新减少无用通信。修复 3读权重按active_adapter索引别写死default# 错误直接读 lora_A.default.weight # 正确用 m.get_adapter_params() 或按当前 active_adapter 取 active m.active_adapter w dict(m.named_parameters())[fbase_model.model.fc1.lora_A.{active}.weight]六、解决方案第二层结构性改进改进 1用set_adapter([...])做加权融合推理# 同时激活多个按权重融合adapter_weights 需与列表等长 m.set_adapter([default, b]) m.set_adapter_merge_weights(True) # 若 PEFT 支持否则手动加权 # 前向自动对多份 lora 求和FSDP 下每份都被分片通信量随 adapter 数增长改进 2FSDP 包裹前先确定最终要用的 adapter 集合from torch.distributed.fsdp import FullyShardedDataParallel as FSDP # 先 add 好所有需要的 adapter再 FSDP 包裹确保分片涵盖全部 m PeftModel.from_pretrained(base, adapter_a) m.add_adapter(b, LoraConfig(r4, target_modules[fc1,fc2])) m m.cuda() m FSDP(m, device_idtorch.cuda.current_device()) # 此刻对所有 lora 参数分片改进 3按需“卸载不用的 adapter”节省显存如果某些 adapter 长期不用训练前delete_adapter移除FSDP 重新包裹或接受其分片仍在但不再参与计算m.delete_adapter(b) # 用完即删减少并存参数七、解决方案第三层断言 / CI 守护import torch import torch.nn as nn import pytest from peft import get_peft_model, LoraConfig class Tiny(nn.Module): def __init__(self): super().__init__() self.fc1 nn.Linear(16, 16) self.fc2 nn.Linear(16, 4) def forward(self, x): return self.fc2(torch.relu(self.fc1(x))) def _make(): m get_peft_model(Tiny(), LoraConfig(r4, lora_alpha8, target_modules[fc1,fc2], lora_dropout0.0)) m.add_adapter(b, LoraConfig(r4, lora_alpha8, target_modules[fc1,fc2], lora_dropout0.0)) return m def test_multiple_adapters_coexist(): m _make() names [n for n, _ in m.named_parameters() if lora_ in n] assert any(default in n for n in names) assert any(.b. in n for n in names) def test_switch_adapter_changes_active(): m _make() m.set_adapter(b) assert m.active_adapter b m.set_adapter(default) assert m.active_adapter default def test_train_only_target_adapter(): m _make() m.set_adapter(b) for n, p in m.named_parameters(): p.requires_grad_( lora_ in n and .b. in n) tr sum(p.numel() for p in m.parameters() if p.requires_grad) # 只有 b 的两个线性层的 lora 参数可训练 assert tr 2 * (4*16 16*4) # 每份 lora A(4*16)B(16*4)两份层 def test_frozen_adapter_not_trainable(): m _make() m.set_adapter(b) for n, p in m.named_parameters(): p.requires_grad_( lora_ in n and .b. in n) a_trainable any(p.requires_grad for n, p in m.named_parameters() if default in n and lora_ in n) assert a_trainable is False这四个测试守护“多 adapter 并存、切换改变 active、只训目标 adapter、冻结的不可训练”。八、排查清单多 LoRA FSDP 出问题时按序查明确语义参数是“并存”的adapter 激活是“运行期切换”的FSDP 静态分片全部参数。显存预期挂 N 个 adapter 就有 N 份 LoRA 参数分片不互相省。切换同步set_adapter后所有 rank 的active_adapter必须一致。按 active 索引读权重别写死lora_A.default。训练冻结非目标 adapterrequires_grad_(目标名 in n)减少无用通信。多 adapter 融合set_adapter([a,b])加权求和FSDP 下通信量随 adapter 数增长。保存确认save_pretrained存的是当前激活 adapter多 adapter 要分别存或用adapter_name。用完删除长期不用的 adapterdelete_adapter减少并存参数。九、小结Question about multiple LoRAs and FSDP的核心矛盾是FSDP 在初始化时对所有存在的参数做静态分片而多 LoRA 的参数是“并存”在模型树上的、adapter 激活是运行期行为混淆二者会引发显存预期错误、切换状态不同步、按固定名字读错副本等问题。最小修复是切换 adapter 时所有 rank 保持active_adapter一致、训练时requires_grad_(目标名 in n)冻结非目标 adapter、读权重按active_adapter索引结构性改进是用set_adapter([a,b])做加权融合、FSDP 包裹前先 add 齐所有 adapter、长期不用的delete_adapter最后用测试守护“多 adapter 并存、切换改变 active、只训目标、冻结不可训”。厘清“参数并存 vs 激活切换”多 LoRA 就能在 FSDP 下正确工作。