Repository navigation
train_lcm_distill_lora_sdxl.py: intermediate validation logs the un-adapted base model; saved LoRA mixes key conventions and drops all attention adapters on load #14548
Description
Describe the bug
examples/consistency_distillation/train_lcm_distill_lora_sdxl.py has two related silent failures, both traced to routing the PEFT state dict through convert_state_dict_to_diffusers before load_lora_weights / save_lora_weights.
1. Intermediate validation never applies the LoRA — it logs the un-adapted base model for the whole run.
log_validation (non-final branch) does:
unet = accelerator.unwrap_model(unet)
state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))
pipeline.load_lora_weights(state_dict)
pipeline.fuse_lora()The extracted keys have no unet. prefix (e.g. down_blocks.0.resnets.0.conv1.lora_A.weight), so StableDiffusionXLPipeline.load_lora_weights matches nothing, prints the "No LoRA keys associated to UNet2DConditionModel found with the prefix='unet'. This is safe to ignore ..." warning, and generates from the bare base model with LCMScheduler at 4 steps.
Evidence from a full 3000-step run with README-default arguments: across all 60 validation points (--validation_steps 50), the 960 logged W&B media files deduplicate to exactly 16 distinct PNGs (4 prompts x 4 images). The image logged at step 49 and the one at step 2999 are byte-identical (same SHA-256: 00e85a5f4dea3dc4f703...). Anyone tuning from the intermediate validation panel is looking at the frozen base model and may wrongly conclude the distillation is not learning.
2. The saved LoRA mixes two key conventions, and every attention adapter is silently dropped on load.
The final save goes through the same converter:
unet_lora_state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))
StableDiffusionXLPipeline.save_lora_weights(args.output_dir, unet_lora_layers=unet_lora_state_dict)convert_state_dict_to_diffusers renames only the attention adapters to the legacy .lora.down.weight / .lora.up.weight convention and passes every other module through in PEFT naming. Our trained pytorch_lora_weights.safetensors therefore holds 1576 keys in two conventions:
| key family | count | example |
|---|---|---|
| attention (legacy naming) | 1120 | unet.down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_k.lora.down.weight |
| conv / FF / proj (PEFT naming) | 456 | unet.down_blocks.0.downsamplers.0.conv.lora_A.weight |
load_lora_weights on this file applies the PEFT-named keys and silently drops all 1120 legacy attention keys. Measured by fusing after a +1.0 perturbation of every lora_B:
conv probe (down_blocks.0.resnets.0.conv1): max |delta| = 0.37
attn probe (down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_q): max |delta| = 0.0
Renaming the 1120 attention keys back to .lora_A/.lora_B and loading the same file makes both probes non-zero (with the actually-trained weights: attn 4.5e-3, conv 4.7e-3), and 4-step inference output changes substantially — so every downstream consumer of the exported file, including the script's own final "test" validation, has been running a LoRA with all attention adapters missing.
Reproduction
No training needed; mirrors the script's exact extract/convert/load lines:
import torch, warnings
warnings.filterwarnings("ignore")
from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel
from diffusers.utils import convert_state_dict_to_diffusers
from peft import LoraConfig, get_peft_model_state_dict
MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
unet = UNet2DConditionModel.from_pretrained(MODEL, subfolder="unet", torch_dtype=torch.float16)
targets = ["to_q","to_k","to_v","to_out.0","proj_in","proj_out","ff.net.0.proj","ff.net.2",
"conv1","conv2","conv_shortcut","downsamplers.0.conv","upsamplers.0.conv","time_emb_proj"]
unet.add_adapter(LoraConfig(r=64, lora_alpha=64, lora_dropout=0.0, target_modules=targets))
with torch.no_grad():
for name, p in unet.named_parameters():
if "lora_B" in name:
p.add_(1.0) # make any successful application unmissable
state = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))
print("sample key:", next(iter(state))) # no `unet.` prefix
pipe = StableDiffusionXLPipeline.from_pretrained(MODEL, torch_dtype=torch.float16)
def probe(pipeline, path):
sd = pipeline.unet.state_dict()
for key in (path + ".weight", path + ".base_layer.weight"):
if key in sd:
return sd[key].clone()
ATTN = "down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_q"
CONV = "down_blocks.0.resnets.0.conv1"
before_attn, before_conv = probe(pipe, ATTN), probe(pipe, CONV)
# Bug 1: the intermediate-validation load path applies nothing at all.
pipe.load_lora_weights(state)
pipe.fuse_lora()
print("bug 1 - peft_config on unet:", hasattr(pipe.unet, "peft_config"))
print("bug 1 - attn delta:", (probe(pipe, ATTN).float() - before_attn.float()).abs().max().item())
print("bug 1 - conv delta:", (probe(pipe, CONV).float() - before_conv.float()).abs().max().item())
# Bug 2: the save->load round trip (the script's final path) drops attention.
import tempfile
pipe2 = StableDiffusionXLPipeline.from_pretrained(MODEL, torch_dtype=torch.float16)
b_attn, b_conv = probe(pipe2, ATTN), probe(pipe2, CONV)
with tempfile.TemporaryDirectory() as tmp:
StableDiffusionXLPipeline.save_lora_weights(tmp, unet_lora_layers=state)
pipe2.load_lora_weights(tmp)
pipe2.fuse_lora()
print("bug 2 - attn delta:", (probe(pipe2, ATTN).float() - b_attn.float()).abs().max().item())
print("bug 2 - conv delta:", (probe(pipe2, CONV).float() - b_conv.float()).abs().max().item())Logs
No LoRA keys associated to UNet2DConditionModel found with the prefix='unet'. This is safe to ignore if LoRA state dict didn't originally have any UNet2DConditionModel related params. You can also try specifying `prefix=None` to resolve the warning. Otherwise, open an issue if you think it's unexpected: https://github.com/huggingface/diffusers/issues/new
bug 1 - peft_config on unet: False
bug 1 - attn delta: 0.0
bug 1 - conv delta: 0.0
bug 2 - attn delta: 0.0
bug 2 - conv delta: 0.3736572265625
Notes on fixes
- Skipping
convert_state_dict_to_diffusersand keeping PEFT-native names throughsave_lora_weights/load_lora_weightslooks like the clean fix for the script; prefixing the intermediate-validation dict withunet.alone is not sufficient (the legacy-renamed attention keys are still dropped; only conv/FF/proj come back). - More broadly,
load_lora_weightsdropping an entire key family with a "safe to ignore" warning is what let both failures ship silently; a hard error (or at least a distinct warning) when a non-empty key family matches no module would have surfaced this immediately.
Happy to share the full 3000-step evidence (W&B media dedup listing, trained safetensors key census, before/after inference grids) if useful.
System Info
- diffusers 0.36.0 (script byte-identical to the v0.36.0 tag)
- peft 0.18.1
- torch 2.9.1+cu129
- Python 3.12, Linux x86_64, single NVIDIA H20
Activity
primorLee commentedon Aug 21, 2026
I checked current main at 2f7e015 and the same conversion is still used in all three relevant paths: intermediate validation, the Accelerate checkpoint save hook, and final export. I would like to prepare a focused fix covering all three sites so checkpoint resumes do not preserve the same mixed key convention.\n\nProposed scope:\n- pass the PEFT-native state dict from get_peft_model_state_dict() directly to load_lora_weights()/save_lora_weights() in this script;\n- add a CPU regression test with dummy LoRA modules that verifies both attention and convolution target keys survive the validation and save/load paths;\n- leave generic loader warning behavior and unrelated training scripts unchanged.\n\nI searched open and recently closed PRs and found no duplicate. This will be an AI-assisted contribution that I will verify and maintain. Per the repository contribution policy, I will wait for a maintainer acknowledgement before opening a PR. Does this scope look right?
Skyyyy0920 commentedon Aug 23, 2026
Thanks for picking this up — the scope looks right to me, and it matches what we measured. Three data points from our verification that may save you a round trip:
Third site confirmed. In v0.36.0 the same conversion appears at intermediate validation (L124), the Accelerate checkpoint save hook (L893–894), and final export (L1445–1446) — so covering all three is the right scope; otherwise a resumed run re-exports the mixed convention exactly as you say.
One caveat for the intermediate-validation path: passing a state dict directly to
pipeline.load_lora_weights()needs theunet.prefix on every key — component attribution is prefix-based, and a bare dict (even PEFT-native naming) matches nothing, which is precisely bug 1. What we verified empirically to load fully is PEFT-native naming plus theunet.prefix (our trained artifact with the 1120 attention keys renamed tolora_A/lora_Bloads with non-zero deltas on both attention and conv probes). The most robust spelling for validation may be a temp-dir round trip throughsave_lora_weights/load_lora_weights, which reuses the exact code path the final export takes.Existing artifacts are recoverable. Files already exported by this script can be fixed with a one-line key remap (
.lora.down.weight→.lora_A.weight,.lora.up.weight→.lora_B.weight); may be worth a note in the fix description for users with trained checkpoints.Happy to run your patch against our full 3000-step trained artifact and the byte-identical-validation-images scenario once a PR is up.
(edited: an escaping mishap stripped the inline code spans from the first version of this comment)
Sonica-B commentedon Oct 25, 2026
Hi @primorLee, are you still planning to send the fix for this?
I reproduced it on
main(d42ebf9) with a tiny SDXL config on CPU:Passing the PEFT state dict at all three sites (validation, save hook, final export) fixes it: 354/354 keys in PEFT format, 177/177 modules loaded, and resume works from both old and new checkpoints. I have a tested branch with regression checks in
test_lcm_lora.py. Happy to review your PR, hand the branch over, or open it myself if you've moved on.@sayakpaul, does this scope look right to you?