【Bug已解决】Feature request: FSDP2 QLoRA 解决方案一、现象长什么样想用 FSDP2 做多卡训练同时用 QLoRA4-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-bitNF4量化存储冻结只训练 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_shard4-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(recurseFalse) ) 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_shardLoRA浮点、可训练标记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, quantizedTrue, trainableFalse), # no_shard ParamRole(lora_A.weight, quantizedFalse, trainableTrue), # shard ParamRole(lora_B.weight, quantizedFalse, trainableTrue), # 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_shardCI 一旦有人把 4-bit 基座也分片test_base_not_sharded立刻变红。八、排查清单FSDP2 QLoRA 报错 / 无省显存时确认是否fully_shard把 4-bit 基座也切了quant_state 错位检查基座是否标记了不分片只有 LoRA 浮点参数被fully_shard确认基座是冻结的requires_gradFalse不参与梯度按第五 / 六节用ParamRole.shard判据统一分片基座应每卡本地完整、本地反量化无需跨卡通信若显存没省确认基座确实 4-bit 且未被全精度副本占用把第七节的 pytest 接进 CI守护4-bit 基座不分片。九、小结FSDP2 QLoRA 缺开箱支持根因是 FSDP2 默认对所有参数含 4-bit 基座分片破坏 QLoRA 基座权重块与 quant_state 一一对应的反量化前提导致 quant_state 错位 / 梯度断。正确配方是4-bit 基座保持本地完整不分片、本地反量化只对 LoRA 浮点参数fully_shard。三层层级第一层只对 LoRA 浮点参数fully_shard跳过 4-bit 基座第二层用ParamRole.shard非量化且可训练作为分片判据统一规划第三层pytest 验证 4-bit 基座不分片、LoRA 分片锁进 CI。核心教训QLoRA 的量化权重是块 元数据绑定的任何分片框架在切它之前都必须确认元数据随块一起切或不切。把量化/冻结参数排除出分片是 FSDP2 QLoRA 能成立的前提——本系列第 504 篇从自动排除机制、本篇从端到端配方两个角度覆盖了它。
【Bug已解决】Feature request: FSDP2 QLoRA 解决方案
【Bug已解决】Feature request: FSDP2 QLoRA 解决方案一、现象长什么样想用 FSDP2 做多卡训练同时用 QLoRA4-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-bitNF4量化存储冻结只训练 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_shard4-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(recurseFalse) ) 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_shardLoRA浮点、可训练标记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, quantizedTrue, trainableFalse), # no_shard ParamRole(lora_A.weight, quantizedFalse, trainableTrue), # shard ParamRole(lora_B.weight, quantizedFalse, trainableTrue), # 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_shardCI 一旦有人把 4-bit 基座也分片test_base_not_sharded立刻变红。八、排查清单FSDP2 QLoRA 报错 / 无省显存时确认是否fully_shard把 4-bit 基座也切了quant_state 错位检查基座是否标记了不分片只有 LoRA 浮点参数被fully_shard确认基座是冻结的requires_gradFalse不参与梯度按第五 / 六节用ParamRole.shard判据统一分片基座应每卡本地完整、本地反量化无需跨卡通信若显存没省确认基座确实 4-bit 且未被全精度副本占用把第七节的 pytest 接进 CI守护4-bit 基座不分片。九、小结FSDP2 QLoRA 缺开箱支持根因是 FSDP2 默认对所有参数含 4-bit 基座分片破坏 QLoRA 基座权重块与 quant_state 一一对应的反量化前提导致 quant_state 错位 / 梯度断。正确配方是4-bit 基座保持本地完整不分片、本地反量化只对 LoRA 浮点参数fully_shard。三层层级第一层只对 LoRA 浮点参数fully_shard跳过 4-bit 基座第二层用ParamRole.shard非量化且可训练作为分片判据统一规划第三层pytest 验证 4-bit 基座不分片、LoRA 分片锁进 CI。核心教训QLoRA 的量化权重是块 元数据绑定的任何分片框架在切它之前都必须确认元数据随块一起切或不切。把量化/冻结参数排除出分片是 FSDP2 QLoRA 能成立的前提——本系列第 504 篇从自动排除机制、本篇从端到端配方两个角度覆盖了它。
