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

[Skyyyy0920]

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_diffusers and keeping PEFT-native names through save_lora_weights / load_lora_weights looks like the clean fix for the script; prefixing the intermediate-validation dict with unet. alone is not sufficient (the legacy-renamed attention keys are still dropped; only conv/FF/proj come back).
• More broadly, load_lora_weights dropping 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

Who can help?

@sayakpaul @linoytsaban

[primorLee]

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?