"""Few-step distillation of the streaming causal CWM by distribution matching (DMD) on the student's own long rollouts (self forcing), as in CausVid / Self Forcing / Self Forcing++. torchrun ... -m rtcwm.train.distill --out DIR --student-lora DF_LORA [--steps 4] [--iters 200] Three LoRAs on one frozen, FSDP-sharded H3 (rtcwm.model.loading.LoRAStack): student the streaming DF model (rtcwm.train.train), sampled chunk by chunk with `--steps` x0-steps teacher the released CWM LoRA, frozen: the real score critic initialised from the released CWM LoRA, trained to denoise the student's samples: the fake score The teacher and the critic run the released model exactly as released: bidirectional over a released 37-latent window (124 frames = 7 VAE clips + the first 2 latents of the 8th, a 5n+2 window), with that window's whole proxy. The VAE is causal inside a clip, so 37 consecutive latents of a chunk stream are exactly the released encoding of those frames (ops/dbg/tail_latents.py) -- no layout mismatch. Each iteration (rollout length R shared by all ranks, drawn from --min-chunks..--max-chunks): 1. every rank draws a clip and a start frame m and streams R chunks from it (rtcwm.model.sampler.rollout, the inference procedure, with the KV cache carried across chunks). As in Self Forcing (mingjia's port, wt-sf-mingjia/sf/self_forcing_model.py), every chunk stops at its own random x0-step (shared by all ranks) and that step's clean prediction is the chunk -- and the history of the next ones. The last --grad-chunks chunks of the teacher window keep the graph of their stopping step, their history entering detached (memory: a graph keeps what it attended to -- its span's proxy keys/values and the history chunks that have since left the window -- alive until the backward pass; four fit). Earlier chunks keep no graph: the window slides, and a graph would keep its own shifted copy of the history's keys/values alive until the backward pass. A history chunk enters the cache as the sampler caches it (rtcwm.model.sampler.HISTORY_FROM_LAST_STEP): noised to the last sampling step's level, the first chunk clean (--rollout-history last-step), or clean. 2. teacher window = chunks R-8 .. R-2 and the first 2 latents of chunk R-1 (Picture 1 = the window's own first frame from the rollout, as the released CWM is used; the start frame's caption; the window's proxy; silent audio), noised to one random level: loss = 0.5 || x - sg(x - (x0_fake - x0_real) / mean|x - x0_real|) ||^2 over the window (DMD gradient) 3. critic: --critic-per-gen flow-matching updates (generator : critic = 1 : 7 by default), each on a random teacher window of the same detached rollout (R-7 of them) with fresh noise -- the rollout, the expensive part, is shared. With --teacher-kind causal the teacher is a streaming DF checkpoint (--teacher-lora, frozen) and the critic starts from it: the teacher window is then the student's own layout -- one 6-chunk span of the rollout (the last one for the generator, a random one for the critic) with the stream's start frame and caption and the span's causal proxy (rule E, rtcwm.model.masks) -- every chunk of it noised to the same level. Long rollouts put the student in the drifted states it reaches at test time, and the teacher's score pulls them back -- training on its own history removes the exposure bias of diffusion forcing. """ from __future__ import annotations import argparse import json import os import random import time from pathlib import Path import torch import torch.distributed as dist from cwm_h3_inference.constants import VIDEO_LATENT_SHAPE, VIDEO_SHIFT from .. import paths from ..model import loading from ..model.layout import CHUNK_LATENTS, SPAN_CHUNKS from ..model.objectives import df_forward, make_window, shift from ..model.packing import Window, run_blocks, video_output from ..model.sampler import rollout, schedule from .dataset import list_samples, load_window, silence, start_frames, stream_chunks from .train import node_group STUDENT_SHIFT = 5.0 # few-step schedule: sigma 1, .94, .83, .63, 0 for 4 steps TEACHER_LATENTS = VIDEO_LATENT_SHAPE[2] # 37 = 7 chunks + 2 TEACHER_CHUNKS = TEACHER_LATENTS // CHUNK_LATENTS + 1 # chunks a teacher window touches: 8 def parse(): ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--out", type=Path, required=True) ap.add_argument("--student-lora", type=Path, required=True, help="streaming DF checkpoint to distil") ap.add_argument("--teacher-lora", type=Path, default=paths.CWM_LORA) ap.add_argument("--critic-lora", type=Path, default=paths.CWM_LORA) ap.add_argument("--steps", type=int, default=4, help="student sampling steps") ap.add_argument("--iters", type=int, default=200) ap.add_argument("--min-chunks", type=int, default=TEACHER_CHUNKS, help="shortest rollout (chunks)") ap.add_argument("--max-chunks", type=int, default=14, help="longest rollout (chunks; 14 = 10 s)") ap.add_argument("--kv-reuse", choices=["on", "off"], default="on", help="rollout: carry the history's keys/values") ap.add_argument("--rollout-history", choices=["last-step", "clean"], default="last-step", help="rollout history: last-step = noised to the last sampling step's level, the history the sampler " "caches by default (HISTORY_FROM_LAST_STEP); clean = one more clean forward") ap.add_argument("--gen-lr", type=float, default=2e-6) ap.add_argument("--critic-lr", type=float, default=4e-6) ap.add_argument("--critic-per-gen", type=int, default=7, help="critic updates per generator update") ap.add_argument("--grad-chunks", type=int, default=4, help="last teacher-window chunks that keep their graph") ap.add_argument("--teacher-kind", choices=["released", "causal"], default="released", help="released: the released CWM on 37-latent bidirectional windows; causal: a streaming DF " "checkpoint (--teacher-lora, also the critic's start) on the student's 6-chunk spans") ap.add_argument("--teacher-picture1", choices=["window", "start"], default="window", help="teacher / critic Picture 1: the window's first frame (rollout) or the stream's start frame") ap.add_argument("--warmup", type=int, default=10) ap.add_argument("--save-every", type=int, default=20) ap.add_argument("--val-clips", type=int, default=12) ap.add_argument("--picture1", choices=["native", "2048"], default="native") ap.add_argument("--seed", type=int, default=0) return ap.parse_args() def park(opts, device): """Move the optimizers' moments to `device`: the CPU while the rollout's graphs fill the GPU (about 10 GB for student + critic), back before a step.""" for opt in opts: for state in opt.state.values(): for key, val in state.items(): if torch.is_tensor(val) and val.dim() > 0: state[key] = val.to(device, non_blocking=True) class Reducer: """Asynchronous gradient all-reduce for replicated LoRA weights, overlapped with backward.""" def __init__(self, params): self.handles = [] for p in params: p.register_post_accumulate_grad_hook( lambda q: self.handles.append(dist.all_reduce(q.grad, op=dist.ReduceOp.AVG, async_op=True))) def wait(self): for h in self.handles: h.wait() self.handles.clear() class Starts: """(clip, start frame m) pairs whose rollout of n chunks stays inside the training spans (no held-out clip, no QC-flagged span) and whose teacher window fits the stream.""" def __init__(self, samples, res: str): spans = {} for c, j in samples: spans.setdefault(c, set()).add(j) self.spans, self.frames = spans, {c: start_frames(c, res) for c in spans} self.n_stream = {c: stream_chunks(c) for c in spans} def draw(self, rng: random.Random, n: int): while True: clip = rng.choice(sorted(self.spans)) ok = [m for m in self.frames[clip] if m + n <= self.n_stream[clip] and all(m + j in self.spans[clip] for j in range(0, n - SPAN_CHUNKS + 1))] if ok: return clip, rng.choice(ok) def released_velocity(model, win: Window, latents, model_t: float, train: bool): """The released (bidirectional) H3 on a released window: latents [1,24,37,H,W] at model time model_t.""" hv = win.video_rows(model, latents) t_rows = torch.cat([win.cond_times(), torch.full((hv.shape[1],), float(model_t))]) tags = torch.cat([win.tags, torch.zeros(hv.shape[1], dtype=torch.int64)]) hidden, temb, inv = run_blocks(model, torch.cat([win.rows, hv], 1), t_rows, tags, torch.cat([win.ca_pos, win.target_pos]), train=train) return video_output(model, hidden, temb, inv, win.rows.shape[1], latents.shape[2]) def causal_velocity(model, win: Window, latents, model_t: float, train: bool): """A streaming DF model on one span (rtcwm.model.objectives.df_forward), every chunk at model time model_t.""" return df_forward(model, win, latents, [float(model_t)] * SPAN_CHUNKS, {}, train=train) def noised(x, gen): """x at a random teacher-schedule level -> (x_t, n, sigma).""" sigma = float(shift(torch.rand(1, generator=gen).clamp(0.02, 0.98), VIDEO_SHIFT)) n = torch.randn(x.shape, generator=gen).to(x.device) return (1 - sigma) * x + sigma * n, n, sigma def dmd_loss(stack, model, win, x, gen, velocity=released_velocity): """DMD generator loss on teacher window x (the chunks that kept their graph carry the gradient).""" xt, _, sigma = noised(x.detach(), gen) with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): stack.use("teacher") x0_real = xt + sigma * velocity(model, win, xt, 1.0 - sigma, train=False).float() stack.use("critic") x0_fake = xt + sigma * velocity(model, win, xt, 1.0 - sigma, train=False).float() p = x.detach() - x0_real grad = (x0_fake - x0_real) / p.abs().mean().clamp_min(1e-6) stack.use("student") # the graph of x is the student's return 0.5 * torch.nn.functional.mse_loss(x.float(), (x.detach() - grad).float()), float(grad.abs().mean()) def critic_loss(stack, model, win, x, gen, velocity=released_velocity): """Flow-matching loss of the critic on a detached window.""" xt, n, sigma = noised(x.detach(), gen) stack.use("critic") with torch.autocast("cuda", dtype=torch.bfloat16): v = velocity(model, win, xt, 1.0 - sigma, train=True) return torch.nn.functional.mse_loss(v.float(), (x.detach() - n).float()) def resume_point(a): """(iterations done, student LoRA, critic LoRA): the latest student/critic pair saved in --out (a requeued or resubmitted run continues from it -- lowprio jobs are preempted and requeued), else the initial ones.""" done = sorted(int(p.stem.split("iter")[1]) for p in a.out.glob("student_iter*.safetensors") if (a.out / p.name.replace("student", "critic")).exists()) if not done: return 0, a.student_lora, a.critic_lora n = done[-1] return n, a.out / f"student_iter{n:05d}.safetensors", a.out / f"critic_iter{n:05d}.safetensors" def main(): a = parse() causal = a.teacher_kind == "causal" velocity = causal_velocity if causal else released_velocity w_chunks = SPAN_CHUNKS if causal else TEACHER_CHUNKS # rollout chunks one teacher window covers assert a.grad_chunks <= w_chunks, (a.grad_chunks, w_chunks) assert TEACHER_CHUNKS <= a.min_chunks <= a.max_chunks, (a.min_chunks, a.max_chunks) dist.init_process_group("nccl") rank, world = dist.get_rank(), dist.get_world_size() local = int(os.environ.get("LOCAL_RANK", 0)) torch.cuda.set_device(local) dev = torch.device("cuda", local) log = (lambda *m: print(f"[{time.strftime('%H:%M:%S')}]", *m, flush=True)) if rank == 0 else (lambda *m: None) group = node_group(rank, world) # NCCL sets up a communicator's buffers (outside the PyTorch allocator) at its first collective. The # first cross-node one would be the generator's gradient all-reduce, with the GPU at its fullest: make # it now, and keep a few GB out of PyTorch's reach for NCCL. torch.cuda.set_per_process_memory_fraction(0.96, dev) dist.all_reduce(torch.zeros(1 << 24, device=dev)) train, val, _ = list_samples(a.val_clips, a.seed, a.picture1) starts = Starts(train, a.picture1) log(f"distill world={world} train_spans={len(train)} clips={len(starts.spans)} val={len(val)} iters={a.iters} " f"student_steps={a.steps} rollout={a.min_chunks}..{a.max_chunks} chunks") if rank == 0: a.out.mkdir(parents=True, exist_ok=True) (a.out / "val_samples.json").write_text(json.dumps(val)) (a.out / "args.json").write_text(json.dumps({k: str(v) for k, v in vars(a).items()}, indent=1)) start, student_lora, critic_lora = resume_point(a) if start: log(f"resuming after iter {start}: {student_lora.name}, {critic_lora.name} (optimizer moments start afresh)") tr = loading.load_base() stack = loading.LoRAStack(tr, dict(student=(student_lora, True), critic=(critic_lora, True), teacher=(a.teacher_lora, False))) loading.shard_blocks(tr, group, dev) stack.to(dev) tr.train() gen_params, critic_params = stack.params("student"), stack.params("critic") reducer = Reducer(gen_params + critic_params) opt_g = torch.optim.AdamW(gen_params, lr=a.gen_lr, betas=(0.0, 0.999), weight_decay=0.01) opt_c = torch.optim.AdamW(critic_params, lr=a.critic_lr, betas=(0.0, 0.999), weight_decay=0.01) warm = lambda i: min(1.0, (i + 1) / a.warmup) # constant after warmup sched_g, sched_c = (torch.optim.lr_scheduler.LambdaLR(o, warm) for o in (opt_g, opt_c)) sigmas = schedule(a.steps, dev, STUDENT_SHIFT) history_noise = float(sigmas[-2]) if a.rollout_history == "last-step" else 0.0 audio = silence() gen = torch.Generator().manual_seed(a.seed * 1000 + rank) rng = random.Random(a.seed * 1000 + rank) log(f"student/critic trainable={sum(p.numel() for p in gen_params) / 1e6:.0f}M each, " f"schedule={[round(float(s), 3) for s in sigmas]}") def step(opt, sched, params): reducer.wait() gn = torch.nn.utils.clip_grad_norm_(params, 1.0) opt.step(); sched.step(); opt.zero_grad(set_to_none=True) return float(gn) t0 = time.time() for it in range(start, a.iters): # the rollout length and the gradient step fix the sequence of FSDP gathers: equal on every rank shared = random.Random(a.seed * 100003 + it) n = shared.randint(a.min_chunks, a.max_chunks) exit_steps = [shared.randrange(a.steps) for _ in range(n)] grad_chunks = list(range(n - a.grad_chunks, n)) clip, m = starts.draw(rng, n) t_it, chunk_times = time.time(), [] park((opt_g, opt_c), "cpu") stack.use("student") with torch.autocast("cuda", dtype=torch.bfloat16): chunks = rollout(tr, lambda j: load_window(clip, m + j, audio, start=m, with_target=False, res=a.picture1), n, sigmas, gen, dev, exit_steps, grad_chunks, kv_reuse=a.kv_reuse == "on", history_noise=history_noise, on_chunk=(lambda k, s: chunk_times.append(s)) if it < 3 else None) t_roll = time.time() - t_it if chunk_times: log(f"iter {it + 1} rollout seconds per chunk: {[round(s, 1) for s in chunk_times]}") def teacher_window(w): """Rollout chunks w .. w+7 as a released window (37 latents) with its conditions. With --teacher-picture1 window the window's Picture 1 is its own first frame -- the first latent of the rollout's chunk w, which the per-clip VAE makes the frame's own encoding to within a few percent -- as the released CWM is used (Picture 1 = the window's first frame): the scores then ask for a continuation of the scene the rollout is in, not of the start frame seen long ago. Causal teacher: rollout chunks w .. w+5 as the student's span m+w (start frame, its caption, causal proxy).""" if causal: sample = load_window(clip, m + w, audio, start=m, with_target=False, res=a.picture1) with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): win = make_window(tr, sample, dev, gen) return win, torch.cat(chunks[w:w + SPAN_CHUNKS], 2) sample = load_window(clip, m + w, audio, start=m, with_target=False, res=a.picture1, n_latents=TEACHER_LATENTS) if a.teacher_picture1 == "window": sample["anchor"] = chunks[w][:, :, :1].detach().cpu() with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): win = make_window(tr, sample, dev, gen) return win, torch.cat(chunks[w:w + TEACHER_CHUNKS], 2)[:, :, :TEACHER_LATENTS] win, x = teacher_window(n - w_chunks) lg, gmag = dmd_loss(stack, tr, win, x, gen, velocity) lg.backward() park((opt_g, opt_c), dev) gn_g = step(opt_g, sched_g, gen_params) del win, x chunks = [c.detach() for c in chunks] lcs = [] for _ in range(a.critic_per_gen): win, x = teacher_window(rng.randrange(n - w_chunks + 1)) lc = critic_loss(stack, tr, win, x, gen, velocity) lc.backward() gn_c = step(opt_c, sched_c, critic_params) lcs.append(float(lc)) stats = torch.tensor([float(lg), sum(lcs) / len(lcs), gmag], device=dev) dist.all_reduce(stats, op=dist.ReduceOp.AVG) if rank == 0: el = time.time() - t0 log(f"iter {it + 1}/{a.iters} R={n} grad_chunks={len(grad_chunks)} dmd={stats[0]:.4f} critic={stats[1]:.4f} " f"|grad|={stats[2]:.3f} gn_g={gn_g:.2e} gn_c={gn_c:.2e} mem={torch.cuda.max_memory_allocated() / 2**30:.0f}G " f"rollout {t_roll:.0f}s, iter {time.time() - t_it:.0f}s, eta={(a.iters - it - 1) * el / (it + 1 - start) / 3600:.2f}h") with open(a.out / "train_log.jsonl", "a") as fh: fh.write(json.dumps(dict(iter=it + 1, chunks=n, exit_steps=exit_steps, grad_chunks=grad_chunks, dmd=float(stats[0]), critic=float(stats[1]), grad=float(stats[2]), gn_g=gn_g, gn_c=gn_c, rollout_sec=t_roll, sec=el)) + "\n") if (it + 1) % a.save_every == 0 or it + 1 == a.iters: for name in ("student", "critic"): p = a.out / f"{name}_iter{it + 1:05d}.safetensors" stack.nets[name].save_weights(str(p), torch.bfloat16, {"base": "cwm", "iter": str(it + 1), "objective": f"dmd-{a.steps}step"}) log("saved", a.out / f"student_iter{it + 1:05d}.safetensors") dist.barrier() dist.destroy_process_group() if __name__ == "__main__": main()