ysharma/image21-turbo-trajectories

Qwen-Image-2.1-Turbo 4-step distillation — trajectories & project state

Status: paused at gate G3-FAIL (quality gate not passed at smoke scale). Budget exhausted.

Goal: train a rank-64 LoRA on Qwen/Qwen-Image-2.1-Turbo so that, with the LoRA on, 4-step generation matches Turbo's own 8-step output from the same noise.

Contents

• smoke/adapter/pytorch_lora_weights.safetensors — the last smoke-trained adapter (128 prompts, 1000 steps, lr 3e-4, endpoint-weighted loss — diverged; NOT a usable deliverable, kept for debugging).
• smoke/preflight_test_1kb.bin — G1 push-preflight artifact.
• extract_prompts.py — prompt extractor (pyarrow column projection on LucasFang/FLUX-Reason-6M, Apache-2.0; images never downloaded). Produces prompts_train.jsonl (~1826 prompts, 28% caption_text) and prompts_eval.jsonl (64 held out: 32 general + 32 quoted-text with quoted_strings for OCR eval).

What was verified (gates)

Gate	Result
G0 plumbing (tiny-random, CPU)	PASS 19/19 — schedule, Euler loop, trajectory format, LoRA train step, save/reload
G1 push preflight	PASS — create_repo(exist_ok=True) + upload_folder + download verify
G2 sampler parity (real Turbo, bf16, L40S)	PASS — own 8-step Euler loop reproduces pipe(...) bit-exact (max diff 0.0); scheduler.sigmas == passed + [0]; num_inference_steps ignored; scheduler uses sigmas verbatim
G3 smoke (quality gate)	FAIL — student@4 LPIPS-to-teacher worse than naive@4 at every scale tested

G3 evidence (all at 512 px, A100/L40S, same noise, LPIPS to teacher@8)

Run	Prompts	Steps	lr	Loss weighting	Student LPIPS	Naive LPIPS
v16	16	100	1e-4	uniform velocity MSE	0.2946	0.2698
v18	16	300	3e-4	uniform	0.3228	0.2698
v19	16	300	3e-4	endpoint (Δσ²)	0.2885	0.2698
scaled	128	1000	3e-4	endpoint	0.8237 (training diverged)	0.2693

Velocity diagnostic (unweighted MSE of predicted velocity vs stored v* on training pairs, adapter on vs off): 0.0575 vs 0.0254 (16-prompt, endpoint-weighted); 2.199 vs 0.025 (128-prompt, diverged). The adapter never improves velocities, even on its own training pairs, and gets worse with more training.

Root causes found and fixed along the way

1. Embeddings must be computed under torch.no_grad(). QwenImage21Pipeline.encode_prompt is not decorated; calling it with grad enabled attaches the text encoder's graph to the embeddings, and every training backward then traverses that shared graph. This caused both the cuda:0 and cpu device mismatch (text encoder offloaded to CPU after encoding) and the "backward through the graph a second time" error. Not a checkpointing problem.
2. AutoencoderKLQwenImage21.decode returns RGBA (4 channels) — slice [:, :3, 0] for RGB.
3. The bf16 GPU decode of this video VAE balloons to ~76 GB of transients at 512 px — decode on CPU instead.
4. Uniform velocity-MSE lets the k=0 pair (Δσ=0.046, huge velocities) dominate while k=3 (Δσ=0.70, the step that lands the final image) gets ~0.4% of the gradient. Endpoint weighting (Δσ²) helped (0.3228 → 0.2885) but did not fix the gate.
5. The upstream DreamBooth trainer (examples/dreambooth/train_dreambooth_lora_qwenimage21.py) trains correctly on the same commit/GPU with the same LoRA config, including with synthetic batches — the harness that works is: standalone QwenImage21Transformer2DModel.from_pretrained, LoRA added before accelerator.prepare, accelerator.accumulate + accelerator.backward, materialized prompt mask.
6. Canonical LoRA publish path: get_peft_model_state_dict → pipe.save_lora_weights(dir, transformer_lora_layers=..., **_collate_lora_metadata({"transformer": transformer})) → pipe.load_lora_weights(dir). Metadata is mandatory — without it the loader rebuilds the config with default alpha (64 instead of 128 → half-strength adapter).
7. The tiny-random repo's scheduler ships use_dynamic_shifting=true, shift_terminal=0.02; the real Turbo uses false/null and passes sigmas verbatim.

Measured timings (A100, 512 px, checkpointing on)

• Teacher trajectory (8-step Euler loop): 0.84 s + 1.9 s/prompt text encoding
• 4-step generation: 0.42 s · Training step: 0.49 s (rank 64, 167.8M LoRA params, fp32 params, AdamW)
• VRAM: 16 GB training peak → 1.3 GB after offload

Next steps (not started)

1. Training-stability sweep at sane lr (1e-4 / 5e-5) at 128-prompt scale — the 3e-4 run diverged (loss spikes to 1.2; adapter velocity MSE 2.2).
2. If a stable run still degrades velocities: the on-policy phase from the original plan (roll the student to σ_2k, teacher two steps from there, train on that) — offline pairs alone appear insufficient at this model/recipe scale.
3. Full plan at 512 px (if the gate passes): trajectories 1.4 h ($3), training 2 epochs 2 h ($5), eval ~$3. Eval at 1 MP with the 512-px LoRA is unverified.

Base model licence: Qwen Research Licence (non-commercial) — derivatives need the licence text, a change notice, and "Built with Qwen", and must not use "Qwen" as the main repo name.