问题诊断 (The Diagnosis)
if accelerator.is_main_process:
transformer = accelerator.unwrap_model(transformer)
transformer = transformer.to(torch.float32)
transformer_lora_state_dict = convert_state_dict_to_diffusers(
get_peft_model_state_dict(transformer)
)
save_file(new_state_dict, save_path)
罪魁祸首就是这行代码:transformer = transformer.to(torch.float32)
发生了什么?
- 训练时的精度: 在整个训练过程中,你的模型是在
fp16 或者 bf16 (由 mixed_precision 参数决定)的混合精度下进行的。这是一个为了速度和显存优化的标准操作。
- 保存时的“好心办坏事”: 在保存 LoRA 权重之前,为了确保最高的精度和兼容性(这是一个老的、安全但过时的习惯),你的脚本强制将整个 Transformer 模型转换回了
fp32 (32位单精度) 格式。
- 结果:
get_peft_model_state_dict 从这个 fp32 模型中提取出来的 LoRA 权重(lora_A 和 lora_B 矩阵),自然也就是 fp32 格式的。最终,你保存到 .safetensors 文件里的是一个 fp32 精度的 LoRA。
为什么 fp32 的 LoRA 在 fp8 上会失效?
这就像是试图把一张未经压缩的、巨大的 RAW 格式照片,直接用一个只为手机 HEIC 格式设计的简单工具去强行压缩。
- FP8 是“极限压缩”:
fp8 (8位浮点数) 是一种极其激进的量化格式,它对权重的数据范围和分布非常敏感。它被设计用来处理那些已经是 fp16 或 bf16 的、“正常范围”的权重。
- FP32 是“高动态范围”:
fp32 的数值范围比 fp16 大得多。一个在 fp32 下看起来很正常的权重,在 fp16 的世界里可能已经是一个需要特殊处理的“极大值”或“极小值”了。
- 转换失败: 当一个为
fp16 -> fp8 设计的推理引擎,突然拿到一个 fp32 的 LoRA 权重时,它在量化过程中很容易出现“溢出” (Overflow) 或 **“下溢” (Underflow)**,导致权重信息大量丢失。结果就是模型输出的图像是黑的、花的,或者完全是噪音。
手术方案:在保存时维持训练精度 (The Surgical Fix)
解决方案很简单:我们只需要在保存时,将 LoRA 权重转换回训练时使用的 fp16 或 bf16 精度,而不是粗暴地转成 fp32。
请将你的脚本最后那个“保存权重”的部分,替换成下面这个“精准转换”的版本:
if accelerator.is_main_process:
transformer = accelerator.unwrap_model(transformer)
transformer_lora_state_dict = get_peft_model_state_dict(transformer)
final_state_dict = {}
logger.info(f"Converting LoRA weights to {weight_dtype} before saving...")
for k, v in transformer_lora_state_dict.items():
final_state_dict[k] = v.to(dtype=weight_dtype)
diffusers_state_dict = convert_state_dict_to_diffusers(final_state_dict)
new_state_dict_for_saving = {}
for k, v in diffusers_state_dict.items():
new_state_dict_for_saving[f"transformer.{k}"] = v
save_path = os.path.join(args.output_dir, "pytorch_lora_weights.safetensors")
save_file(new_state_dict_for_saving, save_path)
logger.info(f"Saved LoRA weights in {weight_dtype} to {save_path}")
accelerator.end_training()
总结:
|
旧方法 (FP32 保存) |
新方法 (精准保存) |
| 保存前转换 |
model.to(torch.float32) |
遍历 state_dict, tensor.to(weight_dtype) |
| 保存的精度 |
强制 fp32 |
与训练精度一致 (fp16/bf16) |
| FP8 兼容性 |
差,易出错 |
好,完美兼容 |
| 文件大小 |
较大 |
较小 (一半) |
把这个修改应用到你的 train_zimage_lora.py 脚本里,重新训练出来的 LoRA,去测试,这次在 fp8 环境下绝对能正常工作了!
这是一个非常高质量的反馈。👍
新的lora需要等我这个25小时的模型训练做完,请等待