huggingface / huggingface/diffusers
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
- Vorherrschende Sprache
- Python
- Sterne
- 34.5k
- Forks
- 7.3k
- Ø Merge
- 3 T. 3 Std.
- Gemergte PRs (30 T.)
- 91
Beschreibung
### 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:
```python
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:
```python
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:
```python
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
Beitragsleitfaden
Rechercherichtung
Start with examples/consistency_distillation/train_lcm_distill_lora_sdxl.py, especially log_validation and the final LoRA save path, then run the provided reproduction to inspect the extracted keys and load/save behavior. Done means intermediate validation applies the adapters and a saved-and-reloaded LoRA retains both attention and non-attention adapters, verified by the stated probes.
Vom Indexierungsmodell aus dem Issue-Text verfasst.
Bewertung
- Tech-Stack
- python, pytorch
- Bereich
- machine-learning
- Issue-Typ
- Bug
- Schwierigkeit
- 4/5
- Geschätzter Aufwand
- 3-5 Tage
- Aktivitätsstatus
- Aktiv
- Klarheit
- Klar beschrieben
- Anfängerfreundlichkeit
- 48/100