【Bug已解决】Feature request: FSDP2 QLoRA 解决方案-育师

【Bug已解决】Feature request: FSDP2 QLoRA 解决方案

一、现象长什么样

想用 FSDP2 做多卡训练，同时用 QLoRA（4-bit 量化基座 + LoRA 适配器）省显存。但翻accelerate/torch文档找不到"FSDP2 + QLoRA"的开箱支持，自己拼起来要么报错要么静默损坏：

# 形态一：fully_shard 把 4-bit 基座也分片，反量化状态错位 RuntimeError： quant_state shape mismatch after fully_shard # 形态二：4-bit 参数被 FSDP2 当成普通参数处理，梯度流断 ValueError： cannot compute grad for quantized param # 形态三：显存没省下来 峰值反而比全精度 FSDP2 更高

最小判据：

触发：FSDP2(fully_shard) + 4-bit 量化基座 + LoRA 现象：quant_state 错位 / 梯度断 / 无省显存 根因：FSDP2 默认对所有参数分片，但 4-bit 基座的量化元数据不能被分片破坏 影响：无法用 FSDP2 QLoRA 做省显存多卡微调

最迷惑的是：FSDP2 对普通模型很香，QLoRA 单卡也很成熟，但两者组合没有现成路径——因为 4-bit 基座的quant_state（缩放因子、零点）是和"整块权重"绑定的，被fully_shard切开后反量化就错。

二、背景

QLoRA 的核心：基座权重用 4-bit（NF4）量化存储，冻结；只训练 LoRA 的A/B低秩矩阵（浮点）。前向时把 4-bit 权重反量化成 bf16 参与计算，反向时梯度只流向 LoRA 参数。

FSDP2 的fully_shard会把参数切分到多卡，并管理 all-gather / reduce-scatter。问题：

4-bit 基座参数（Params4bit）带着quant_state（量化元数据）。fully_shard若把它当普通参数切分，会破坏"权重块与 quant_state 的对应"——反量化需要整块权重 + 对应缩放，切开后缩放对不上；
4-bit 参数是冻结的、不应有梯度，也不应被 FSDP2 的 all-gather 通信管理（它是常量，可被各卡本地反量化）；
LoRA 的A/B是浮点、需要梯度、可以（也建议）分片以省显存。

正确的 FSDP2 QLoRA 配方是：4-bit 基座保持本地完整（不被 fully_shard 切分，各卡持有完整量化权重，本地反量化），只对 LoRA 浮点参数做 fully_shard。这样：

基座省了 8x 显存（4-bit），且每卡本地反量化无需跨卡通信；
LoRA 参数被分片，多卡可训更大 rank 的 LoRA；
反量化状态不被破坏。

根因是"FSDP2 默认对所有参数（含 4-bit 基座）分片，破坏了 quant_state"。

三、根因

抽象成代码（示意）：

def fsdp2_qlora_naive(model): # BUG：对整模型 fully_shard，4-bit 基座也被切 for m in model.modules(): if has_params(m): fully_shard(m) # 4-bit 基座的 quant_state 被切开 -> 错位

根因链条：

QLoRA 基座是 4-bit +quant_state，需"整块对应"；
fully_shard默认切分所有参数，破坏 quant_state 对应；
反量化缩放对不上 -> 数值错 / RuntimeError；
冻结的 4-bit 参数被纳入通信管理，浪费且没必要；
正确做法：基座本地完整、只分片 LoRA。

一句话：FSDP2 默认分片所有参数，把 4-bit 基座的 quant_state 切坏，QLoRA 不可用。

四、最小可运行复现

用纯 Python 模拟"量化权重被切分后缩放对应错"：

# repro_fsdp2_qlora.py class QuantBlock: def __init__(self, weight, scales): self.weight = weight # 4-bit 权重（块） self.scales = scales # 每块一个缩放 def dequant(block): # 反量化需要 weight 块与 scales 一一对应 if len(block.weight) != len(block.scales): raise RuntimeError("quant_state 与权重块不匹配") return [w * s for w, s in zip(block.weight, block.scales)] def shard_block(block, shards): # BUG：把权重块切开但 scales 没跟着切 return block.weight[:shards], block.scales # scales 仍是整体 -> 错位 def main(): block = QuantBlock(weight=[1,2,3,4], scales=[0.1,0.1,0.1,0.1]) w_shard, scales = shard_block(block, 2) bad = QuantBlock(w_shard, scales) try: dequant(bad) except RuntimeError as e: print("复现成功 ->", e) if __name__ == "__main__": main()

运行输出：

复现成功 -> quant_state 与权重块不匹配

4-bit 权重块被切、scales 没跟随，反量化错位，正是真实 bug 的抽象。

五、解决方案（第一层：最小直接修复）

最小且必须的一步：只对 LoRA 浮点参数做fully_shard，跳过 4-bit 基座。基座保持本地完整，各卡本地反量化：

# fix_layer1.py from torch.distributed.fsdp import fully_shard def fsdp2_qlora(model): for name, module in model.named_modules(): # 只 shard 含可训练浮点参数的模块（LoRA），跳过 4-bit 基座 has_trainable_float = any( p.requires_grad and not is_quantized(p) for p in module.parameters(recurse=False) ) if has_trainable_float: fully_shard(module) # 4-bit 基座：不 fully_shard，保持本地完整 return model def is_quantized(p): return hasattr(p, "quant_state") or type(p).__name__ == "Params4bit"

要点：

用is_quantized识别 4-bit 参数，跳过其模块；
只对 LoRA 浮点参数fully_shard，分片省显存；
基座每卡本地反量化，无需跨卡通信，quant_state 不被破坏。

六、解决方案（第二层：结构性改进）

把"QLoRA + FSDP2 的分片决策"做成显式策略：基座（量化、冻结）标记no_shard，LoRA（浮点、可训练）标记shard，由策略统一驱动：

# fix_layer2.py from dataclasses import dataclass, field from typing import List @dataclass class ParamRole: name: str quantized: bool trainable: bool @property def shard(self) -> bool: # 只有"浮点且可训练"的参数才分片；量化/冻结的不分片 return (not self.quantized) and self.trainable class QLoRAFsdp2Planner: def __init__(self, roles: List[ParamRole]): self.roles = roles def plan(self): return {r.name: ("shard" if r.shard else "no_shard") for r in self.roles} # 用法 roles = [ ParamRole("base.weight", quantized=True, trainable=False), # no_shard ParamRole("lora_A.weight", quantized=False, trainable=True), # shard ParamRole("lora_B.weight", quantized=False, trainable=True), # shard ] planner = QLoRAFsdp2Planner(roles) print(planner.plan()) # -> {'base.weight':'no_shard', 'lora_A.weight':'shard', 'lora_B.weight':'shard'}

要点：

ParamRole.shard用"非量化且可训练"作为分片判据，语义清晰；
QLoRAFsdp2Planner统一产出分片计划，4-bit 基座天然no_shard；
任何新参数类型只需填quantized/trainable，无需改分片逻辑。

七、解决方案（第三层：断言 / CI 守护）

写 pytest 验证"4-bit 基座不分片、LoRA 分片"：

# test_fsdp2_qlora.py import pytest def is_quantized(p): return getattr(p, "quantized", False) def decide_shard(params): plan = {} for name, p in params.items(): plan[name] = "shard" if (not is_quantized(p) and p["trainable"]) else "no_shard" return plan def test_base_not_sharded(): params = {"base.weight": {"quantized": True, "trainable": False}} plan = decide_shard(params) assert plan["base.weight"] == "no_shard" def test_lora_sharded(): params = {"lora_A.weight": {"quantized": False, "trainable": True}} plan = decide_shard(params) assert plan["lora_A.weight"] == "shard" def test_quant_state_preserved(): # 基座不分片 -> quant_state 完整 plan = decide_shard({"base.weight": {"quantized": True, "trainable": False}}) assert plan["base.weight"] == "no_shard"

CI 一旦有人把 4-bit 基座也分片，test_base_not_sharded立刻变红。

八、排查清单

FSDP2 QLoRA 报错 / 无省显存时：

确认是否fully_shard把 4-bit 基座也切了（quant_state 错位）；
检查基座是否标记了"不分片"，只有 LoRA 浮点参数被fully_shard；
确认基座是冻结的（requires_grad=False），不参与梯度；
按第五 / 六节用ParamRole.shard判据统一分片；
基座应每卡本地完整、本地反量化，无需跨卡通信；
若显存没省，确认基座确实 4-bit 且未被全精度副本占用；
把第七节的 pytest 接进 CI，守护"4-bit 基座不分片"。

九、小结

FSDP2 QLoRA 缺开箱支持，根因是 FSDP2 默认对所有参数（含 4-bit 基座）分片，破坏 QLoRA 基座"权重块与 quant_state 一一对应"的反量化前提，导致 quant_state 错位 / 梯度断。正确配方是：4-bit 基座保持本地完整（不分片、本地反量化），只对 LoRA 浮点参数fully_shard。

三层层级：