wileewang/longlive-stepdistill-blocksize-loras

LongLive step-distillation LoRAs

4-step distillation acceleration LoRAs for Wan2.2-TI2V-5B. Two families live here.

---

Family A — recipe LoRAs (2026-08, current best)

These are the ones to use. They mount LoRA on self-attention + FFN only (cross-attention is deliberately left un-adapted) and train with timestep_shift 10 under autoregressive self-rollout.

file	mounted	layers	params	block size	AR blocks
recipe_selfffn_shift10_AR_block8.pt	self_attn + ffn	180	228,065,280	8	3
recipe_selfffn_shift10_AR_block4.pt	self_attn + ffn	180	228,065,280	4	6

Both are rank 128 / alpha 128, 2000 steps, seed 1, trained on 24 latents with local_attn_size 24. The two configs differ in exactly one line (num_frame_per_block), so they are a clean block-size comparison.

Why self-attention + FFN only

Measured against the base model's 50-step output (DreamSim, 10 prompts, noise floor 0.0030 from three retrained seeds):

mounted	params	distance to 50-step target
self + ffn	228,065,280	0.2437
everything (self+cross+ffn)	322,437,120	0.2546
cross + ffn	228,065,280	0.2567
ffn only	133,693,440	0.2648
self + cross (no ffn)	188,743,680	0.2706

Dropping cross-attention is not a compromise for size — it scores better than mounting everything, at 29% fewer parameters. Adding cross-attention costs 3.6× the noise floor.

block 8 vs block 4

A genuine trade-off, both directions statistically significant (paired per prompt, n=10):

	block 8	block 4	winner	p
seam discontinuity at the two shared seams (frames 28, 60)	0.0105	0.0053	block 4, 10:0	0.002
drift from first to last frame	0.124	0.164	block 8, 9:1	0.021
whole-clip mean frame-to-frame distance	0.0066	0.0051	block 4, 3:7	0.344
pixel motion	4.08	3.25	—	0.109

Block 4 stitches each seam more smoothly but has 2.5× more seams: at 93 frames it has 5 (indices 12/28/44/60/76) versus block 8's 2 (28/60), and block 4's three extra seams score 0.0068 — 48% above its own interior, so they are real seams, not free. Block 8 accumulates less drift over the clip. Note block 4 also moves less (motion 3.25 vs 4.08), so part of "smoother" may be "less happening"; the seam result survives that caveat (10:0) but the whole-clip mean does not.

Inference cost is nearly identical: rendering the same 10 clips took 16m23s (block 8) vs 16m55s (block 4) — block 4 runs 6 sequential blocks instead of 3, but each is half as long.

---

Family B — block size × rollout ablation (2026-06)

Five LoRAs mounting all 300 linear layers (self + cross + ffn) at timestep_shift 5, varying block size and whether the student rolls out autoregressively.

file	block size	student rollout
block8_single.pt	8	single bidirectional (baseline)
block16_single.pt	16	single bidirectional
block16_AR8.pt	16	autoregressive (2×8 blocks)
block24_single.pt	24	single bidirectional
block24_AR8.pt	24	autoregressive (3×8 blocks)

Teacher + critic are bidirectional (no AR mask). Family A and Family B are not directly comparable — they differ in mount set, timestep_shift, and data pipeline.

---

Loading these — read this before you run anything

Format. Every .pt is a dict with keys generator_lora, critic_lora, step. Inference reads only generator_lora. Keys look like base_model.model.blocks.0.self_attn.q.lora_A.weight (PEFT layout).

The one thing that will silently break

Family A mounts 180 layers, not 300. If your inference code mounts the default full set, PEFT's set_peft_model_state_dict loads with strict=False — it discards the mismatched keys without raising, leaves lora_B at its zero init, and you render the un-adapted base model at 4 steps. That looks like a bad result, not like an error. We lost a week of conclusions to exactly this failure mode.

Two guards are worth porting:

1. config-vs-model: assert the number of mounted layers equals what the config declares (adapter.expected_layers), and stop the run if not.
2. weights-vs-mount: before loading, intersect the checkpoint's module names with the model's mounted module names, and refuse to render if the overlap is zero.

Inference config — Family A

model_kwargs:
  model_name: Wan2.2-TI2V-5B
  timestep_shift: 10.0        # 10, NOT the 5.0 used by Family B
  num_frame_per_block: 8      # 8 for ..._block8.pt, 4 for ..._block4.pt
  local_attn_size: 24
data:
  image_or_video_shape: [1, 24, 48, 44, 80]   # 24 latents -> 93 frames
inference:
  sampling_steps: 4
  guidance_scale: 1.0
  sink_size: 0
  multi_shot_rope_offset: 8
checkpoints:
  lora_ckpt: <path to the .pt>
adapter:
  type: lora
  rank: 128
  alpha: 128
  dropout: 0.0
  exclude_modules: [cross_attn]     # THIS is what makes it 180 layers
  expected_layers:
    generator: 180                  # gate: wrong mount stops the run
logging:
  seed: 0

Three values must match the file you load, or the result is silently wrong rather than loudly broken:

• timestep_shift: 10.0 — Family A was distilled at 10. Loading it at 5 puts the four denoising steps at different points on the noise schedule than the ones it was trained for.
• num_frame_per_block — 8 or 4, matching the filename. This sets the AR block size at inference, so it also decides where the seams land.
• exclude_modules: [cross_attn] — without it you mount 300 layers and 120 of them stay at zero init.

Frame count and seam positions

93 frames = 4 × (24 − 1) + 1. Latent 0 decodes to a single frame; every later latent to four. So latent i > 0 covers frames 4i−3 … 4i, and chunk k hands over between frames 4·block·k − 4 and 4·block·k − 3:

• block 8 → 3 chunks, seams between frames 28│29 and 60│61
• block 4 → 6 chunks, seams at 12│13, 28│29, 44│45, 60│61, 76│77

The naive latent i → frame 4i gives 32 and 64, which is wrong by 4 frames — enough to miss the seam entirely with a ±2 frame window, and enough to flip the ranking of arms.

If you mount with plain PEFT instead

import torch
from peft import LoraConfig, get_peft_model, set_peft_model_state_dict

sd = torch.load(path, map_location="cpu", weights_only=True)["generator_lora"]

targets = [n for n, m in transformer.named_modules()
           if isinstance(m, torch.nn.Linear)
           and any(t in n for t in ("self_attn", "ffn"))      # cross_attn excluded
           and "blocks." in n]
assert len(targets) == 180, f"expected 180 mounted layers, got {len(targets)}"

model = get_peft_model(transformer, LoraConfig(
    r=128, lora_alpha=128, lora_dropout=0.0, target_modules=targets))

mounted = {n.split(".lora_A")[0] for n, _ in model.named_parameters() if ".lora_A" in n}
in_ckpt = {k.split(".lora_A")[0].replace("base_model.model.", "")
           for k in sd if ".lora_A" in k}
assert mounted & in_ckpt, "checkpoint does not match the mounted adapter"

set_peft_model_state_dict(model, sd)

alpha == rank, and neither use_rslora nor use_dora is set, so the PEFT scale is alpha / r = 1.0.

---

Training settings, for reproduction

	Family A	Family B
base	Wan2.2-TI2V-5B	same
mounted	self_attn + ffn (180)	all (300)
rank / alpha	128 / 128	128 / 128
timestep_shift	10.0	5.0
training latents	24	8 / 16 / 24
local_attn_size	24	varies
AR (generator_is_causal)	true	varies
steps	2000	2000
lr (generator / critic)	1e-5 / 2e-6	same
batch size	1 × 8 GPUs	same
seed	1	0

The mount knob is applied to the critic as well as the generator (apply_to_critic: true), so an arm that fails to learn has had both its student and its critic capacity changed — do not attribute such a failure to the student alone.

Family A used seed: 1 because a seed of 0 gets replaced by an unrecorded random number at startup, and the data sampler seed was also fixed (it previously came from the wall clock and differed per rank, so shards overlapped). With 30000 training prompts and only 16000 samples drawn in 2000 steps, the sampler seed decides which prompts a run ever sees — Family B's runs therefore each saw a different subset, and comparisons across the two families carry that extra confound.