FP8 LoRA 训练在整块包装的模型上梯度丢失,权重保持初始值
維護者通常 1 天內回覆
評估
- 難度
- 4/5
- 預估耗時
- 3-5 天
- 新手友好度
- 49/100
- Issue 類型
- 缺陷
- 描述清晰度
- 描述清楚
- 活躍度
- 活躍
研究方向
Start with diffsynth/core/vram/layers.py and compare AutoWrappedModule.cast_to with the override in AutoWrappedNonRecurseModule.cast_to; the issue includes a minimal reproduction using a small wrapped block and PEFT LoRA. Run that reproduction and check whether lora_B.grad is present, then verify that LoRA parameters update for both wrapper types while frozen weights still use converted copies.
由索引模型根據 Issue 內容生成。
描述
现象
--fp8_models 做 LoRA 训练时,循环能正常跑完,loss 也有正常数值,但 LoRA 权重从未更新。
在 Fibo-2.0-dev 上跑完 12800 step、每 500 step 存一次,得到 26 个检查点。这些文件字节完全相同。其中 364 个 lora_B 的最大绝对值是 0,lora_A 仍是 PEFT 的初始值。拿到 ComfyUI 里 strength 1.0 的出图和底模是同一张。
官方 examples/fibo/model_training/lora/Fibo-2.0-dev-T2I.sh 没有开 --fp8_models,所以示例路径碰不到。16GB 显卡放不下 BF16 的 Fibo DiT,才会走文档里的 FP8 存储。
原因
diffsynth/core/vram/layers.py 里 AutoWrappedModule.cast_to 在存储 dtype 和计算 dtype 不同时执行:
return copy.deepcopy(module).to(dtype=dtype, device=device)
FP8 训练的配置是 offload / onload / preparing 为 float8_e4m3fn,computation 为 bfloat16。每次 forward 都会整块 deepcopy。
Fibo 的显存映射把整个 FiboTransformerBlock 和 FiboPerceiverResampler 包成 AutoWrappedModule,递归在这里停止,里面的 nn.Linear 不会再包成 AutoWrappedLinear。PEFT 把 LoRA 注入到 .module 内部。反传更新的是这份复制品,原参数的 grad 是 None,优化器不会改它们。存盘读到的就是初始 LoRA。
--find_unused_parameters 会让这些没有梯度的参数不报错,所以训练看起来是成功的。
影响范围
不是所有模型。Qwen-Image、Z-Image、FLUX、FLUX.2、FLUX.3 的 DiT 按每个 Linear 包成 AutoWrappedLinear,LoRA 加在类型转换外面。官方 Qwen FP8 LoRA 示例走的是这条,不受影响。
会丢梯度的是「整块包装 + FP8」:
FiboTransformerBlock、FiboPerceiverResampler:AutoWrappedModule,整块 deepcopy。- Wan 主 DiT 的
DiTBlock:AutoWrappedNonRecurseModule。它会共享子模块里的AutoTorchModule,但 PEFT 的 LoRA 层本身不是AutoTorchModule,仍会被复制。下面的小例子把grad打成了None。 - Vace 的
DiTBlock、Wan S2V 的 block 也是整块AutoWrappedModule。
docs/zh/Training/FP8_Precision.md 写的是:冻结权重可以用 FP8 存放,LoRA 的梯度仍然更新原参数。现在这条对整块包装的模型不成立。
最小复现
不必下载完整权重。存储用 float16、计算用 float32,走的是同一个 cast_to 分支。
import torch
import torch.nn as nn
from peft import LoraConfig, inject_adapter_in_model
from diffsynth.core.vram.layers import AutoWrappedModule
class Block(nn.Module):
def __init__(self):
super().__init__()
self.to_q = nn.Linear(8, 8, bias=False)
def forward(self, x):
return self.to_q(x)
block = Block()
for p in block.parameters():
p.requires_grad = False
block.to(dtype=torch.float16)
kw = dict(
offload_dtype=torch.float16, offload_device="cpu",
onload_dtype=torch.float16, onload_device="cpu",
preparing_dtype=torch.float16, preparing_device="cpu",
computation_dtype=torch.float32, computation_device="cpu",
)
wrapped = AutoWrappedModule(block, **kw, name="block")
wrapped = inject_adapter_in_model(
LoraConfig(r=2, lora_alpha=2, target_modules=["to_q"]), wrapped
)
for p in wrapped.parameters():
if p.requires_grad:
p.data = p.data.float()
y = wrapped(torch.randn(2, 8))
y.float().pow(2).mean().backward()
for n, p in wrapped.named_parameters():
if "lora_B" in n:
print(n, None if p.grad is None else float(p.grad.abs().max()))
结果是 lora_B 仍为 0,grad 为 None。
Wan 主 DiT 的包装同样会丢。先把 to_q 包成 AutoWrappedLinear,再把 block 包成 AutoWrappedNonRecurseModule,然后同样注入 LoRA,lora_B.grad 还是 None。
修复方向
cast_to 保持现在的复制和 dtype 转换。复制完成之后,把副本里 requires_grad=True 的参数重新指回原模块上的 Parameter。冻结权重继续用转换后的计算精度副本。没有可训练参数时,行为和现在一样,推理不受影响。
AutoWrappedNonRecurseModule.cast_to 重写了这个函数,需要同样处理,否则 Wan 主 DiT 还是接不上梯度。
- 主要語言
- Python
- 星號
- 13.3k
- 分支
- 1.3k
- 平均合併
- 1 天 14 小時
- 30 天內合併 PR
- 43
環境準備
這個專案沒有提供開發容器、Dockerfile 或貢獻指南,環境需要你自己搭建:先看它的 README,通用步驟見我們的新手貢獻指南。
從這裡開始
- 先讀完整個 Issue,再讀專案的貢獻指南。
- 在 Issue 下留言說明你要接手 —— 這能避免兩個人做同樣的事。
- Fork 儲存庫,在一個分支上完成修改。
- 送出 Pull Request,並在描述裡引用這個 Issue 編號。
modelscope/DiffSynth-Studio 的其他 Issue
-
Partial download surfaces as "Cannot detect the model type" rather than a download error可能重新可做 關聯的 PR 已關閉且未合併。 未關閉
難度 1/5 1-3 小時 新手友好度 78/100
modelscope/DiffSynth-Studio#1668 ·
維護者通常 1 天內回覆
-
bfloat16训练时段错误: 数组越界的一种方案可能重新可做 關聯的 PR 已關閉且未合併。 未關閉
難度 2/5 1-3 小時 新手友好度 72/100
modelscope/DiffSynth-Studio#1499 · 1 則留言 ·
維護者通常 1 天內回覆
-
難度 1/5 1 小時以內 新手友好度 78/100
modelscope/DiffSynth-Studio#1373 · 5 則留言 · 1 個 reaction ·
維護者通常 1 天內回覆
-
難度 3/5 1-2 天 新手友好度 48/100
modelscope/DiffSynth-Studio#1790 ·
維護者通常 1 天內回覆
-
OPD做步数蒸馏的疑问未關閉
難度 4/5 3-5 天 新手友好度 30/100
modelscope/DiffSynth-Studio#1786 · 3 則留言 ·
維護者通常 1 天內回覆
查看 modelscope/DiffSynth-Studio 的全部 Issue
相似的 Issue
-
難度 2/5 1-3 小時 新手友好度 82/100
維護者通常 2 天內回覆
-
難度 2/5 1-3 小時 新手友好度 86/100
-
Validate finite output clamping limits in StepPredictor.prepare_clamping_params to prevent silent NaN/Inf propagation可能已有人在做 @muskan-creates352 今天認領。 未關閉
難度 2/5 1-3 小時 新手友好度 84/100
mllam/neural-lam#848 ·
維護者通常 2 天內回覆
-
infra
難度 1/5 1 小時以內 新手友好度 88/100
fab-geocommuns/RNB-coeur#1166 ·
維護者通常 1 天內回覆
-
難度 2/5 1-3 小時 新手友好度 78/100
deepspeedai/DeepSpeed#8951 ·
維護者通常 2 天內回覆