diff --git a/data_pipeline/cfg_ablation.py b/data_pipeline/cfg_ablation.py new file mode 100644 index 000000000..f103733e4 --- /dev/null +++ b/data_pipeline/cfg_ablation.py @@ -0,0 +1,241 @@ +#!/usr/bin/env python3 +"""CFG-ablation eval on a trackwan checkpoint. + +Reproduces ``fastvideo/train/callbacks/track_validation.py::_sample`` denoise +formula EXACTLY, but with configurable ``w_text`` and ``w_motion`` so we can +isolate which CFG branch is causing artifacts. + +Runs each val sample under N configs; uploads all videos to one wandb run for +side-by-side visual comparison. + +Usage (single-GPU, on any held alloc): + python data_pipeline/cfg_ablation.py \ + --model-dir /mnt/lustre/vlm-s4duan/exports/synth_stage2_paperLR_ckpt400 \ + --yaml examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_paperLR.yaml \ + --val-parquet /mnt/lustre/vlm-s4duan/val_examples_mixed/combined_parquet_dataset \ + --wandb-run-name cfg_abl_paperLR_ckpt400 +""" +from __future__ import annotations + +import argparse +import glob +import os +import sys +import time +from pathlib import Path + +import imageio.v2 as imageio +import numpy as np +import pyarrow.parquet as pq +import torch + +sys.path.insert(0, os.path.join(os.path.dirname(__file__))) +import trackwan_infer as twi # noqa: E402 +from fastvideo.forward_context import set_forward_context # noqa: E402 +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( # noqa: E402 + FlowMatchEulerDiscreteScheduler, +) + + +def load_val_samples(parquet_dir: str, text_len: int) -> list: + from fastvideo.dataset.dataloader.schema import pyarrow_schema_i2v_track + from fastvideo.dataset.utils import collate_rows_from_parquet_schema + fs = sorted(glob.glob(os.path.join(parquet_dir, "**", "*.parquet"), recursive=True)) + if not fs: + raise FileNotFoundError(parquet_dir) + rows = [] + for f in fs: + rows.extend(pq.read_table(f).to_pylist()) + batch = collate_rows_from_parquet_schema(rows, pyarrow_schema_i2v_track, + text_padding_length=int(text_len), cfg_rate=0.0) + infos = batch.get("info_list") or [{} for _ in rows] + return [{ + "id": str(infos[i].get("id", f"clip{i:03d}")) if i < len(infos) else f"clip{i:03d}", + "text_embedding": batch["text_embedding"][i:i + 1].clone(), + "text_attention_mask": batch["text_attention_mask"][i:i + 1].clone(), + "first_frame_latent": batch["first_frame_latent"][i:i + 1].clone(), + "clip_feature": batch["clip_feature"][i:i + 1].clone(), + "vae_latent": batch["vae_latent"][i:i + 1].clone(), + "track_points": batch["track_points"][i:i + 1].clone(), + "track_visibility": batch["track_visibility"][i:i + 1].clone(), + "object_ids": batch["object_ids"][i:i + 1].clone() if "object_ids" in batch else None, + "track_weights": batch["track_weights"][i:i + 1].clone() if "track_weights" in batch else None, + "caption": str(infos[i].get("caption", "") if i < len(infos) else ""), + } for i in range(len(rows))] + + +@torch.no_grad() +def generate_with_cfg(model, sample: dict, *, w_text: float, w_motion: float, + num_steps: int, seed: int, + null_text: torch.Tensor | None = None, + null_text_mask: torch.Tensor | None = None, + no_track: bool = False) -> np.ndarray: + """Reproduce track_validation._sample denoise EXACTLY, w/ configurable CFG. + + If ``null_text`` is provided, use it as the ∅ text (canonical Wan/T5 convention + = UMT5-encoded empty string). Otherwise fall back to zeros_like(txt). + If ``no_track`` is True, pass tp=None/tv=None to every forward (matches the + val callback's `track_val/no_track` panel). + """ + device = model.device + dtype = torch.bfloat16 + + ff = sample["first_frame_latent"].to(device, dtype) + txt = sample["text_embedding"].to(device, dtype) + mask = sample["text_attention_mask"].to(device, dtype) + clip = sample["clip_feature"].to(device, dtype) + if no_track: + tp = None + tv = None + else: + # CRITICAL: sparse-sample tracks the same way training + track_validation do. + # Feeding the raw 2500 tracks is out-of-distribution and produces melted output. + tp_raw = sample["track_points"].float() + tv_raw = sample["track_visibility"].float() + oid = sample.get("object_ids") + tw = sample.get("track_weights") + aug_gen = torch.Generator(device=device).manual_seed(int(seed)) + tp_s, tv_s = model._augment_tracks(tp_raw, tv_raw, aug_gen, object_ids=oid, track_weights=tw) + if tp_s is None or tv_s is None: # motion_drop fired -> fall back to raw + tp_s, tv_s = tp_raw, tv_raw + tp = tp_s.to(device, dtype) + tv = tv_s.to(device, dtype) + cond20 = model._build_i2v_cond_concat(ff) + + _, _, T, H, W = ff.shape + gen = torch.Generator(device="cpu").manual_seed(int(seed)) + latents = torch.randn((1, 16, T, H, W), generator=gen, dtype=torch.float32).to(device) + sched = FlowMatchEulerDiscreteScheduler(shift=float(model.timestep_shift)) + sched.set_timesteps(int(num_steps), device=device) + + do_text_cfg = (w_text != 1.0) + do_motion_cfg = (w_motion != 1.0) and not no_track # no tracks -> motion CFG is a no-op + + def _fwd(text_e, tp_e, tv_e, mi, ts_): + with torch.autocast(device.type, dtype=dtype), \ + set_forward_context(current_timestep=ts_, attn_metadata=None): + return model.transformer(hidden_states=mi, encoder_hidden_states=text_e, + encoder_attention_mask=mask, timestep=ts_, + encoder_hidden_states_image=clip, + track_points=tp_e, track_visibility=tv_e, + return_dict=False) + + if do_text_cfg: + if null_text is not None: + # Broadcast the [seq_len, dim] null embedding to match txt shape [B, seq_len, dim] + nl = null_text.to(device, dtype) + if nl.dim() == 2: + nl = nl.unsqueeze(0).expand_as(txt).contiguous() + txt_null = nl + else: + txt_null = torch.zeros_like(txt) + else: + txt_null = None + + for tt in sched.timesteps: + mi = torch.cat([latents.to(dtype), cond20], dim=1) + ts_ = tt.reshape(1).to(device, dtype) + v_full = _fwd(txt, tp, tv, mi, ts_) + if not do_text_cfg and not do_motion_cfg: + v = v_full + elif do_text_cfg and not do_motion_cfg: + v_no_text = _fwd(txt_null, tp, tv, mi, ts_) + v = v_no_text + w_text * (v_full - v_no_text) + elif do_motion_cfg and not do_text_cfg: + v_no_motion = _fwd(txt, None, None, mi, ts_) + v = v_no_motion + w_motion * (v_full - v_no_motion) + else: + # joint text+motion CFG (matches track_validation._sample) + v_no_text = _fwd(txt_null, tp, tv, mi, ts_) + v_no_motion = _fwd(txt, None, None, mi, ts_) + alpha = w_text / max(w_text + w_motion, 1e-6) + v_base = alpha * v_no_text + (1.0 - alpha) * v_no_motion + v = v_base + w_text * (v_full - v_no_text) + w_motion * (v_full - v_no_motion) + latents = sched.step(v.float(), tt, latents.float(), return_dict=False)[0] + + return twi.decode_to_pixels(model, latents) + + +def main() -> None: + p = argparse.ArgumentParser() + p.add_argument("--model-dir", required=True) + p.add_argument("--yaml", required=True) + p.add_argument("--val-parquet", required=True) + p.add_argument("--wandb-project", default="wantrack-bidir") + p.add_argument("--wandb-run-name", required=True) + p.add_argument("--out-dir", default="/mnt/lustre/vlm-s4duan/cfg_ablation") + p.add_argument("--steps", type=int, default=30) + p.add_argument("--seed", type=int, default=1000) + p.add_argument("--fps", type=int, default=24) + p.add_argument("--configs", default="3.0,1.5;1.0,1.5;1.0,1.0", + help="';'-sep list of 'w_text,w_motion' pairs") + p.add_argument("--null-text-pt", + help="Path to precomputed UMT5(``'') embedding .pt file. If given, this is " + "used as the ∅ text branch instead of zeros_like(txt).") + p.add_argument("--no-track", action="store_true", + help="Force track_points=None throughout denoise. Reproduces the val " + "callback's `track_val/no_track` panel.") + args = p.parse_args() + + null_text = None + null_text_mask = None + if args.null_text_pt: + blob = torch.load(args.null_text_pt, map_location="cpu") + null_text = blob["embedding"] + null_text_mask = blob.get("attention_mask") + print(f"[abl] loaded null-text embedding: shape={list(null_text.shape)} " + f"norm={null_text.norm().item():.3f}", flush=True) + + import wandb + Path(args.out_dir).mkdir(parents=True, exist_ok=True) + + print(f"[abl] loading model {args.model_dir}", flush=True) + model, tc = twi.load_trackwan(args.model_dir, args.yaml) + text_len = int(getattr(tc.data, "text_padding_length", 256)) + + print(f"[abl] loading val samples from {args.val_parquet}", flush=True) + samples = load_val_samples(args.val_parquet, text_len) + print(f"[abl] {len(samples)} val samples", flush=True) + + configs = [tuple(float(x) for x in pair.split(",")) + for pair in args.configs.split(";")] + print(f"[abl] configs: {configs}", flush=True) + + run = wandb.init(project=args.wandb_project, name=args.wandb_run_name, + config={"model_dir": args.model_dir, "steps": args.steps, + "configs": [list(c) for c in configs]}) + + for i, s in enumerate(samples): + sid = s["id"] or f"clip{i:03d}" + print(f"[abl] === sample {i}: {sid} ===", flush=True) + wandb_log = {"sample": i, "caption": s["caption"][:120]} + for wt, wm in configs: + t0 = time.time() + frames = generate_with_cfg(model, s, w_text=wt, w_motion=wm, + num_steps=args.steps, seed=args.seed + i, + null_text=null_text, null_text_mask=null_text_mask, + no_track=args.no_track) + dt = time.time() - t0 + tag = f"wt{wt}_wm{wm}".replace(".", "p") + fn = Path(args.out_dir) / f"{sid}_{tag}.mp4" + imageio.mimsave(str(fn), frames, fps=args.fps, macro_block_size=1) + wandb_log[f"gen_{tag}"] = wandb.Video(str(fn), fps=args.fps, format="mp4") + print(f"[abl] {tag}: {dt:.1f}s -> {fn.name}", flush=True) + wandb.log(wandb_log) + + # also render GT reference once + for i, s in enumerate(samples): + try: + ref = twi.decode_reference(model, s["vae_latent"].to(model.device, torch.bfloat16)) + fn = Path(args.out_dir) / f"{s['id']}_gt.mp4" + imageio.mimsave(str(fn), ref, fps=args.fps, macro_block_size=1) + wandb.log({"sample": i, "reference_gt": wandb.Video(str(fn), fps=args.fps, format="mp4")}) + except Exception as e: + print(f"[abl] gt render failed for sample {i}: {e}", flush=True) + + wandb.finish() + print("[abl] DONE", flush=True) + + +if __name__ == "__main__": + main() diff --git a/data_pipeline/check_gate_growth.py b/data_pipeline/check_gate_growth.py new file mode 100644 index 000000000..9876c674c --- /dev/null +++ b/data_pipeline/check_gate_growth.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Report the WanTrack "gate" — patch_embedding.weight[:, 36:] — from a DCP checkpoint. + +The overfit stages (A/B) exist to grow this slot from EXACTLY ZERO up to roughly the scale the +successful 1.3B merged init reached (std ~0.011), co-adapting it with track_encoder. If it is +still ~0 there is no bootstrap happening and step C would merge nothing; if it explodes past the +pretrained channel scale (~0.0124) the track pathway is overwhelming the image prior. + +Loads ONLY the tensors of interest (DCP resharding is per-tensor), so it costs ~MBs, not the +~200GB a full checkpoint load would. + +Usage: python data_pipeline/check_gate_growth.py [more dirs...] +""" +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import torch +import torch.distributed.checkpoint as dcp +from torch.distributed.checkpoint import FileSystemReader + +BASE_IN = 36 +# NOTE: the live module nests the conv, so a TRAINED checkpoint stores +# 'patch_embedding.proj.weight', while the converter writes 'patch_embedding.weight' into the +# init safetensors. Try both — silently matching neither is how the earlier LWLR experiment +# ended up measuring nothing. +PE_KEYS = ["patch_embedding.proj.weight", "patch_embedding.weight"] +KEYS = PE_KEYS + [ + "track_encoder.proj.weight", + "track_encoder.temporal_conv.weight", +] + + +def _find(md_keys, want: str) -> str | None: + """DCP nests model tensors under a role prefix (e.g. 'student.'); match on the suffix.""" + exact = [k for k in md_keys if k == want] + if exact: + return exact[0] + cands = [k for k in md_keys if k.endswith(want)] + return cands[0] if cands else None + + +def report(ckpt: Path) -> None: + dcp_dir = ckpt / "dcp" if (ckpt / "dcp").is_dir() else ckpt + reader = FileSystemReader(str(dcp_dir)) + md = reader.read_metadata() + planned = md.state_dict_metadata + sd = {} + resolved = {} + for want in KEYS: + k = _find(list(planned.keys()), want) + if k is None: + continue + meta = planned[k] + sd[k] = torch.empty(tuple(meta.size), dtype=meta.properties.dtype) + resolved[want] = k + if not sd: + print(f"{ckpt.name}: none of {KEYS} found in checkpoint metadata") + return + dcp.load(sd, checkpoint_id=str(dcp_dir)) + + out = [f"{ckpt.name:>18s}"] + pe_k = next((resolved[k] for k in PE_KEYS if k in resolved), None) + if pe_k is None: + out.append("!! patch_embedding NOT FOUND — cannot report gate") + else: + pe = sd[pe_k].float() + out.append(f"gate[:, {BASE_IN}:] std={pe[:, BASE_IN:].std():.6f}") + out.append(f"pretrained std={pe[:, :BASE_IN].std():.6f}") + for w in ("track_encoder.proj.weight", "track_encoder.temporal_conv.weight"): + if w in resolved: + out.append(f"{w.split('.')[1]}_norm={sd[resolved[w]].float().norm():.4f}") + print(" ".join(out)) + + +def main() -> None: + if len(sys.argv) < 2: + print(__doc__) + sys.exit(1) + for key, default in [("RANK", "0"), ("LOCAL_RANK", "0"), ("WORLD_SIZE", "1"), + ("MASTER_ADDR", "127.0.0.1"), ("MASTER_PORT", "29555")]: + os.environ.setdefault(key, default) + print(f"reference: 1.3B merged-from-overfit gate std = 0.011418 (grown from 0.000000)") + for p in sys.argv[1:]: + try: + report(Path(p)) + except Exception as e: # keep going across checkpoints + print(f"{Path(p).name}: ERROR {type(e).__name__}: {e}") + + +if __name__ == "__main__": + main() diff --git a/data_pipeline/convert_trackwan_init_v2.py b/data_pipeline/convert_trackwan_init_v2.py new file mode 100644 index 000000000..f0059f4da --- /dev/null +++ b/data_pipeline/convert_trackwan_init_v2.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Build a WanTrack init checkpoint from a base Wan diffusers model, with several +strategies for the patch_embedding track-channel slot and the track_encoder. + +Modes for patch_embedding[:, base_in:] (the new 16 in-channels): + --pe-init zero : zero-pad (original stage-1 recipe) + --pe-init random : normal-init with std matched to base first-N channels + +Modes for track_encoder.{proj, temporal_conv}: + --track-src default : default Conv init (used by original convert) + --track-src : lift both convs from an existing WanTrack safetensors + (e.g. a trained 1.3B ckpt). Bias is copied when present + and its shape matches. + +Usage: + python data_pipeline/convert_trackwan_init_v2.py \ + --base /mnt/lustre/vlm-s4duan/models/Wan2.1-I2V-14B-720P-Diffusers \ + --out /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_random_init \ + --id-dim 64 --pe-init random + + python data_pipeline/convert_trackwan_init_v2.py \ + --base /mnt/lustre/vlm-s4duan/models/Wan2.1-I2V-14B-720P-Diffusers \ + --out /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_partial_merged \ + --id-dim 64 --pe-init random \ + --track-src /mnt/lustre/vlm-s4duan/exports/synth_stage2_paperLR_ckpt600/transformer/model.safetensors +""" +from __future__ import annotations + +import argparse +import json +import os +import shutil +from pathlib import Path + +import torch +from safetensors.torch import load_file, save_file +from torch import nn + +NEW_IN_CHANNELS = 52 +TRACK_CHANNELS = 16 +VAE_T_COMP = 4 +DEFAULT_ID_DIM = 128 + + +def build_track_encoder_default(id_dim: int, dtype: torch.dtype, use_bias: bool = False) -> dict[str, torch.Tensor]: + """Default (non-zero) Conv init for track_encoder.proj + temporal_conv.""" + temporal_conv = nn.Conv3d(id_dim, TRACK_CHANNELS, + kernel_size=(VAE_T_COMP, 1, 1), stride=(VAE_T_COMP, 1, 1), + bias=use_bias) + proj = nn.Conv3d(TRACK_CHANNELS, TRACK_CHANNELS, kernel_size=1, bias=use_bias) + d = { + "track_encoder.temporal_conv.weight": temporal_conv.weight.detach().to(dtype).clone(), + "track_encoder.proj.weight": proj.weight.detach().to(dtype).clone(), + } + if use_bias: + d["track_encoder.temporal_conv.bias"] = temporal_conv.bias.detach().to(dtype).clone() + d["track_encoder.proj.bias"] = proj.bias.detach().to(dtype).clone() + return d + + +def load_src_state(src_path: str) -> dict[str, torch.Tensor]: + """Load a source state dict from a .safetensors FILE or a diffusers model DIR. + + A 14B export is sharded across several *.safetensors, so accept a directory (either the + model root or its transformer/ subdir) and merge every shard. + """ + p = Path(src_path) + if p.is_file(): + return load_file(str(p)) + tdir = p / "transformer" if (p / "transformer").is_dir() else p + shards = sorted(tdir.glob("*.safetensors")) + if not shards: + raise FileNotFoundError(f"no *.safetensors under {tdir}") + state: dict[str, torch.Tensor] = {} + for s in shards: + state.update(load_file(str(s))) + print(f"[src] loaded {len(state)} tensors from {len(shards)} shard(s) in {tdir}") + return state + + +def lift_track_encoder_from_src(src_path: str, id_dim: int, dtype: torch.dtype) -> dict[str, torch.Tensor]: + """Load track_encoder.{proj, temporal_conv}.{weight, bias} from an existing WanTrack ckpt. + + Verifies shape compatibility. If bias is missing in the source we skip it; if the + source shapes don't match the expected shapes we fall back to default init for that + key and print a warning. + """ + src_state = load_src_state(src_path) + expected_shapes = { + "track_encoder.temporal_conv.weight": (TRACK_CHANNELS, id_dim, VAE_T_COMP, 1, 1), + "track_encoder.proj.weight": (TRACK_CHANNELS, TRACK_CHANNELS, 1, 1, 1), + "track_encoder.temporal_conv.bias": (TRACK_CHANNELS,), + "track_encoder.proj.bias": (TRACK_CHANNELS,), + } + picked: dict[str, torch.Tensor] = {} + for k, want in expected_shapes.items(): + if k not in src_state: + continue + got = tuple(src_state[k].shape) + if got != want: + print(f"[track-src] SHAPE MISMATCH on {k}: src={got} vs expected={want}; SKIPPING (will need fallback)") + continue + picked[k] = src_state[k].detach().to(dtype).clone() + print(f"[track-src] lifted {k} shape={got}") + # If either weight is missing, we still need to keep the model loadable. + fallback = build_track_encoder_default(id_dim, dtype, use_bias=False) + for k, v in fallback.items(): + if k not in picked: + print(f"[track-src] {k} not in source -> default init") + picked[k] = v + return picked + + +def std_of_first_n(w: torch.Tensor, n: int) -> float: + """Std of the pretrained first-N in-channels slice of patch_embedding.weight.""" + with torch.no_grad(): + s = w[:, :n].float().std().item() + return float(s) + + +def main() -> None: + p = argparse.ArgumentParser() + p.add_argument("--base", required=True, help="Base Wan diffusers model dir.") + p.add_argument("--out", required=True, help="Output diffusers model dir.") + p.add_argument("--id-dim", type=int, default=DEFAULT_ID_DIM) + p.add_argument("--pe-init", choices=("zero", "random"), default="zero", + help="patch_embedding init strategy for the ADDED track channels [:, base_in:]. " + "Base pretrained channels [:, :base_in] are always preserved.") + p.add_argument("--pe-random-seed", type=int, default=1234) + p.add_argument("--track-src", default=None, + help="Path to a .safetensors (or diffusers model dir) with track_encoder.* to lift.") + p.add_argument("--pe-src", default=None, + help="Path to a .safetensors (or diffusers model dir) to lift patch_embedding.weight[:, base_in:] " + "from — i.e. the TRACK SLOT only; the pretrained [:, :base_in] channels always come from " + "--base. Use together with --track-src pointing at the SAME source so the encoder and the " + "track slot stay CO-ADAPTED (lifting the encoder alone is worse than random init).") + p.add_argument("--use-bias-defaults", action="store_true", + help="When --track-src is not set, build track_encoder convs WITH bias (matches TRACKWAN_TRACK_BIAS=1 training).") + args = p.parse_args() + + id_dim = int(args.id_dim) + TRACK_CONFIG = { + "id_dim": id_dim, + "track_channels": TRACK_CHANNELS, + "vae_spatial_compression": 8, + "vae_temporal_compression": VAE_T_COMP, + "max_track_id": 100_000, + "zero_init_head": False, + } + + base = Path(args.base) + out = Path(args.out) + if not (base / "transformer" / "config.json").exists(): + raise FileNotFoundError(f"{base}/transformer/config.json not found") + out.mkdir(parents=True, exist_ok=True) + + # 1) Symlink every top-level entry except transformer/ + for entry in os.listdir(base): + if entry == "transformer": + continue + src_p = (base / entry).resolve() + dst = out / entry + if dst.is_symlink() or dst.exists(): + if dst.is_dir() and not dst.is_symlink(): + shutil.rmtree(dst) + else: + dst.unlink() + os.symlink(src_p, dst) + + # 2) Rewrite transformer/config.json + tdir = out / "transformer" + tdir.mkdir(exist_ok=True) + cfg = json.loads((base / "transformer" / "config.json").read_text()) + base_in = int(cfg["in_channels"]) + cfg["in_channels"] = NEW_IN_CHANNELS + cfg["track_config"] = TRACK_CONFIG + (tdir / "config.json").write_text(json.dumps(cfg, indent=2)) + + # 3) Load full base transformer state + sf_files = sorted((base / "transformer").glob("*.safetensors")) + if not sf_files: + raise FileNotFoundError(f"No safetensors under {base}/transformer") + state: dict[str, torch.Tensor] = {} + for sf in sf_files: + state.update(load_file(str(sf))) + dtype = next(iter(state.values())).dtype + + # 4) Grow patch_embedding.weight from base_in -> 52 input channels + pe_key = "patch_embedding.weight" + if pe_key not in state: + raise KeyError(f"{pe_key} not found in base transformer") + w = state[pe_key] # [hidden, base_in, 1, 2, 2] + if w.shape[1] != base_in: + raise ValueError(f"{pe_key} in_ch {w.shape[1]} != config in_channels {base_in}") + new_w = torch.empty((w.shape[0], NEW_IN_CHANNELS, *w.shape[2:]), dtype=w.dtype) + new_w[:, :base_in] = w # pretrained channels first + if args.pe_init == "zero": + new_w[:, base_in:] = 0 + print(f"[convert] pe-init=zero -> new channels are zero") + else: + std = std_of_first_n(w, base_in) + g = torch.Generator().manual_seed(int(args.pe_random_seed)) + noise = torch.empty_like(new_w[:, base_in:], dtype=torch.float32) + noise.normal_(mean=0.0, std=std, generator=g) + new_w[:, base_in:] = noise.to(w.dtype) + print(f"[convert] pe-init=random N(0, {std:.5f}) matching first-{base_in} std") + # 4b) Optionally overwrite ONLY the track slot [:, base_in:] from a trained source. + # This is the other half of the "merged" recipe: the track slot and the track_encoder were + # co-adapted during the overfit, so they must be lifted TOGETHER (--pe-src + --track-src from + # the same ckpt). Lifting the encoder alone leaves it talking to a random projection. + if args.pe_src: + src_pe_state = load_src_state(args.pe_src) + if pe_key not in src_pe_state: + raise KeyError(f"{pe_key} not found in --pe-src {args.pe_src}") + src_pe = src_pe_state[pe_key] + want = (w.shape[0], NEW_IN_CHANNELS, *w.shape[2:]) + if tuple(src_pe.shape) != want: + raise ValueError(f"--pe-src {pe_key} shape {tuple(src_pe.shape)} != expected {want}") + slot = src_pe[:, base_in:].to(w.dtype) + new_w[:, base_in:] = slot + print(f"[pe-src] lifted {pe_key}[:, {base_in}:] from source " + f"(std={slot.float().std().item():.5f}, absmax={slot.float().abs().max().item():.5f})") + + state[pe_key] = new_w + print(f"[convert] patch_embedding.weight: {tuple(w.shape)} -> {tuple(new_w.shape)}") + with torch.no_grad(): + print(f"[convert] pe[:, :{base_in}] std={new_w[:, :base_in].float().std().item():.5f} (pretrained) | " + f"pe[:, {base_in}:] std={new_w[:, base_in:].float().std().item():.5f} (track slot)") + + # 5) Add track_encoder + if args.track_src: + te = lift_track_encoder_from_src(args.track_src, id_dim, dtype) + else: + te = build_track_encoder_default(id_dim, dtype, use_bias=args.use_bias_defaults) + for k, v in te.items(): + state[k] = v + print(f"[convert] added {len(te)} track_encoder tensors") + + save_file(state, str(tdir / "diffusion_pytorch_model.safetensors"), metadata={"format": "pt"}) + print(f"[convert] wrote {tdir/'diffusion_pytorch_model.safetensors'} ({len(state)} keys)") + print(f"[convert] done -> {out}") + + +if __name__ == "__main__": + main() diff --git a/data_pipeline/decode_wansyn32k.py b/data_pipeline/decode_wansyn32k.py new file mode 100644 index 000000000..7fed2cac6 --- /dev/null +++ b/data_pipeline/decode_wansyn32k.py @@ -0,0 +1,162 @@ +#!/usr/bin/env python3 +"""Decode FastVideo Wan2.2-Syn-121x704x1280_32k parquet latents to 480x832 mp4s. + +The HF dataset stores Wan 2.2 5B VAE latents (48-channel, 16x spatial, 4x +temporal). We can't feed these to our Wan 2.1 WanTrack model directly, so we +decode → pixels → downsize to our training resolution 480x832 → save mp4. +Downstream (SAM + CoTracker + i2v_track preprocess) then treats these as if they +were raw synth mp4s. + +Idempotent + resumable: skips samples whose mp4 already exists. + +Slurm sharded via ``--num-shards``/``--shard``. Each worker processes +``parquets[shard::num_shards]``. +""" +from __future__ import annotations + +import argparse +import glob +import json +import os +import time +from pathlib import Path + +import imageio.v2 as imageio +import numpy as np +import pyarrow.parquet as pq +import torch +from diffusers import AutoencoderKLWan +from PIL import Image + + +def _resize_frames(frames: np.ndarray, out_h: int, out_w: int) -> np.ndarray: + """Center-crop-then-resize a (T,H,W,3) uint8 stack to (T,out_h,out_w,3).""" + T, H, W, _ = frames.shape + src_ar = W / H + dst_ar = out_w / out_h + if src_ar > dst_ar: + new_w = int(round(H * dst_ar)) + x0 = (W - new_w) // 2 + cropped = frames[:, :, x0:x0 + new_w, :] + else: + new_h = int(round(W / dst_ar)) + y0 = (H - new_h) // 2 + cropped = frames[:, y0:y0 + new_h, :, :] + out = np.empty((T, out_h, out_w, 3), dtype=np.uint8) + for t in range(T): + out[t] = np.asarray(Image.fromarray(cropped[t]).resize((out_w, out_h), Image.BILINEAR)) + return out + + +@torch.no_grad() +def decode_row(vae: AutoencoderKLWan, latent_bytes: bytes, latent_shape: list[int], + device: torch.device, dtype: torch.dtype) -> np.ndarray: + """Decode one row's Wan 2.2 5B latent to a [T,704,1280,3] uint8 pixel stack.""" + lat = np.frombuffer(latent_bytes, dtype=np.float32).reshape(latent_shape) + lat = torch.from_numpy(lat).to(device=device, dtype=dtype).unsqueeze(0) # [1,C,T,H,W] + # unnormalize per-channel + m = torch.tensor(vae.config.latents_mean, device=device, dtype=dtype).view(1, -1, 1, 1, 1) + s = torch.tensor(vae.config.latents_std, device=device, dtype=dtype).view(1, -1, 1, 1, 1) + lat = lat * s + m + pix = vae.decode(lat, return_dict=False)[0] # [1,3,T_out,H*16,W*16] in [-1,1] + pix = pix.clamp(-1, 1).float() + pix = ((pix + 1.0) * 127.5).round().to(torch.uint8) + pix = pix[0].permute(1, 2, 3, 0).cpu().numpy() # [T,H,W,3] + return pix + + +def main() -> None: + p = argparse.ArgumentParser() + p.add_argument("--parquet-dir", required=True, help="dir containing train/Part_*/latents_chunk_*.parquet") + p.add_argument("--vae-dir", required=True, help="Wan 2.2 5B VAE dir") + p.add_argument("--out-dir", required=True, help="output root; writes videos/vid_*.mp4 + meta/*.json") + p.add_argument("--out-h", type=int, default=480) + p.add_argument("--out-w", type=int, default=832) + p.add_argument("--fps", type=int, default=24) + p.add_argument("--shard", type=int, default=0) + p.add_argument("--num-shards", type=int, default=1) + p.add_argument("--dtype", default="bfloat16", choices=["float32", "bfloat16", "float16"]) + p.add_argument("--limit-rows", type=int, default=None, help="cap rows for smoke test") + args = p.parse_args() + + device = torch.device("cuda") + dtype = getattr(torch, args.dtype) + + vids_dir = Path(args.out_dir) / "videos" + meta_dir = Path(args.out_dir) / "meta" + vids_dir.mkdir(parents=True, exist_ok=True) + meta_dir.mkdir(parents=True, exist_ok=True) + + files = sorted(glob.glob(os.path.join(args.parquet_dir, "**", "*.parquet"), recursive=True)) + if not files: + raise FileNotFoundError(f"no *.parquet under {args.parquet_dir}") + my_files = files[args.shard::args.num_shards] + print(f"[dec {args.shard}/{args.num_shards}] {len(my_files)} parquets", flush=True) + + vae = AutoencoderKLWan.from_pretrained(args.vae_dir, torch_dtype=dtype).to(device).eval() + + n_ok = 0 + n_skip = 0 + n_err = 0 + t0 = time.time() + row_budget = args.limit_rows or 10**12 + processed_rows = 0 + for fk, f in enumerate(my_files): + try: + table = pq.read_table(f) + except Exception as e: # skip unreadable parquets + print(f"[dec {args.shard}] parquet read fail {f}: {e}", flush=True) + continue + # Dataset's `id` field only encodes chunk index (not Part), so the same id + # appears in many Parts. Prefix Part directory so vid_ids are unique. + part_name = Path(f).parent.name # e.g. "Part_57" + rows = table.to_pylist() + for r in rows: + if processed_rows >= row_budget: + break + processed_rows += 1 + row_key = str(r.get("id") or r.get("file_name") or f"{fk:05d}_{n_ok:04d}") + vid_id = f"{part_name}_{row_key}" + vid_path = vids_dir / f"{vid_id}.mp4" + meta_path = meta_dir / f"{vid_id}.json" + if vid_path.exists() and meta_path.exists(): + n_skip += 1 + continue + try: + pix = decode_row(vae, r["vae_latent_bytes"], list(r["vae_latent_shape"]), + device=device, dtype=dtype) + resized = _resize_frames(pix, args.out_h, args.out_w) # [T, out_h, out_w, 3] + # atomic mp4 write (imageio infers format from extension → keep .mp4) + tmp = vid_path.parent / f".{vid_path.name}.tmp.mp4" + imageio.mimsave(str(tmp), resized, fps=args.fps, codec="libx264", + quality=7, macro_block_size=1) + tmp.rename(vid_path) + meta = { + "path": vid_path.name, + "id": vid_id, + "cap": [str(r.get("caption") or "")], + "fps": float(args.fps), + "num_frames": int(resized.shape[0]), + "resolution": [int(args.out_h), int(args.out_w)], + "src_resolution": [int(r.get("height", 704)), int(r.get("width", 1280))], + "src_dataset": "FastVideo/Wan2.2-Syn-121x704x1280_32k", + } + meta_path.write_text(json.dumps(meta)) + n_ok += 1 + if n_ok % 20 == 0: + dt = time.time() - t0 + rate = n_ok / max(dt, 1e-9) + print(f"[dec {args.shard}] ok={n_ok} skip={n_skip} err={n_err} " + f"rate={rate:.2f}/s", flush=True) + except Exception as e: # skip a bad row, keep going + n_err += 1 + print(f"[dec {args.shard}] row {vid_id} FAILED: {e}", flush=True) + if processed_rows >= row_budget: + break + + print(f"[dec {args.shard}] DONE ok={n_ok} skip={n_skip} err={n_err} " + f"in {(time.time()-t0)/60:.1f}min", flush=True) + + +if __name__ == "__main__": + main() diff --git a/data_pipeline/lwlr_experiment.py b/data_pipeline/lwlr_experiment.py new file mode 100644 index 000000000..678c9602d --- /dev/null +++ b/data_pipeline/lwlr_experiment.py @@ -0,0 +1,246 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Layer-wise learning-rate (LWLR) bootstrap experiment. + +Question: on a 14B random-init WanTrack model, does giving the track pathway +(track_encoder.* + patch_embedding.weight[:, 36:]) a 100x higher LR than the +rest actually let it grow into functional gradient magnitudes within a small +number of steps? Or is the encoder stuck regardless? + +Runs a short training loop on real OpenVid samples, tracking weight+grad norms +of the track pathway at each step. Compares against the "known good" magnitude +observed in 1.3B_merged (per-sample track_encoder.proj grad ~0.5). +""" +from __future__ import annotations + +import argparse +import glob +import json +import os +import random +import sys +from pathlib import Path + +import numpy as np +import pyarrow.parquet as pq +import torch +import torch.nn.functional as F +from safetensors.torch import load_file as st_load + +# Distributed init BEFORE importing fastvideo (same pattern as gradient_flow_analysis) +def _init_distributed_1gpu() -> None: + os.environ.setdefault("RANK", "0") + os.environ.setdefault("LOCAL_RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29513") + os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + from fastvideo.distributed.parallel_state import ( + maybe_init_distributed_environment_and_model_parallel, + ) + maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1) + + +_init_distributed_1gpu() + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from gradient_flow_analysis import ( # noqa: E402 + build_model_from_diffusers, + load_samples, + build_i2v_cond_concat, + flow_matching_noisy, + Sample, +) + + +TRACK_PATH_PARAMS = ( + "track_encoder.proj.weight", + "track_encoder.temporal_conv.weight", + "track_encoder.proj.bias", # may not exist depending on TRACKWAN_TRACK_BIAS + "track_encoder.temporal_conv.bias", +) + + +def split_param_groups(model, high_lr: float, base_lr: float) -> list[dict]: + """Give track pathway (track_encoder.* + patch_embedding[:, 36:]) a high LR. + + patch_embedding.weight is a SINGLE tensor of shape [hidden, 52, 1, 2, 2] — we can't + give sub-slice a different LR through standard param groups. Workaround: put the + whole patch_embedding in a "hybrid" group with high LR (the base 36 channels will + also move faster, but by only 100x for a short warmup that's small — 100 steps of + high LR is still << stage-1 4800 steps of normal LR). + """ + high_group, base_group = [], [] + named = dict(model.named_parameters()) + for n, p in named.items(): + if not p.requires_grad: + continue + # High-LR pathway: encoder convs (+ bias if present) + the whole patch_embedding.weight + is_encoder = n.startswith("track_encoder.") + is_patch_embed = (n == "patch_embedding.weight") + if is_encoder or is_patch_embed: + high_group.append(p) + else: + base_group.append(p) + return [ + {"params": high_group, "lr": high_lr, "name": "track_pathway"}, + {"params": base_group, "lr": base_lr, "name": "dit_body"}, + ] + + +def snapshot_norms(model, sample_grad_i: int) -> dict: + """Return weight/grad norms for the diagnostic params.""" + out = {} + for n, p in model.named_parameters(): + if not p.requires_grad: + continue + # Track params of interest + keep = ( + n.startswith("track_encoder.") + or n == "patch_embedding.weight" + ) + if not keep: + continue + w_norm = p.detach().float().norm().item() + g_norm = p.grad.detach().float().norm().item() if p.grad is not None else 0.0 + # For patch_embedding, split into base/track slices + if n == "patch_embedding.weight": + w_base = p.detach()[:, :36].float() + w_track = p.detach()[:, 36:].float() + out["patch_embedding.weight[:, :36]"] = { + "w_norm": w_base.norm().item(), "g_norm": 0.0, + } + out["patch_embedding.weight[:, 36:]"] = { + "w_norm": w_track.norm().item(), "g_norm": 0.0, + } + if p.grad is not None: + g_base = p.grad.detach()[:, :36].float() + g_track = p.grad.detach()[:, 36:].float() + out["patch_embedding.weight[:, :36]"]["g_norm"] = g_base.norm().item() + out["patch_embedding.weight[:, 36:]"]["g_norm"] = g_track.norm().item() + else: + out[n] = {"w_norm": w_norm, "g_norm": g_norm} + return out + + +def one_step(model, s: Sample, device: torch.device, flow_shift: float, + from_ctx_ns) -> float: + """Forward+backward on one sample; grads accumulate.""" + vae = s.vae_latent.unsqueeze(0).to(device, dtype=torch.bfloat16) + ff = s.first_frame_latent.unsqueeze(0).to(device, dtype=torch.bfloat16) + clip = s.clip_feature.unsqueeze(0).to(device, dtype=torch.bfloat16) + text = s.text_embedding.unsqueeze(0).to(device, dtype=torch.bfloat16) + mask = torch.ones(1, s.text_embedding.shape[0], device=device, dtype=torch.bfloat16) + tp = s.track_points.unsqueeze(0).to(device, dtype=torch.bfloat16) + tv = s.track_visibility.unsqueeze(0).to(device, dtype=torch.bfloat16) + + num_latent_t = vae.shape[2] + expected_T = (num_latent_t - 1) * 4 + 1 + tp = tp[:, :expected_T] + tv = tv[:, :expected_T] + + sigma_u = float(torch.rand(1).item()) + noise = torch.randn_like(vae, dtype=torch.float32) + clean_f = vae.float() + noisy_f, _, ts_val = flow_matching_noisy(clean_f, noise, sigma_u, + flow_shift=flow_shift) + noisy = noisy_f.to(torch.bfloat16) + cond20 = build_i2v_cond_concat(ff, vae_temporal_compression=4) + hs = torch.cat([noisy, cond20], dim=1) + ts = torch.tensor([ts_val], device=device, dtype=torch.bfloat16) + + with torch.autocast(device_type="cuda", dtype=torch.bfloat16), from_ctx_ns(current_timestep=ts, attn_metadata=None): + pred = model( + hidden_states=hs, encoder_hidden_states=text, encoder_attention_mask=mask, + timestep=ts, encoder_hidden_states_image=clip, + track_points=tp, track_visibility=tv, return_dict=False, + ) + target = noise - clean_f + loss = F.mse_loss(pred.float(), target.float()) + loss.backward() + return float(loss.item()) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="/mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_random_init") + ap.add_argument("--n-samples", type=int, default=16) + ap.add_argument("--n-steps", type=int, default=60) + ap.add_argument("--batch-accum", type=int, default=4, help="gradient accumulation steps per optim step") + ap.add_argument("--high-lr", type=float, default=1e-3, help="LR for track pathway") + ap.add_argument("--base-lr", type=float, default=1e-5, help="LR for the rest of DiT") + ap.add_argument("--flow-shift", type=float, default=6.0) + ap.add_argument("--parquet-glob", type=str, + default="/mnt/lustre/vlm-s4duan/openvid_1m/combined_parquet_dataset/shard_*/**/*.parquet") + ap.add_argument("--out", type=str, + default="/mnt/lustre/vlm-s4duan/gradient_analysis_14b/lwlr_trace.json") + ap.add_argument("--seed", type=int, default=42) + args = ap.parse_args() + + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + + from fastvideo.forward_context import set_forward_context + device = torch.device("cuda") + + samples = load_samples(args.parquet_glob, args.n_samples, seed=args.seed) + print(f"loaded {len(samples)} samples") + assert len(samples) >= args.batch_accum, "need enough samples for accumulation" + + print(f"loading {args.model} ...") + model = build_model_from_diffusers(args.model, device) + + # Set up param groups + param_groups = split_param_groups(model, high_lr=args.high_lr, base_lr=args.base_lr) + n_high = sum(p.numel() for p in param_groups[0]["params"]) + n_base = sum(p.numel() for p in param_groups[1]["params"]) + print(f"track_pathway params: {n_high/1e6:.2f}M (LR={args.high_lr:g})") + print(f"dit_body params: {n_base/1e6:.2f}M (LR={args.base_lr:g})") + optim = torch.optim.AdamW(param_groups, weight_decay=0.0) # no wd so we see raw drift + + trace: list[dict] = [] + # step 0 (before any update) + print("\n=== step 0 (init) ===") + norms0 = snapshot_norms(model, -1) + for k, v in norms0.items(): + print(f" {k:<45s} w={v['w_norm']:.4f} g={v['g_norm']:.4f}") + trace.append({"step": 0, "loss": None, "norms": norms0}) + + step = 0 + accum = 0 + losses_accum = [] + for outer in range(args.n_steps * args.batch_accum): + s = samples[outer % len(samples)] + loss = one_step(model, s, device, args.flow_shift, set_forward_context) + losses_accum.append(loss) + accum += 1 + if accum < args.batch_accum: + continue + step += 1 + # snapshot BEFORE step (grads are accumulated over the batch) + gs = snapshot_norms(model, -1) + optim.step() + optim.zero_grad(set_to_none=True) + avg_loss = float(np.mean(losses_accum)) + losses_accum = [] + accum = 0 + if step % 5 == 0 or step <= 5: + print(f"\n=== step {step} (loss={avg_loss:.4f}) ===") + for k, v in gs.items(): + marker = "" + if k == "track_encoder.proj.weight": + marker = " <-- 1.3B_merged saw per_sample_norm ~0.50" + elif k == "patch_embedding.weight[:, 36:]": + marker = " <-- 1.3B_merged saw per_sample_norm ~0.73" + print(f" {k:<45s} w={v['w_norm']:.4f} g={v['g_norm']:.4f}{marker}") + trace.append({"step": step, "loss": avg_loss, "norms": gs}) + + Path(args.out).parent.mkdir(parents=True, exist_ok=True) + with open(args.out, "w") as f: + json.dump({"config": vars(args), "trace": trace}, f, indent=2) + print(f"\nwrote trace -> {args.out}") + + +if __name__ == "__main__": + main() diff --git a/data_pipeline/precompute_null_text.py b/data_pipeline/precompute_null_text.py new file mode 100644 index 000000000..217398c0d --- /dev/null +++ b/data_pipeline/precompute_null_text.py @@ -0,0 +1,81 @@ +#!/usr/bin/env python3 +"""Precompute the UMT5 embedding of an empty string. + +This is the canonical Wan/T5 null-text convention: feed "" to the text encoder +and use the resulting non-zero embedding as the ∅ in classifier-free guidance. +Contrast with `torch.zeros_like(text_embedding)`, which is a different tensor +distribution that the base model was never asked to handle. + +Saves a dict: + {"embedding": [seq_len, dim], "attention_mask": [seq_len]} +so the val callback can just index into these when it needs a null text. + +Usage: + python data_pipeline/precompute_null_text.py \ + --model-dir /mnt/lustre/vlm-s4duan/exports/synth_stage2_paperLR_ckpt400 \ + --output /mnt/lustre/vlm-s4duan/exports/null_text_umt5.pt +""" +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch + + +def main() -> None: + p = argparse.ArgumentParser() + p.add_argument("--model-dir", required=True, + help="Diffusers export dir containing text_encoder/ and tokenizer/") + p.add_argument("--output", required=True) + p.add_argument("--text-len", type=int, default=256, + help="Pad/truncate to this length (matches training text_padding_length)") + p.add_argument("--device", default="cuda") + args = p.parse_args() + + from transformers import AutoTokenizer + try: + from transformers import UMT5EncoderModel # type: ignore + except ImportError: + UMT5EncoderModel = None # type: ignore + from transformers import AutoModel + + model_dir = Path(args.model_dir) + tok = AutoTokenizer.from_pretrained(model_dir / "tokenizer") + if UMT5EncoderModel is not None: + try: + enc = UMT5EncoderModel.from_pretrained(model_dir / "text_encoder", torch_dtype=torch.bfloat16) + except Exception: + enc = AutoModel.from_pretrained(model_dir / "text_encoder", torch_dtype=torch.bfloat16) + else: + enc = AutoModel.from_pretrained(model_dir / "text_encoder", torch_dtype=torch.bfloat16) + enc = enc.to(args.device).eval() + + tokens = tok( + [""], + max_length=args.text_len, + padding="max_length", + truncation=True, + return_tensors="pt", + return_attention_mask=True, + ) + input_ids = tokens["input_ids"].to(args.device) + attn_mask = tokens["attention_mask"].to(args.device) + print(f"tokenized '' -> {input_ids.shape}, non-pad tokens = {int(attn_mask.sum().item())}") + print(f" first 5 token ids: {input_ids[0, :5].tolist()}") + + with torch.no_grad(): + out = enc(input_ids=input_ids, attention_mask=attn_mask) + # take last_hidden_state; strip the batch dim + hs = out.last_hidden_state[0].float().cpu() # [seq_len, dim] + am = attn_mask[0].cpu() + print(f"embedding shape={list(hs.shape)} norm={hs.norm().item():.3f} mean={hs.mean().item():.5f}") + print(f"attn_mask shape={list(am.shape)} sum={int(am.sum().item())}") + + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + torch.save({"embedding": hs, "attention_mask": am, "text_len": args.text_len}, args.output) + print(f"[null-text] saved -> {args.output}") + + +if __name__ == "__main__": + main() diff --git a/data_pipeline/run_decode_wansyn32k_slurm.sh b/data_pipeline/run_decode_wansyn32k_slurm.sh new file mode 100755 index 000000000..7d8b7f199 --- /dev/null +++ b/data_pipeline/run_decode_wansyn32k_slurm.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# Wan 2.2 5B VAE decode of the 32k HF latent dataset -> 480x832 mp4s. +# Sharded across NODES * GPUS * PROCS_PER_GPU workers via SLURM_PROCID. +set -euo pipefail +WORK=/mnt/lustre/vlm-s4duan +: "${PARQUET_DIR:?}"; : "${OUT_DIR:?}"; : "${VAE_DIR:?}" +NODES=${NODES:-8}; GPUS=${GPUS:-4}; PROCS_PER_GPU=${PROCS_PER_GPU:-2} +TASKS_PER_NODE=$(( GPUS * PROCS_PER_GPU )) +CPUS_PER_TASK=$(( 128 / TASKS_PER_NODE )); [ "$CPUS_PER_TASK" -lt 1 ] && CPUS_PER_TASK=1 +NUM_SHARDS=$(( NODES * TASKS_PER_NODE )) + +mkdir -p "$OUT_DIR" "$WORK/logs" +echo "nodes=$NODES gpus/node=$GPUS procs/gpu=$PROCS_PER_GPU -> $NUM_SHARDS workers" + +sbatch -N "$NODES" --gres=gpu:$GPUS --ntasks-per-node=$TASKS_PER_NODE --exclusive \ + --cpus-per-task=$CPUS_PER_TASK --mem=0 -t 12:00:00 -J decode_wansyn32k \ + --chdir="$WORK/FastVideo" -o "$WORK/logs/decode_wansyn32k_%j.out" -e "$WORK/logs/decode_wansyn32k_%j.out" \ + --wrap "srun --chdir=$WORK/FastVideo \ + bash -lc 'source .venv/bin/activate && \ + export HOME=$WORK HF_HOME=$WORK/.hf TORCH_HOME=$WORK/.torch \ + MPLCONFIGDIR=$WORK/.mpl TOKENIZERS_PARALLELISM=false && \ + export CUDA_VISIBLE_DEVICES=\$(( SLURM_LOCALID % $GPUS )) && \ + python data_pipeline/decode_wansyn32k.py \ + --parquet-dir $PARQUET_DIR --vae-dir $VAE_DIR --out-dir $OUT_DIR \ + --num-shards $NUM_SHARDS --shard \$SLURM_PROCID'" diff --git a/data_pipeline/run_segment_synth_slurm.sh b/data_pipeline/run_segment_synth_slurm.sh new file mode 100755 index 000000000..0faec2a57 --- /dev/null +++ b/data_pipeline/run_segment_synth_slurm.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# SAM 2.1-b+ segmentation of frame-0 for our synth mp4s. Adds object_ids + +# track_weights to each tracks .npz. Idempotent (skips npz already labeled). +set -euo pipefail +WORK=/mnt/lustre/vlm-s4duan +: "${DATA_DIR:?}" +NODES=${NODES:-4}; GPUS=${GPUS:-4}; PROCS_PER_GPU=${PROCS_PER_GPU:-2} +TASKS_PER_NODE=$(( GPUS * PROCS_PER_GPU )) +CPUS_PER_TASK=$(( 128 / TASKS_PER_NODE )); [ "$CPUS_PER_TASK" -lt 1 ] && CPUS_PER_TASK=1 +NUM_SHARDS=$(( NODES * TASKS_PER_NODE )) + +mkdir -p "$WORK/logs" +echo "nodes=$NODES gpus/node=$GPUS procs/gpu=$PROCS_PER_GPU -> $NUM_SHARDS workers" + +sbatch -N "$NODES" --gres=gpu:$GPUS --ntasks-per-node=$TASKS_PER_NODE --exclusive \ + --cpus-per-task=$CPUS_PER_TASK --mem=0 -t 4:00:00 -J segment_synth \ + --chdir="$WORK/FastVideo" -o "$WORK/logs/segment_synth_%j.out" -e "$WORK/logs/segment_synth_%j.out" \ + --wrap "srun --chdir=$WORK/FastVideo \ + bash -lc 'source .venv/bin/activate && \ + export TORCH_HOME=$WORK/.torch HF_HOME=$WORK/.hf HOME=$WORK MPLCONFIGDIR=$WORK/.mpl \ + YOLO_CONFIG_DIR=$WORK/.ultralytics TOKENIZERS_PARALLELISM=false && \ + export CUDA_VISIBLE_DEVICES=\$(( SLURM_LOCALID % $GPUS )) && \ + python data_pipeline/segment_tracks.py \ + --data-dir $DATA_DIR \ + --num-shards $NUM_SHARDS --shard \$SLURM_PROCID'" diff --git a/data_pipeline/run_tracks_synth_slurm.sh b/data_pipeline/run_tracks_synth_slurm.sh new file mode 100755 index 000000000..be4f478fe --- /dev/null +++ b/data_pipeline/run_tracks_synth_slurm.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# CoTracker extraction for our synth mp4s. Same pattern as run_tracks_slurm.sh but +# with the aspect/lowres filters disabled (synth videos are exactly 720x1280) and +# H/W set to our training resolution 480x832. +set -euo pipefail +WORK=/mnt/lustre/vlm-s4duan +: "${VIDEO_LIST:?}"; : "${OUT_DIR:?}" +NODES=${NODES:-4}; GPUS=${GPUS:-4}; PROCS_PER_GPU=${PROCS_PER_GPU:-2} +FPS=${FPS:-24}; NUM_FRAMES=${NUM_FRAMES:-121}; GRID=${GRID:-50} +HEIGHT=${HEIGHT:-480}; WIDTH=${WIDTH:-832} +TASKS_PER_NODE=$(( GPUS * PROCS_PER_GPU )) +CPUS_PER_TASK=$(( 128 / TASKS_PER_NODE )); [ "$CPUS_PER_TASK" -lt 1 ] && CPUS_PER_TASK=1 + +mkdir -p "$OUT_DIR" "$WORK/logs" +echo "nodes=$NODES gpus/node=$GPUS procs/gpu=$PROCS_PER_GPU -> $((NODES*TASKS_PER_NODE)) workers at ${HEIGHT}x${WIDTH}" + +sbatch -N "$NODES" --gres=gpu:$GPUS --ntasks-per-node=$TASKS_PER_NODE --exclusive \ + --cpus-per-task=$CPUS_PER_TASK --mem=0 -t 6:00:00 -J cotracker_synth \ + --chdir="$WORK/FastVideo" -o "$WORK/logs/cotracker_synth_%j.out" -e "$WORK/logs/cotracker_synth_%j.out" \ + --wrap "srun --chdir=$WORK/FastVideo \ + bash -lc 'source .venv/bin/activate && export TORCH_HOME=$WORK/.torch TOKENIZERS_PARALLELISM=false && \ + python data_pipeline/extract_tracks_mp.py \ + --video-list $VIDEO_LIST --out-dir $OUT_DIR \ + --gpus-per-node $GPUS --fps $FPS --num-frames $NUM_FRAMES --grid-size $GRID \ + --height $HEIGHT --width $WIDTH --min-height 0 --aspect-tol 10.0'" diff --git a/data_pipeline/segment_tracks.py b/data_pipeline/segment_tracks.py index 7e855063a..6cfe2efed 100644 --- a/data_pipeline/segment_tracks.py +++ b/data_pipeline/segment_tracks.py @@ -90,6 +90,8 @@ def main() -> None: help="drop masks smaller than this fraction of the frame (fights over-segmentation)") p.add_argument("--max-masks", type=int, default=0, help="keep only the N largest masks (0 = keep all)") p.add_argument("--limit", type=int, default=None) + p.add_argument("--shard", type=int, default=0, help="this shard index (0-based)") + p.add_argument("--num-shards", type=int, default=1, help="total shards for parallel processing") args = p.parse_args() # FastSAM-*.pt -> FastSAM loader; everything else (SAM/SAM2/SAM2.1/MobileSAM) -> SAM loader. @@ -103,6 +105,9 @@ def main() -> None: items = json.loads(manifest_path.read_text()) if manifest_path.exists() else [] if args.limit: items = items[:args.limit] + if args.num_shards > 1: + items = items[args.shard::args.num_shards] + print(f"[seg] shard {args.shard}/{args.num_shards}: processing {len(items)} items", flush=True) n_ok = 0 for k, item in enumerate(items, 1): vpath = args.data_dir / args.videos_subdir / item["path"] diff --git a/data_pipeline/two_track_val.py b/data_pipeline/two_track_val.py new file mode 100644 index 000000000..7a97469cb --- /dev/null +++ b/data_pipeline/two_track_val.py @@ -0,0 +1,325 @@ +#!/usr/bin/env python3 +"""2-track validation ablation. + +For each val sample, keep exactly 2 tracks (one per top-motion object) and zero +the rest. This is the extreme sparse-inference case: what the user actually gets +when they draw 2 traces at inference. Compares checkpoints side-by-side. + +Uses paper CFG (wt=3.0, wm=1.5) unless overridden. Runs `_sample`-style denoise +matching track_validation.py exactly. + +Usage: + python data_pipeline/two_track_val.py \ + --model-dir /mnt/lustre/vlm-s4duan/exports/merged_bias_ckpt4800 \ + --yaml examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_paperLR.yaml \ + --val-parquet /mnt/lustre/vlm-s4duan/val_examples_mixed/combined_parquet_dataset \ + --wandb-run-name twotrack_mb4800 +""" +from __future__ import annotations + +import argparse +import glob +import os +import sys +import time +from pathlib import Path + +import imageio.v2 as imageio +import numpy as np +import pyarrow.parquet as pq +import torch + +sys.path.insert(0, os.path.join(os.path.dirname(__file__))) +import trackwan_infer as twi # noqa: E402 +from fastvideo.forward_context import set_forward_context # noqa: E402 +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( # noqa: E402 + FlowMatchEulerDiscreteScheduler, +) + + +def load_val_samples(parquet_dir: str, text_len: int) -> list: + from fastvideo.dataset.dataloader.schema import pyarrow_schema_i2v_track + from fastvideo.dataset.utils import collate_rows_from_parquet_schema + fs = sorted(glob.glob(os.path.join(parquet_dir, "**", "*.parquet"), recursive=True)) + if not fs: + raise FileNotFoundError(parquet_dir) + rows = [] + for f in fs: + rows.extend(pq.read_table(f).to_pylist()) + batch = collate_rows_from_parquet_schema(rows, pyarrow_schema_i2v_track, + text_padding_length=int(text_len), cfg_rate=0.0) + infos = batch.get("info_list") or [{} for _ in rows] + return [{ + "id": str(infos[i].get("id", f"clip{i:03d}")) if i < len(infos) else f"clip{i:03d}", + "file_name": rows[i].get("file_name", ""), + "text_embedding": batch["text_embedding"][i:i + 1].clone(), + "text_attention_mask": batch["text_attention_mask"][i:i + 1].clone(), + "first_frame_latent": batch["first_frame_latent"][i:i + 1].clone(), + "clip_feature": batch["clip_feature"][i:i + 1].clone(), + "vae_latent": batch["vae_latent"][i:i + 1].clone(), + "track_points": batch["track_points"][i:i + 1].clone(), # [1,T,N,2] normalised + "track_visibility": batch["track_visibility"][i:i + 1].clone(), + "object_ids": batch["object_ids"][i:i + 1].clone() if "object_ids" in batch else None, + "caption": str(infos[i].get("caption", "") if i < len(infos) else ""), + } for i in range(len(rows))] + + +def pick_two_tracks(sample: dict, seed: int, num_bg: int = 0, + num_fg: int = 2) -> tuple[torch.Tensor, torch.Tensor, list]: + """Return (tp_2, tv_2, info) — ``num_fg`` foreground tracks + ``num_bg`` bg anchor tracks. + + Foreground picks require the track to STAY IN FRAME [0,1] for all 121 frames — no + point in showing traces that fly off the canvas. + """ + tp = sample["track_points"].float() # [1,T,N,2] + tv = sample["track_visibility"].float() # [1,T,N] + oid = sample["object_ids"] # [1,N] + if oid is None: + raise ValueError("no object_ids in val sample") + oid = oid[0].to(torch.int64) # [N] + + tp0 = tp[0] # [T, N, 2] + N = tp0.shape[1] + motion_per_track = torch.linalg.norm(tp0[-1] - tp0[0], dim=-1) # [N] + + # Which tracks stay fully within [0,1] on both x and y across ALL frames? + in_frame = ((tp0 >= 0.0) & (tp0 <= 1.0)).all(dim=-1).all(dim=0) # [N] + + # Rank object_ids by max in-frame within-object track motion. + fg_oids = torch.unique(oid[oid >= 0]) + obj_best = [] # (oid, best_track_idx, best_motion) + for o in fg_oids.tolist(): + idxs = ((oid == o) & in_frame & (tv[0, 0] > 0.5)).nonzero(as_tuple=True)[0] + if idxs.numel() == 0: + continue + m = motion_per_track[idxs] + best = idxs[int(torch.argmax(m).item())] + obj_best.append((o, int(best.item()), float(m.max().item()))) + obj_best.sort(key=lambda x: -x[2]) + picked_tracks = [idx for _, idx, _ in obj_best[:num_fg]] + info = [{"oid": int(o), "motion_max": mm, "track_idx": pk} + for (o, pk, mm) in obj_best[:num_fg]] + rng = np.random.RandomState(seed) + vis0 = tv[0, 0] # [N] visibility at frame 0 + + # Add num_bg background anchor tracks — near-static (bottom of motion), visible at frame 0. + bg_picked = [] + if num_bg > 0: + bg_pool = ((oid == -1) & (vis0 > 0.5)).nonzero(as_tuple=True)[0] + if bg_pool.numel() > 0: + # Sort bg tracks by motion (ascending), take lowest-motion ones and pick uniformly + bg_motion = motion_per_track[bg_pool] + n_pool = int(min(bg_pool.numel(), max(num_bg * 20, num_bg))) + top_idxs = torch.argsort(bg_motion)[:n_pool] + candidate = bg_pool[top_idxs] + if candidate.numel() >= num_bg: + sel = rng.choice(candidate.numel(), num_bg, replace=False) + bg_picked = [int(candidate[k].item()) for k in sel] + else: + bg_picked = [int(k.item()) for k in candidate] + for k in bg_picked: + info.append({"oid": -1, "motion_max": float(motion_per_track[k].item()), "track_idx": k}) + + # Build zeroed visibility: only picked tracks visible. + tv_new = torch.zeros_like(tv) + for k in picked_tracks + bg_picked: + tv_new[0, :, k] = tv[0, :, k] + + return tp, tv_new, info + + +@torch.no_grad() +def sample_denoise(model, sample: dict, tp: torch.Tensor, tv: torch.Tensor, + *, w_text: float, w_motion: float, num_steps: int, seed: int) -> np.ndarray: + """Reproduce track_validation._sample denoise EXACTLY.""" + device = model.device + dtype = torch.bfloat16 + + ff = sample["first_frame_latent"].to(device, dtype) + txt = sample["text_embedding"].to(device, dtype) + mask = sample["text_attention_mask"].to(device, dtype) + clip = sample["clip_feature"].to(device, dtype) + tp_d = tp.to(device, dtype) + tv_d = tv.to(device, dtype) + cond20 = model._build_i2v_cond_concat(ff) + + _, _, T, H, W = ff.shape + gen = torch.Generator(device="cpu").manual_seed(int(seed)) + latents = torch.randn((1, 16, T, H, W), generator=gen, dtype=torch.float32).to(device) + sched = FlowMatchEulerDiscreteScheduler(shift=float(model.timestep_shift)) + sched.set_timesteps(int(num_steps), device=device) + + cfg_on = (w_text != 1.0) or (w_motion != 1.0) + txt_null = torch.zeros_like(txt) if cfg_on else None + + def _fwd(text_e, tp_e, tv_e, mi, ts_): + with torch.autocast(device.type, dtype=dtype), \ + set_forward_context(current_timestep=ts_, attn_metadata=None): + return model.transformer(hidden_states=mi, encoder_hidden_states=text_e, + encoder_attention_mask=mask, timestep=ts_, + encoder_hidden_states_image=clip, + track_points=tp_e, track_visibility=tv_e, + return_dict=False) + + for tt in sched.timesteps: + mi = torch.cat([latents.to(dtype), cond20], dim=1) + ts_ = tt.reshape(1).to(device, dtype) + v_full = _fwd(txt, tp_d, tv_d, mi, ts_) + if not cfg_on: + v = v_full + elif w_text != 1.0 and w_motion == 1.0: + v_no_text = _fwd(txt_null, tp_d, tv_d, mi, ts_) + v = v_no_text + w_text * (v_full - v_no_text) + elif w_text == 1.0 and w_motion != 1.0: + v_no_motion = _fwd(txt, None, None, mi, ts_) + v = v_no_motion + w_motion * (v_full - v_no_motion) + else: + v_no_text = _fwd(txt_null, tp_d, tv_d, mi, ts_) + v_no_motion = _fwd(txt, None, None, mi, ts_) + alpha = w_text / max(w_text + w_motion, 1e-6) + v_base = alpha * v_no_text + (1.0 - alpha) * v_no_motion + v = v_base + w_text * (v_full - v_no_text) + w_motion * (v_full - v_no_motion) + latents = sched.step(v.float(), tt, latents.float(), return_dict=False)[0] + + return twi.decode_to_pixels(model, latents) + + +def draw_track_overlay(video: np.ndarray, tp: torch.Tensor, tv: torch.Tensor, + info: list | None = None) -> np.ndarray: + """Draw kept tracks on the video. Foreground = red/blue, background anchors = gray.""" + from PIL import Image, ImageDraw + T, H, W, _ = video.shape + tp_np = tp[0].cpu().numpy() # [T, N, 2] + tv_np = tv[0].cpu().numpy() # [T, N] + N = tp_np.shape[1] + kept = [k for k in range(N) if tv_np[:, k].max() > 0.5] + fg_colors = [(255, 60, 60), (60, 200, 255)] + bg_color = (180, 180, 180) + # figure out which are fg vs bg by looking at info if given + is_bg_by_idx = {} + if info is not None: + for entry in info: + is_bg_by_idx[entry["track_idx"]] = (entry["oid"] == -1) + out = [] + fg_seen = 0 + for t in range(T): + img = Image.fromarray(np.ascontiguousarray(video[t])) + dr = ImageDraw.Draw(img) + fg_seen = 0 + for k in kept: + if tv_np[t, k] <= 0.5: + continue + x = float(tp_np[t, k, 0] * W) + y = float(tp_np[t, k, 1] * H) + if is_bg_by_idx.get(k, False): + col = bg_color + r = 3 + else: + col = fg_colors[fg_seen % len(fg_colors)] + fg_seen += 1 + r = 5 + dr.ellipse([x - r, y - r, x + r, y + r], fill=col, outline=(255, 255, 255)) + out.append(np.asarray(img)) + return np.stack(out) + + +def main() -> None: + p = argparse.ArgumentParser() + p.add_argument("--model-dir", required=True) + p.add_argument("--yaml", required=True) + p.add_argument("--val-parquet", required=True) + p.add_argument("--wandb-project", default="wantrack-bidir") + p.add_argument("--wandb-run-name", required=True) + p.add_argument("--out-dir", default="/mnt/lustre/vlm-s4duan/two_track_val") + p.add_argument("--w-text", type=float, default=3.0) + p.add_argument("--w-motion", type=float, default=1.5) + p.add_argument("--steps", type=int, default=30) + p.add_argument("--seed", type=int, default=1000) + p.add_argument("--fps", type=int, default=24) + p.add_argument("--num-fg-tracks", type=int, default=2, + help="Number of foreground (moving object) tracks to keep.") + p.add_argument("--num-bg-tracks", type=int, default=0, + help="Extra static bg anchor tracks in addition to the fg ones.") + p.add_argument("--swap-text-from", type=int, default=None, + help="For every sample, replace text_embedding+mask with THIS sample idx's text.") + p.add_argument("--swap-clip-from", type=int, default=None, + help="For every sample, replace clip_feature with THIS sample idx's clip.") + p.add_argument("--swap-ff-from", type=int, default=None, + help="For every sample, replace first_frame_latent with THIS sample idx's ff.") + args = p.parse_args() + + out_dir = Path(args.out_dir) / args.wandb_run_name + out_dir.mkdir(parents=True, exist_ok=True) + + import wandb + + print(f"[2trk] loading model {args.model_dir}", flush=True) + model, tc = twi.load_trackwan(args.model_dir, args.yaml) + text_len = int(getattr(tc.data, "text_padding_length", 256)) + print(f"[2trk] text_len={text_len}", flush=True) + + print(f"[2trk] loading val samples", flush=True) + samples = load_val_samples(args.val_parquet, text_len) + print(f"[2trk] {len(samples)} val samples", flush=True) + + run = wandb.init(project=args.wandb_project, name=args.wandb_run_name, + config={"model_dir": args.model_dir, "w_text": args.w_text, + "w_motion": args.w_motion, "steps": args.steps}) + + for i, s in enumerate(samples): + sid = s["id"] or f"clip{i:03d}" + print(f"\n[2trk] === sample {i}: {sid} ({s['file_name'][:40]}) ===", flush=True) + # Optional cross-sample field swaps (isolate which condition drives quality) + if args.swap_text_from is not None: + src = samples[args.swap_text_from] + s = {**s, "text_embedding": src["text_embedding"].clone(), + "text_attention_mask": src["text_attention_mask"].clone()} + print(f"[2trk] SWAPPED text from sample {args.swap_text_from}", flush=True) + if args.swap_clip_from is not None: + src = samples[args.swap_clip_from] + s = {**s, "clip_feature": src["clip_feature"].clone()} + print(f"[2trk] SWAPPED clip from sample {args.swap_clip_from}", flush=True) + if args.swap_ff_from is not None: + src = samples[args.swap_ff_from] + s = {**s, "first_frame_latent": src["first_frame_latent"].clone()} + print(f"[2trk] SWAPPED first_frame_latent from sample {args.swap_ff_from}", flush=True) + tp2, tv2, info = pick_two_tracks(s, seed=args.seed + i, + num_fg=args.num_fg_tracks, num_bg=args.num_bg_tracks) + print(f"[2trk] picked: {info}", flush=True) + + t0 = time.time() + frames = sample_denoise(model, s, tp2, tv2, w_text=args.w_text, w_motion=args.w_motion, + num_steps=args.steps, seed=args.seed + i) + dt = time.time() - t0 + print(f"[2trk] denoise {dt:.1f}s -> frames {frames.shape}", flush=True) + + overlay = draw_track_overlay(frames, tp2, tv2, info) + fn_raw = out_dir / f"{sid}_gen.mp4" + fn_ov = out_dir / f"{sid}_gen_overlay.mp4" + imageio.mimsave(str(fn_raw), frames, fps=args.fps, macro_block_size=1) + imageio.mimsave(str(fn_ov), overlay, fps=args.fps, macro_block_size=1) + + # GT for reference + try: + ref = twi.decode_reference(model, s["vae_latent"].to(model.device, torch.bfloat16)) + ref_ov = draw_track_overlay(ref, tp2, tv2, info) + fn_ref = out_dir / f"{sid}_gt_overlay.mp4" + imageio.mimsave(str(fn_ref), ref_ov, fps=args.fps, macro_block_size=1) + wandb.log({ + f"sample_{i:02d}_gen": wandb.Video(str(fn_ov), fps=args.fps, format="mp4"), + f"sample_{i:02d}_gt": wandb.Video(str(fn_ref), fps=args.fps, format="mp4"), + "sample": i, "caption": s["caption"][:200], "picked": str(info), + }) + except Exception as e: + print(f"[2trk] gt fail: {e}", flush=True) + wandb.log({ + f"sample_{i:02d}_gen": wandb.Video(str(fn_ov), fps=args.fps, format="mp4"), + "sample": i, "caption": s["caption"][:200], "picked": str(info), + }) + + wandb.finish() + print("[2trk] DONE", flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/inference/gradio/trackwan/app_multi_trace.py b/examples/inference/gradio/trackwan/app_multi_trace.py index 375577689..e9b3bd15e 100644 --- a/examples/inference/gradio/trackwan/app_multi_trace.py +++ b/examples/inference/gradio/trackwan/app_multi_trace.py @@ -1,17 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 -"""TrackWan multi-trace demo — pick a preprocessed openvid clip, draw multiple traces -with mouse (Space to record), generate. +"""TrackWan interactive rollout demo. + +Two point-input modes on the same canvas: + - **Trace mode** — hover, press *Space* to record a 5-second (121-frame) drag. + - **Anchor mode** — click once to drop a static point (a track whose (x,y) is + constant across all 121 frames). Design (matches training-time validation exactly): - Loads conditioning (text_embedding, clip_feature, first_frame_latent) DIRECTLY from - the preprocessed openvid parquet — no on-the-fly encoding. Same path as + the preprocessed openvid parquet. Same path as ``fastvideo/train/callbacks/track_validation.py``. - First frame image shown to the user is decoded from the VAE latent (``decode_reference``). - - Traces: press Space to record a 5s (121-frame) trajectory. Multi-trace stacks. - - Background points: sampled from SAM foreground-avoiding pixels of frame-0. Click a dot - to delete just it; buttons for resample / clear all. - - Generation: uses the sample's precomputed conditioning + user's track_points, then runs - the SAME motion-CFG denoise loop as ``TrackValidationCallback._sample``. + - Both traces and anchors are packed into a single ``track_points[T,N,2]`` tensor. + Anchors get their (x,y) broadcast across T=121; traces are linearly interpolated. + - Generation: the SAME motion-CFG denoise loop as ``TrackValidationCallback._sample`` + (MotionStream Eq. 2, wt=3.0, wm=1.5 by default). Launch:: @@ -42,53 +45,10 @@ import numpy as np from PIL import Image REPO = Path(__file__).resolve().parents[4] -SAM_WEIGHT_CANDIDATES = [ - "/mnt/lustre/vlm-s4duan/models/seg/sam2.1_b.pt", - "/mnt/lustre/vlm-s4duan/FastVideo/sam2.1_b.pt", - "/mnt/lustre/vlm-s4duan/models/seg/sam2_b.pt", -] NUM_FRAMES = 121 FPS = 24 -# ============================================================================= -# SAM -# ============================================================================= -def load_sam(weight: str | None = None): - from ultralytics import SAM - if weight is None: - for w in SAM_WEIGHT_CANDIDATES: - if os.path.exists(w): - weight = w - break - if weight is None: - raise FileNotFoundError(f"No SAM weight found in {SAM_WEIGHT_CANDIDATES}") - return SAM(weight) - - -def sam_foreground_mask(sam, img: np.ndarray, imgsz: int = 640) -> np.ndarray: - res = sam(img, imgsz=imgsz, conf=0.4, iou=0.9, retina_masks=True, verbose=False) - if not res or res[0].masks is None: - return np.zeros(img.shape[:2], bool) - m = res[0].masks.data.cpu().numpy().astype(bool) - if m.shape[0] and m.shape[1:] != img.shape[:2]: - import torch - t = torch.from_numpy(m.astype(np.uint8))[None].float() - t = torch.nn.functional.interpolate(t, size=img.shape[:2], mode="nearest")[0] - m = t.numpy().astype(bool) - return m.any(0) if m.shape[0] > 0 else np.zeros(img.shape[:2], bool) - - -def sample_bg_points_normalized(fg_mask: np.ndarray, n: int, seed: int = 0) -> list: - ys, xs = np.where(~fg_mask) - if len(xs) == 0: - return [] - rng = np.random.default_rng(seed) - idx = rng.choice(len(xs), size=min(n, len(xs)), replace=False) - H, W = fg_mask.shape - return [[float(xs[i]) / W, float(ys[i]) / H] for i in idx] - - # ============================================================================= # frame b64 + track construction # ============================================================================= @@ -98,15 +58,23 @@ def img_to_b64(arr: np.ndarray) -> str: return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode() -def build_track_points(traces_json: str, bg_pts: list) -> tuple[np.ndarray, np.ndarray, int, int]: - """List-of-normalized-trajectories → (track_points[T,N,2], track_visibility[T,N], n_tr, n_bg).""" +def build_track_points(traces_json: str, anchors_json: str) -> tuple[np.ndarray, np.ndarray, int, int]: + """(traces, anchors) → (track_points[T,N,2], track_visibility[T,N], n_tr, n_an). + + Each trace is a list of normalized (x,y) samples at ~24 Hz over 5 s → interpolated to T=121. + Each anchor is a single normalized (x,y) → broadcast across all T frames. + """ T = NUM_FRAMES try: traces = json.loads(traces_json) if traces_json else [] except Exception: traces = [] - n_tr, n_bg = len(traces), len(bg_pts) - N = n_tr + n_bg + try: + anchors = json.loads(anchors_json) if anchors_json else [] + except Exception: + anchors = [] + n_tr, n_an = len(traces), len(anchors) + N = n_tr + n_an if N == 0: return np.zeros((T, 0, 2), np.float32), np.zeros((T, 0), np.float32), 0, 0 tp = np.zeros((T, N, 2), np.float32) @@ -122,14 +90,14 @@ def build_track_points(traces_json: str, bg_pts: list) -> tuple[np.ndarray, np.n tgt = np.linspace(0, arr.shape[0] - 1, T) tp[:, i, 0] = np.interp(tgt, src, arr[:, 0]) tp[:, i, 1] = np.interp(tgt, src, arr[:, 1]) - for j, (bx, by) in enumerate(bg_pts): - tp[:, n_tr + j, 0] = bx - tp[:, n_tr + j, 1] = by - return tp, tv, n_tr, n_bg + for j, (ax, ay) in enumerate(anchors): + tp[:, n_tr + j, 0] = ax + tp[:, n_tr + j, 1] = ay + return tp, tv, n_tr, n_an # ============================================================================= -# canvas JS (multi-trace + click-to-delete bg points) +# canvas JS (traces via Space + anchors via click, mode radio decides) # ============================================================================= def canvas_js(num_frames: int, fps: int) -> str: palette = ["#ff3c3c", "#3cd23c", "#3c78ff", "#ffc83c", "#c83cc8", "#3cdcdc", "#ff821e", "#9696ff"] @@ -138,32 +106,40 @@ def canvas_js(num_frames: int, fps: int) -> str: " const ctx = cvs.getContext('2d');\n" f" const PAL = {json.dumps(palette)};\n" f" const N = {num_frames}; const DT = {int(1000 / fps)};\n" - " const st = { img:new Image(), mouse:[0.5,0.5], rec:false, current:[], traces:[], bg:[] };\n" + " const st = { img:new Image(), mouse:[0.5,0.5], rec:false, current:[], " + "traces:[], anchors:[], mode:'trace' };\n" + " window._st = st;\n" " cvs.width = 832; cvs.height = 480;\n" " function commit_traces(){\n" " const tb = document.querySelector('#traces_json textarea');\n" " if (tb) { tb.value = JSON.stringify(st.traces); tb.dispatchEvent(new Event('input',{bubbles:true})); }\n" " }\n" - " function commit_bg(){\n" - " const tb = document.querySelector('#bg_pts_json textarea');\n" - " if (tb) { tb.value = JSON.stringify(st.bg); tb.dispatchEvent(new Event('input',{bubbles:true})); }\n" + " function commit_anchors(){\n" + " const tb = document.querySelector('#anchors_json textarea');\n" + " if (tb) { tb.value = JSON.stringify(st.anchors); tb.dispatchEvent(new Event('input',{bubbles:true})); }\n" " }\n" " function draw(){\n" " ctx.clearRect(0,0,cvs.width,cvs.height);\n" " if (st.img.complete && st.img.width) ctx.drawImage(st.img, 0, 0, cvs.width, cvs.height);\n" - " st.bg.forEach(p => { const x=p[0]*cvs.width, y=p[1]*cvs.height;\n" - " ctx.beginPath(); ctx.arc(x,y,3,0,7); ctx.fillStyle='#c8c8c8'; ctx.fill();\n" - " ctx.strokeStyle='#000'; ctx.stroke(); });\n" + " // anchors: solid gray dots with dark ring\n" + " st.anchors.forEach((p, i) => { const x=p[0]*cvs.width, y=p[1]*cvs.height;\n" + " ctx.beginPath(); ctx.arc(x,y,7,0,7); ctx.fillStyle='#888'; ctx.fill();\n" + " ctx.strokeStyle='#000'; ctx.lineWidth=2; ctx.stroke();\n" + " ctx.fillStyle='#fff'; ctx.font='bold 10px sans-serif'; ctx.textAlign='center'; ctx.textBaseline='middle';\n" + " ctx.fillText('A'+(i+1), x, y);\n" + " });\n" + " // traces: colored polyline with start/end markers\n" " st.traces.forEach((traj, i) => {\n" " const color = PAL[i % PAL.length];\n" " ctx.strokeStyle=color; ctx.lineWidth=3; ctx.beginPath();\n" " traj.forEach((p, j) => { const x=p[0]*cvs.width, y=p[1]*cvs.height; j?ctx.lineTo(x,y):ctx.moveTo(x,y); });\n" " ctx.stroke();\n" " const s = traj[0], e = traj[traj.length-1];\n" - " ctx.fillStyle=color; ctx.beginPath(); ctx.arc(s[0]*cvs.width, s[1]*cvs.height, 6, 0, 7); ctx.fill(); ctx.strokeStyle='#000'; ctx.stroke();\n" + " ctx.fillStyle=color; ctx.beginPath(); ctx.arc(s[0]*cvs.width, s[1]*cvs.height, 6, 0, 7); ctx.fill(); ctx.strokeStyle='#000'; ctx.lineWidth=1; ctx.stroke();\n" " ctx.fillStyle='#ff3333'; ctx.beginPath(); ctx.arc(e[0]*cvs.width, e[1]*cvs.height, 6, 0, 7); ctx.fill(); ctx.strokeStyle='#000'; ctx.stroke();\n" - " ctx.fillStyle='#000'; ctx.font='13px sans-serif'; ctx.fillText('#'+(i+1), s[0]*cvs.width+8, s[1]*cvs.height-8);\n" + " ctx.fillStyle='#000'; ctx.font='13px sans-serif'; ctx.textAlign='left'; ctx.textBaseline='alphabetic'; ctx.fillText('T'+(i+1), s[0]*cvs.width+8, s[1]*cvs.height-8);\n" " });\n" + " // live-drawing current trace\n" " if (st.rec && st.current.length) {\n" " const color = PAL[st.traces.length % PAL.length];\n" " ctx.strokeStyle=color; ctx.lineWidth=3; ctx.beginPath();\n" @@ -180,28 +156,34 @@ def canvas_js(num_frames: int, fps: int) -> str: " const r = cvs.getBoundingClientRect();\n" " const cx = (e.clientX - r.left) * (cvs.width / r.width);\n" " const cy = (e.clientY - r.top) * (cvs.height / r.height);\n" - " let idx = -1, best = 12*12;\n" - " st.bg.forEach((p, i) => { const dx = p[0]*cvs.width - cx, dy = p[1]*cvs.height - cy;\n" - " const d2 = dx*dx + dy*dy; if (d2 < best) { best = d2; idx = i; } });\n" - " if (idx >= 0) { st.bg.splice(idx, 1); commit_bg(); draw(); }\n" + " if (st.mode === 'anchor') {\n" + " // toggle delete-if-near, else place\n" + " let idx = -1, best = 14*14;\n" + " st.anchors.forEach((p, i) => { const dx = p[0]*cvs.width - cx, dy = p[1]*cvs.height - cy;\n" + " const d2 = dx*dx + dy*dy; if (d2 < best) { best = d2; idx = i; } });\n" + " if (idx >= 0) st.anchors.splice(idx, 1);\n" + " else st.anchors.push([cx / cvs.width, cy / cvs.height]);\n" + " commit_anchors(); draw();\n" + " }\n" " });\n" + " window._setmode = (m) => { st.mode = m || 'trace'; };\n" " window._setbg = (b64) => {\n" " if (!b64) return;\n" " const im = new Image();\n" " im.onload = () => { const ar = im.width / im.height; cvs.width = 832; cvs.height = Math.round(832 / ar); st.img = im; draw(); };\n" " im.src = b64;\n" " };\n" - " window._setbg_pts = (j) => { try { st.bg = JSON.parse(j) || []; } catch(e) { st.bg = []; } draw(); };\n" - " window._clear = () => {\n" + " window._clear_traces = () => {\n" " st.traces = []; st.current = []; st.rec = false; commit_traces(); draw();\n" - " const stt = document.getElementById('recstatus'); if (stt) stt.textContent = 'cleared. Space for trace #1';\n" + " const stt = document.getElementById('recstatus'); if (stt) stt.textContent = 'traces cleared. Space for trace #1';\n" " };\n" - " window._undo = () => {\n" + " window._undo_trace = () => {\n" " if (st.traces.length) st.traces.pop();\n" " commit_traces(); draw();\n" - " const stt = document.getElementById('recstatus'); if (stt) stt.textContent = st.traces.length + ' committed. Space for another';\n" + " const stt = document.getElementById('recstatus'); if (stt) stt.textContent = st.traces.length + ' trace(s). Space for another';\n" " };\n" - " window._clearbg = () => { st.bg = []; commit_bg(); draw(); };\n" + " window._clear_anchors = () => { st.anchors = []; commit_anchors(); draw(); };\n" + " window._undo_anchor = () => { if (st.anchors.length) st.anchors.pop(); commit_anchors(); draw(); };\n" " window._startRec = () => {\n" " if (st.rec) return;\n" " st.rec = true; st.current = []; let n = 0;\n" @@ -221,6 +203,7 @@ def canvas_js(num_frames: int, fps: int) -> str: " if (e.code !== 'Space' || e.repeat) return;\n" " const tag = document.activeElement && document.activeElement.tagName;\n" " if (tag === 'INPUT' || tag === 'TEXTAREA') return;\n" + " if (st.mode !== 'trace') return; // Space only records in trace mode\n" " e.preventDefault(); window._startRec();\n" " });\n" " draw();\n" @@ -234,35 +217,42 @@ def build_ui(state): canvas_html = ("
" "" "
" - "pick a clip, click Load frame, then hover and press Space to record trace #1
" + "pick a clip, click Load frame, then choose a mode below
" "") labels = state["labels"] - with gr.Blocks(title="TrackWan multi-trace demo") as demo: + with gr.Blocks(title="TrackWan interactive rollout") as demo: gr.Markdown( - "## TrackWan — multi-trace generation (parquet-conditioned)\n" - "1. Pick a preprocessed openvid clip → **Load frame**.\n" - "2. Hover over the frame and press **Space** to record a 5-second trace.\n" - "3. Repeat Space for more traces. Optionally include background points. **Generate**.") + "## TrackWan — interactive rollout demo\n" + "1. Pick a preprocessed clip → **Load frame**.\n" + "2. Choose a mode:\n" + " - **Trace** — hover on the frame and press **Space** to record a 5 s drag.\n" + " - **Anchor** — click on the frame to drop a *static* point. Click an existing anchor to remove it.\n" + "3. **Generate**. Denoise formula matches training's validation exactly (MotionStream Eq. 2).") with gr.Row(): with gr.Column(scale=1): clip = gr.Dropdown(labels, value=labels[0] if labels else None, label="Preprocessed clip") caption = gr.Textbox(label="Caption (from parquet)", interactive=False, lines=2) - load_btn = gr.Button("Load frame (runs SAM)", variant="primary") - sam_on = gr.Checkbox(True, label="Run SAM (needed for background points)") - gr.Markdown("**Background points (SAM foreground-avoiding)**") - bg_on = gr.Checkbox(True, label="Include background points") - bg_n = gr.Slider(0, 50, 20, step=1, label="Num background points") - with gr.Row(): - resample_bg_btn = gr.Button("Resample") - clear_bg_btn = gr.Button("Clear all bg pts") - gr.Markdown("_Tip: **click a bg dot** on the frame to delete just that one._") - gr.Markdown("**Trace drawing**") + load_btn = gr.Button("Load frame", variant="primary") + + gr.Markdown("**Input mode**") + mode = gr.Radio(["trace", "anchor"], value="trace", + label="Mode", + info="Trace = Space to record a moving drag. Anchor = single click for a static point.") + + gr.Markdown("**Traces** (moving)") with gr.Row(): rec_btn = gr.Button("● Record (Space)") - undo_btn = gr.Button("Undo") - clear_btn = gr.Button("Clear") - trace_info = gr.Markdown("_no traces yet._") + undo_trace_btn = gr.Button("Undo trace") + clear_trace_btn = gr.Button("Clear traces") + + gr.Markdown("**Anchors** (static)") + with gr.Row(): + undo_anchor_btn = gr.Button("Undo anchor") + clear_anchor_btn = gr.Button("Clear anchors") + + info_md = gr.Markdown("_no traces or anchors yet._") + gr.Markdown("**Generation** (denoise formula matches training's validation exactly)") seed_in = gr.Slider(0, 9999, 1000, step=1, label="Seed") steps_in = gr.Slider(10, 60, 30, step=1, label="Denoise steps") @@ -278,13 +268,12 @@ def build_ui(state): # hidden bridges ff_b64 = gr.Textbox(elem_id="ff_b64", visible=False) - bg_pts_json = gr.Textbox(elem_id="bg_pts_json", visible=False, value="[]") traces_json = gr.Textbox(elem_id="traces_json", visible=False, value="[]") + anchors_json = gr.Textbox(elem_id="anchors_json", visible=False, value="[]") # server-side per-page state st_clip_idx = gr.State(0) st_frame = gr.State(None) - st_fg = gr.State(None) # ---------------- caption on clip change ---------------- def _cap(clip_name): @@ -292,35 +281,18 @@ def build_ui(state): return i, state["samples"][i]["caption"][:400] clip.change(_cap, [clip], [st_clip_idx, caption]) - # ---------------- load frame (decode ref frame from VAE latent + SAM) ---------------- - sam_cache: dict[int, np.ndarray] = {} + # ---------------- load frame (decode ref frame from VAE latent) ---------------- first_frame_cache: dict[int, np.ndarray] = {} - def _load(clip_name, sam_on_, bg_n_, bg_on_): + def _load(clip_name): i = state["label_to_idx"].get(clip_name, 0) if i not in first_frame_cache: first_frame_cache[i] = state["decode_first_frame"](i) frame = first_frame_cache[i] - if sam_on_: - if i in sam_cache: - fg = sam_cache[i] - msg = f"_clip loaded ({state['samples'][i]['caption'][:60]}...) — SAM cached, {int(fg.sum())} fg px_" - else: - t0 = time.time() - try: - fg = sam_foreground_mask(state["sam"], frame) - sam_cache[i] = fg - except Exception as e: - fg = np.zeros(frame.shape[:2], bool) - msg = f"_clip loaded — SAM {time.time()-t0:.1f}s, {int(fg.sum())} fg px_" - else: - fg = np.zeros(frame.shape[:2], bool) - msg = "_clip loaded (SAM skipped)_" - bg = sample_bg_points_normalized(fg, int(bg_n_)) if (bg_on_ and fg.any()) else [] - return i, frame, fg, img_to_b64(frame), json.dumps(bg), msg + msg = f"_clip loaded ({state['samples'][i]['caption'][:60]}...)_" + return i, frame, img_to_b64(frame), msg - load_btn.click(_load, [clip, sam_on, bg_n, bg_on], - [st_clip_idx, st_frame, st_fg, ff_b64, bg_pts_json, status]) + load_btn.click(_load, [clip], [st_clip_idx, st_frame, ff_b64, status]) # preload first clip on startup def _preload(): @@ -328,58 +300,46 @@ def build_ui(state): if i not in first_frame_cache: first_frame_cache[i] = state["decode_first_frame"](i) frame = first_frame_cache[i] - fg = np.zeros(frame.shape[:2], bool) - return (i, state["samples"][i]["caption"][:400], frame, fg, img_to_b64(frame), - "[]", "_default clip preloaded (SAM off). Click **Load frame** to enable bg points._") - demo.load(_preload, None, - [st_clip_idx, caption, st_frame, st_fg, ff_b64, bg_pts_json, status]) - - # ---------------- bg controls ---------------- - state.setdefault("bg_seed", [0]) - def _update_bg(fg, on_, n_, fresh: bool = False): - if fg is None or not on_ or not fg.any(): - return "[]" - if fresh: - state["bg_seed"][0] += 1 - return json.dumps(sample_bg_points_normalized(fg, int(n_), seed=state["bg_seed"][0])) - for w in (bg_on, bg_n): - w.change(_update_bg, [st_fg, bg_on, bg_n], bg_pts_json) - resample_bg_btn.click(lambda fg, on_, n_: _update_bg(fg, on_, n_, fresh=True), - [st_fg, bg_on, bg_n], bg_pts_json) - clear_bg_btn.click(None, None, None, js="() => window._clearbg && window._clearbg()") + return (i, state["samples"][i]["caption"][:400], frame, img_to_b64(frame), + "_default clip preloaded. Draw a trace (Space) or place anchors (click)._") + demo.load(_preload, None, [st_clip_idx, caption, st_frame, ff_b64, status]) # ---------------- JS bridges ---------------- ff_b64.change(None, ff_b64, None, js="(b) => window._setbg && window._setbg(b)") - bg_pts_json.change(None, bg_pts_json, None, js="(j) => window._setbg_pts && window._setbg_pts(j)") + mode.change(None, mode, None, js="(m) => window._setmode && window._setmode(m)") rec_btn.click(None, None, None, js="() => window._startRec && window._startRec()") - undo_btn.click(None, None, None, js="() => window._undo && window._undo()") - clear_btn.click(None, None, None, js="() => window._clear && window._clear()") + undo_trace_btn.click(None, None, None, js="() => window._undo_trace && window._undo_trace()") + clear_trace_btn.click(None, None, None, js="() => window._clear_traces && window._clear_traces()") + undo_anchor_btn.click(None, None, None, js="() => window._undo_anchor && window._undo_anchor()") + clear_anchor_btn.click(None, None, None, js="() => window._clear_anchors && window._clear_anchors()") - def _tinfo(t): + def _tinfo(t, a): try: - n = len(json.loads(t) if t else []) + n_t = len(json.loads(t) if t else []) except Exception: - n = 0 - return "_no traces yet. Press **Space** to record trace #1._" if n == 0 \ - else f"_{n} trace(s) committed. **Space** for another, or **Generate**._" - traces_json.change(_tinfo, traces_json, trace_info) + n_t = 0 + try: + n_a = len(json.loads(a) if a else []) + except Exception: + n_a = 0 + if n_t == 0 and n_a == 0: + return "_no traces or anchors yet._" + return f"_{n_t} trace(s) + {n_a} anchor(s) → {n_t + n_a} total track(s)_" + traces_json.change(_tinfo, [traces_json, anchors_json], info_md) + anchors_json.change(_tinfo, [traces_json, anchors_json], info_md) # ---------------- generate (MATCHES TrackValidationCallback._sample exactly) ---------------- - def _generate(clip_idx, traces_str, bg_str, seed, steps, w_t, w_m): + def _generate(clip_idx, traces_str, anchors_str, seed, steps, w_t, w_m): def logln(log, msg): line = time.strftime("%H:%M:%S ") + msg print(line, flush=True) return (log + "\n" + line) if log else line log = "" - try: - bg = json.loads(bg_str) if bg_str else [] - except Exception: - bg = [] - tp_np, tv_np, n_tr, n_bg = build_track_points(traces_str, bg) - log = logln(log, f"clip idx={clip_idx}: {n_tr} trace(s) + {n_bg} bg point(s) → N={tp_np.shape[1]}") + tp_np, tv_np, n_tr, n_an = build_track_points(traces_str, anchors_str) + log = logln(log, f"clip idx={clip_idx}: {n_tr} trace(s) + {n_an} anchor(s) → N={tp_np.shape[1]}") if tp_np.shape[1] == 0: - return None, "_record at least one trace, or enable bg points_", log + return None, "_record at least one trace or place at least one anchor_", log ts = time.strftime("%Y%m%d_%H%M%S") req_dir = Path(state["out_dir"]) / f"req_{ts}" @@ -388,7 +348,7 @@ def build_ui(state): (req_dir / "meta.json").write_text(json.dumps({ "clip_idx": int(clip_idx), "seed": int(seed), "steps": int(steps), "w_text": float(w_t), "w_motion": float(w_m), - "num_traces": n_tr, "num_bg": n_bg, + "num_traces": n_tr, "num_anchors": n_an, }, indent=2)) log = logln(log, f"req -> {req_dir.name}, denoising ...") @@ -406,7 +366,7 @@ def build_ui(state): return None, f"_generation failed: {e}_", logln(log, f"EXC: {e}") gen_btn.click(_generate, - [st_clip_idx, traces_json, bg_pts_json, seed_in, steps_in, w_text, w_motion], + [st_clip_idx, traces_json, anchors_json, seed_in, steps_in, w_text, w_motion], [out_video, status, log_box]) demo.load(None, None, None, js=canvas_js(NUM_FRAMES, FPS)) @@ -490,10 +450,7 @@ def make_generator(model, tc, samples): # Efficient parquet loader (reads only enough files to get N rows, NOT all 4494) # ============================================================================= def _load_first_n_from_parquet(data_path: str, n: int, text_len: int) -> list: - """Grab the first N rows from a preprocessed parquet dataset (like - ``trackwan_infer.load_conditioning_from_parquet``, but stops after N rows so - 259k-row datasets don't OOM the loader). - """ + """Grab the first N rows from a preprocessed parquet dataset.""" import glob import pyarrow.parquet as pq from fastvideo.dataset.dataloader.schema import pyarrow_schema_i2v_track @@ -508,7 +465,7 @@ def _load_first_n_from_parquet(data_path: str, n: int, text_len: int) -> list: if len(rows) >= n: break sel = rows[:n] - print(f"[app] loaded {len(sel)} rows from {min(len(files), 1 + (n // max(1, (len(sel)//1))))} parquet file(s)", flush=True) + print(f"[app] loaded {len(sel)} rows from parquet", flush=True) batch = collate_rows_from_parquet_schema(sel, pyarrow_schema_i2v_track, text_padding_length=int(text_len), cfg_rate=0.0) infos = batch.get("info_list") or [{} for _ in sel] @@ -535,21 +492,16 @@ def main() -> None: ap.add_argument("--model-dir", required=True, help="diffusers-exported ckpt dir") ap.add_argument("--yaml", required=True, help="training yaml") ap.add_argument("--data-path", required=True, - help="preprocessed openvid parquet root (combined_parquet_dataset)") + help="preprocessed parquet root (combined_parquet_dataset)") ap.add_argument("--num-clips", type=int, default=30, help="how many parquet clips to preload") ap.add_argument("--out-dir", default="/mnt/lustre/vlm-s4duan/multi_trace_out") ap.add_argument("--host", default="0.0.0.0") ap.add_argument("--port", type=int, default=7864) - ap.add_argument("--sam-weight", default=None) args = ap.parse_args() os.makedirs(args.out_dir, exist_ok=True) state: dict = {"out_dir": args.out_dir} - # SAM (works on GPU when CUDA_VISIBLE_DEVICES set) - print(f"[app] loading SAM ...", flush=True) - state["sam"] = load_sam(args.sam_weight) - # Load model print(f"[app] loading model from {args.model_dir} ...", flush=True) sys.path.insert(0, str(REPO / "data_pipeline")) @@ -559,14 +511,12 @@ def main() -> None: state["model"] = model state["tc"] = tc - # Load parquet conditioning EFFICIENTLY: stop as soon as we have enough rows. - # (trackwan_infer.load_conditioning_from_parquet reads ALL 4494 files → OOM on 259k rows.) + # Load parquet conditioning efficiently print(f"[app] loading {args.num_clips} parquet clips ...", flush=True) samples = _load_first_n_from_parquet(args.data_path, args.num_clips, text_len) labels = [] label_to_idx = {} for i, s in enumerate(samples): - # Trim caption to a short unique label for the dropdown cap = (s["caption"] or f"clip_{i}").strip().replace("\n", " ") lab = f"{i:03d} {cap[:60]}" labels.append(lab) diff --git a/examples/train/acquire_nodes.sh b/examples/train/acquire_nodes.sh new file mode 100644 index 000000000..ce0c59db3 --- /dev/null +++ b/examples/train/acquire_nodes.sh @@ -0,0 +1,80 @@ +#!/usr/bin/env bash +# Acquire a single gang allocation of N nodes on the Slinky (k8s-backed) pool. +# +# Why this is not just `sbatch -N`: the k8s autoscaler destroys idle nodes, so the pool only +# contains as many nodes as there is current demand. Slurm REJECTS (does not queue) a job asking +# for more nodes than physically exist in the partition: +# "Batch job submission failed: Requested node configuration is not available" +# So we have to manufacture demand first, then grab the gang allocation: +# 1. keep N+SLACK short 1-node warmups queued -> autoscaler grows the pool +# 2. probe with `sbatch --test-only -N` until Slurm says it *could* schedule it +# 3. submit the real -N hold (now accepted; pends behind the warmups) +# 4. drop the warmups so the hold starts immediately +# +# Note: other people's jobs (e.g. the demo hold) occupy nodes too, so the pool has to grow to +# roughly N + (nodes held by others) before a -N request becomes satisfiable. +# +# Usage: NODES=12 JOB=wan14b_hold bash examples/train/acquire_nodes.sh +set -uo pipefail +WORK=/mnt/lustre/vlm-s4duan +NODES="${NODES:-12}" +JOB="${JOB:-wan14b_hold}" +TIME="${TIME:-120:00:00}" +SLACK="${SLACK:-3}" # extra warmups beyond NODES, to out-pace other users' holds +WARM_SECS="${WARM_SECS:-900}" +MAX_ROUNDS="${MAX_ROUNDS:-60}" +mkdir -p "$WORK/logs" +say() { echo "[acquire $(date +%H:%M:%S)] $*"; } + +EXIST=$(squeue -h -u "$USER" -n "${JOB}" -o '%i' 2>/dev/null | head -1) +if [ -n "$EXIST" ]; then say "reusing existing hold $EXIST"; echo "ALLOC=$EXIST"; exit 0; fi + +WANT_WARM=$(( NODES + SLACK )) + +for round in $(seq 1 "$MAX_ROUNDS"); do + pool=$(sinfo -h -p all -N -o '%N' 2>/dev/null | sort -u | wc -l) + + # Probe: can Slurm schedule -N$NODES at all? (--test-only never actually submits) + probe=$(sbatch --test-only -N"$NODES" --gres=gpu:4 --ntasks-per-node=1 --exclusive \ + -t "$TIME" -p all --wrap='hostname' 2>&1 | head -1) + + if echo "$probe" | grep -q "to start at"; then + say "pool=$pool — probe OK, submitting real -N$NODES hold" + OUT=$(sbatch -N"$NODES" --gres=gpu:4 --ntasks-per-node=1 --exclusive -t "$TIME" \ + -p all -J "$JOB" --requeue --chdir="$WORK" -o "$WORK/logs/${JOB}_%j.out" \ + --wrap='srun sleep infinity' 2>&1) + if echo "$OUT" | grep -q "Submitted batch"; then + ALLOC=$(echo "$OUT" | grep -oE '[0-9]+' | head -1) + say "SUBMITTED hold jobid=$ALLOC — clearing warmups so it can start" + scancel -u "$USER" -n warmup 2>/dev/null + for i in $(seq 1 180); do + st=$(squeue -h -j "$ALLOC" -o '%t' 2>/dev/null) + if [ "$st" = R ]; then + say "RUNNING on: $(squeue -h -j "$ALLOC" -o '%N')" + echo "ALLOC=$ALLOC"; exit 0 + fi + [ -z "$st" ] && { say "hold $ALLOC vanished"; break; } + sleep 10 + done + say "hold $ALLOC queued but not started yet (jobid kept)"; echo "ALLOC=$ALLOC"; exit 0 + fi + say "real submit rejected despite OK probe: $(echo "$OUT" | tail -1)" + else + say "pool=$pool — probe says not schedulable yet" + fi + + # Keep warmup pressure up so the autoscaler keeps growing the pool. + have=$(squeue -h -u "$USER" -n warmup -o '%i' 2>/dev/null | wc -l) + need=$(( WANT_WARM - have )) + if [ "$need" -gt 0 ]; then + say "submitting $need warmup(s) (have $have, want $WANT_WARM) to drive scale-up" + for _ in $(seq 1 "$need"); do + sbatch -N1 --gres=gpu:4 --exclusive -p all -t 00:20:00 -J warmup \ + -o /dev/null --chdir="$WORK" --wrap="srun sleep $WARM_SECS" >/dev/null 2>&1 + done + fi + sleep 30 +done + +say "FAILED to acquire $NODES nodes after $MAX_ROUNDS rounds" +exit 1 diff --git a/examples/train/chain_stepAB.sh b/examples/train/chain_stepAB.sh new file mode 100644 index 000000000..5fca72681 --- /dev/null +++ b/examples/train/chain_stepAB.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# Wait for step A to finish, then seed and launch step B — so the held allocation never sits +# idle between the two overfit stages. +# +# "A is done" means BOTH: no step-A launcher process is alive, AND a COMPLETE checkpoint-1000 +# exists (dcp/.metadata present). Requiring both avoids launching B off a half-written +# checkpoint if A died mid-save. +set -uo pipefail +WORK=/mnt/lustre/vlm-s4duan +REPO=$WORK/FastVideo +ALLOC="${ALLOC:-728}" +A_JOB="${A_JOB:-wan14b_stepA_v2}" +A_CKPT="$WORK/wantrack_14b_synth_sparse_fixed_out/checkpoint-1000" +: "${WANDB_API_KEY:?export WANDB_API_KEY}" +say() { echo "[chain $(date +%H:%M:%S)] $*"; } + +say "waiting for step A to reach a complete $A_CKPT ..." +while :; do + if [ -f "$A_CKPT/dcp/.metadata" ]; then + # Checkpoint is complete; wait for the launcher to actually exit before reusing the nodes. + # Use the launcher's PID file, NOT `pgrep -f run_wan14b_held.sh` — that pattern also matches + # any watcher/shell whose command line contains the string, so the check never goes false + # and the chain waits forever on nodes that are already idle. + A_RUNFILE="$WORK/logs/${A_JOB}.running" + if [ ! -f "$A_RUNFILE" ] || ! kill -0 "$(cat "$A_RUNFILE" 2>/dev/null)" 2>/dev/null; then + say "step A finished and $A_CKPT is complete" + break + fi + say "checkpoint-1000 complete, waiting for step A launcher (pid $(cat "$A_RUNFILE")) to exit ..." + fi + [ -n "$(squeue -h -j "$ALLOC" -o '%t' 2>/dev/null)" ] || { say "ALLOC $ALLOC vanished — aborting chain"; exit 1; } + sleep 60 +done + +say "seeding step B output dir from A's checkpoint-1000" +bash "$REPO/examples/train/run_stepB_seed.sh" || { say "seed failed"; exit 1; } + +say "launching step B (random-track overfit, WANTRACK_FIXED_SAMPLE=0)" +cd "$REPO" +exec env ALLOC="$ALLOC" NODES=4 JOB=wan14b_stepB_v2 PORT=30918 \ + CFG=examples/train/scenario/worldmodel/finetune_wantrack_synth_sparse_random_14b.yaml \ + WANTRACK_FIXED_SAMPLE=0 WANTRACK_FREEZE_HEAD=0 TRACKWAN_TRACK_BIAS=1 \ + WANDB_API_KEY="$WANDB_API_KEY" \ + bash examples/train/run_wan14b_held.sh diff --git a/examples/train/run_stepB_seed.sh b/examples/train/run_stepB_seed.sh new file mode 100644 index 000000000..08c721c29 --- /dev/null +++ b/examples/train/run_stepB_seed.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +# Seed step B's output dir with step A's final checkpoint, so B starts from A while its config +# can still use resume_from_checkpoint: latest (which is what makes crash-restarts safe). +# +# Hardlinked (cp -al), not copied: a 14B training-state checkpoint is hundreds of GB (bf16 +# weights + fp32 master + Adam m/v), and both dirs live on the same Lustre filesystem. +# Checkpoints are write-once, so sharing inodes is safe here. +set -euo pipefail +WORK=/mnt/lustre/vlm-s4duan +STEP="${STEP:-1000}" +SRC="${SRC:-$WORK/wantrack_14b_synth_sparse_fixed_out/checkpoint-$STEP}" +DST_DIR="${DST_DIR:-$WORK/wantrack_14b_synth_sparse_random_out}" +DST="$DST_DIR/checkpoint-$STEP" + +[ -d "$SRC" ] || { echo "[B-seed] missing $SRC"; exit 1; } +[ -f "$SRC/dcp/.metadata" ] || { echo "[B-seed] $SRC incomplete (no dcp/.metadata) — refusing"; exit 1; } + +if [ -d "$DST" ]; then + echo "[B-seed] $DST already exists — leaving it alone" +else + mkdir -p "$DST_DIR" + echo "[B-seed] hardlinking $SRC -> $DST" + cp -al "$SRC" "$DST" +fi +echo "[B-seed] checkpoints now in $DST_DIR:" +ls -d "$DST_DIR"/checkpoint-* 2>/dev/null | sed 's/^/ /' diff --git a/examples/train/run_stepC_merge.sh b/examples/train/run_stepC_merge.sh new file mode 100644 index 000000000..d63123d61 --- /dev/null +++ b/examples/train/run_stepC_merge.sh @@ -0,0 +1,74 @@ +#!/usr/bin/env bash +# STEP C — build the 14B "merged" init from the step-B overfit. +# +# Takes the 14B-native overfit (steps A+B) and transplants BOTH halves of the track pathway +# into a FRESH 14B base: +# * track_encoder.{proj,temporal_conv}.{weight,bias} (--track-src) +# * patch_embedding.weight[:, 36:] — the track slot (--pe-src) +# while patch_embedding.weight[:, :36] stays the pretrained I2V weights from --base. +# +# Both halves come from the SAME source on purpose: they were co-adapted during the overfit. +# Lifting the encoder alone (the "partial_merged" experiment) measured WORSE than random init, +# because the trained encoder ends up feeding a random projection. +# +# Usage: [STEP=2000] bash examples/train/run_stepC_merge.sh +set -euo pipefail +WORK=/mnt/lustre/vlm-s4duan +REPO=$WORK/FastVideo +STEP="${STEP:-2000}" +SRC_CKPT="${SRC_CKPT:-$WORK/wantrack_14b_synth_sparse_random_out/checkpoint-$STEP}" +EXPORT_DIR="${EXPORT_DIR:-$WORK/exports/overfit_14b_random_ckpt$STEP}" +OUT="${OUT:-$WORK/models/trackwan_14b_i2v_d64_merged_from_overfit_bias}" +BASE="${BASE:-$WORK/models/Wan2.1-I2V-14B-720P-Diffusers}" + +cd "$REPO" +source .venv/bin/activate +export HOME=$WORK HF_HOME=$WORK/.hf TORCH_HOME=$WORK/.torch MPLCONFIGDIR=$WORK/.mpl \ + PYTHONPATH=$REPO TOKENIZERS_PARALLELISM=false NCCL_CUMEM_ENABLE=0 \ + TRACKWAN_TRACK_BIAS=1 + +[ -d "$SRC_CKPT" ] || { echo "[C] missing $SRC_CKPT"; exit 1; } +[ -f "$SRC_CKPT/dcp/.metadata" ] || { echo "[C] $SRC_CKPT is an INCOMPLETE checkpoint (no dcp/.metadata)"; exit 1; } + +# 1) DCP -> diffusers (single process; DCP reshards automatically) +if [ ! -d "$EXPORT_DIR/transformer" ]; then + echo "[C] exporting $SRC_CKPT -> $EXPORT_DIR" + python -m fastvideo.train.entrypoint.dcp_to_diffusers \ + --checkpoint "$SRC_CKPT" --output-dir "$EXPORT_DIR" +else + echo "[C] reusing existing export $EXPORT_DIR" +fi + +# 2) fresh 14B base + co-adapted encoder AND track slot from that export +echo "[C] building merged init -> $OUT" +python data_pipeline/convert_trackwan_init_v2.py \ + --base "$BASE" --out "$OUT" --id-dim 64 --pe-init random \ + --track-src "$EXPORT_DIR" --pe-src "$EXPORT_DIR" + +# 3) verify: pretrained channels untouched, track slot + encoder match the overfit +echo "[C] verifying merged init against source ..." +python - "$OUT" "$EXPORT_DIR" "$BASE" <<'PY' +import sys, glob +import torch +from safetensors.torch import load_file + +def load(d): + st = {} + for f in sorted(glob.glob(f"{d}/transformer/*.safetensors")): + st.update(load_file(f)) + return st + +out, src, base = (load(p) for p in sys.argv[1:4]) +pe_o, pe_s, pe_b = out["patch_embedding.weight"], src["patch_embedding.weight"], base["patch_embedding.weight"] +print(f" pe shape {tuple(pe_o.shape)} (base {tuple(pe_b.shape)})") +print(f" pe[:, :36] == base : {torch.equal(pe_o[:, :36], pe_b)} std={pe_o[:, :36].float().std():.5f}") +print(f" pe[:, 36:] == overfit slot : {torch.equal(pe_o[:, 36:], pe_s[:, 36:])} std={pe_o[:, 36:].float().std():.5f}") +ok = torch.equal(pe_o[:, :36], pe_b) and torch.equal(pe_o[:, 36:], pe_s[:, 36:]) +for k in sorted(k for k in out if "track_encoder" in k): + same = k in src and torch.equal(out[k], src[k]) + ok &= same + print(f" {k:42s} lifted={same} norm={out[k].float().norm():.5f}") +print(" RESULT:", "OK — encoder and track slot are co-adapted" if ok else "MISMATCH") +sys.exit(0 if ok else 1) +PY +echo "[C] done -> $OUT" diff --git a/examples/train/run_synth_stage2_slurm.sh b/examples/train/run_synth_stage2_slurm.sh new file mode 100755 index 000000000..405ce62fa --- /dev/null +++ b/examples/train/run_synth_stage2_slurm.sh @@ -0,0 +1,43 @@ +#!/usr/bin/env bash +# Synth stage-2: MotionStream-style fine-tuning on 49k Wan2.2 synth data. +# Init: merged_bias_ckpt4800 (our stage-1 end on openvid). +# Recipe: TRACK_DROP=0.5 + MOTION_DROP=0.3 + PMASK=0.2 + freeze head. +# 600 steps -> ~1.55 epochs on combined 49k (paper claims "~1 epoch" for stage-2). +# Two VARIANTs: paperLR (1e-6, paper spec) or 5x (5e-6, 5x paper, 2x below stage-1's 1e-5). +set -euo pipefail +VARIANT="${VARIANT:-paperLR}" +NODES="${1:-4}"; shift || true +: "${WANDB_API_KEY:?export WANDB_API_KEY before launching}" + +case "$VARIANT" in + paperLR) + CFG=examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_paperLR.yaml + JOB=synth_stage2_paperLR + PORT=30700 + ;; + 5x) + CFG=examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_5x.yaml + JOB=synth_stage2_5x + PORT=30800 + ;; + *) echo "unknown VARIANT=$VARIANT"; exit 1 ;; +esac + +export WANTRACK_AUG=1 +export WANTRACK_SPARSE=1 +export WANTRACK_EXTRA_RANDOM=20 +export WANTRACK_EXTRA_MODE=random +export WANTRACK_TRACK_DROP=0.5 +export WANTRACK_MOTION_DROP=0.3 +export WANTRACK_PMASK=0.2 +export WANTRACK_MASK_CHUNK=8 +export WANTRACK_TEXT_DROP=0 +export WANTRACK_FIXED_SAMPLE=0 +export WANTRACK_DEBUG="${WANTRACK_DEBUG:-1}" +export WANTRACK_FREEZE_HEAD=1 +export TRACKWAN_TRACK_BIAS=1 + +export WANDB_MODE=online + +CFG="$CFG" JOB="$JOB" PORT="$PORT" \ + bash "$(dirname "$0")/run_openvid_bidir_held.sh" "$NODES" "$@" diff --git a/examples/train/run_wan14b_held.sh b/examples/train/run_wan14b_held.sh new file mode 100644 index 000000000..49a951a46 --- /dev/null +++ b/examples/train/run_wan14b_held.sh @@ -0,0 +1,137 @@ +#!/usr/bin/env bash +# Run a Wan2.1-14B WanTrack stage inside an ALREADY-HELD Slurm allocation. +# +# Same held-alloc + auto-restart pattern as run_openvid_bidir_held.sh, but every WANTRACK_* +# knob is parameterized instead of hardcoded, because the four stages of the 14B teacher +# pipeline need DIFFERENT sampling/masking settings: +# +# A (fixed overfit) : WANTRACK_FIXED_SAMPLE=1 FREEZE_HEAD=0 +# B (random overfit) : WANTRACK_FIXED_SAMPLE=0 FREEZE_HEAD=0 +# D (stage-1 openvid): FREEZE_HEAD=1, no dropping +# E (stage-2 synth) : FREEZE_HEAD=1 TRACK_DROP=0.5 MOTION_DROP=0.3 PMASK=0.2 MASK_CHUNK=8 +# +# FREEZE_HEAD is the important one: the overfit stages MUST train track_encoder (that is the +# whole point — co-adapting it with the patch_embedding track slot), while stages D/E freeze it. +# +# Usage: +# ALLOC=728 NODES=4 JOB=wan14b_stepA PORT=30910 \ +# CFG=examples/train/scenario/worldmodel/finetune_wantrack_synth_sparse_fixed_14b.yaml \ +# WANTRACK_FIXED_SAMPLE=1 WANTRACK_FREEZE_HEAD=0 \ +# WANDB_API_KEY=... bash examples/train/run_wan14b_held.sh +set -uo pipefail +WORK=/mnt/lustre/vlm-s4duan +REPO=$WORK/FastVideo +CFG="${CFG:?set CFG}" +NODES="${NODES:-4}"; GPUS=4 +JOB="${JOB:-wan14b}" +PORT="${PORT:-30900}" +ALLOC="${ALLOC:?set ALLOC to a running held allocation jobid}" +WANDB_RUN_ID="${WANDB_RUN_ID:-${JOB}}" +: "${WANDB_API_KEY:?export WANDB_API_KEY}" +TOTAL_GPUS=$(( NODES * GPUS )) +mkdir -p "$WORK/logs" + +OUTPUT_DIR=$(grep -oE 'output_dir:[^#]*' "$REPO/$CFG" | head -1 | sed 's/.*output_dir:[[:space:]]*//; s/"//g' | xargs) + +# Per-stage WANTRACK defaults (overridable from the environment). +W_AUG="${WANTRACK_AUG:-1}" +W_SPARSE="${WANTRACK_SPARSE:-1}" +W_EXTRA_RANDOM="${WANTRACK_EXTRA_RANDOM:-20}" +W_EXTRA_MODE="${WANTRACK_EXTRA_MODE:-random}" +W_PMASK="${WANTRACK_PMASK:-0}" +W_MASK_CHUNK="${WANTRACK_MASK_CHUNK:-0}" +W_TRACK_DROP="${WANTRACK_TRACK_DROP:-0}" +W_MOTION_DROP="${WANTRACK_MOTION_DROP:-0}" +W_TEXT_DROP="${WANTRACK_TEXT_DROP:-0}" +W_FIXED="${WANTRACK_FIXED_SAMPLE:-0}" +W_FREEZE="${WANTRACK_FREEZE_HEAD:-0}" +W_BIAS="${TRACKWAN_TRACK_BIAS:-1}" + +# A crash mid-save leaves checkpoint-N/dcp without .metadata; drop it so resume falls back to +# the last COMPLETE checkpoint instead of crash-looping on "metadata is None". +clean_bad_ckpt() { + [ -n "$OUTPUT_DIR" ] || return 0 + local latest + latest=$(ls -d "$OUTPUT_DIR"/checkpoint-* 2>/dev/null | sed 's/.*checkpoint-//' | sort -n | tail -1) + [ -n "$latest" ] || return 0 + local d="$OUTPUT_DIR/checkpoint-$latest" + if [ ! -f "$d/dcp/.metadata" ]; then + echo "[held] checkpoint-$latest incomplete (no dcp/.metadata) — removing" + rm -rf "$d" + fi +} + +[ "$(squeue -h -j "$ALLOC" -o '%t' 2>/dev/null)" = R ] || { echo "[held] alloc $ALLOC not running"; exit 1; } +NODELIST=$(squeue -h -j "$ALLOC" -o '%N') +# Pin the exact subset of the held allocation this stage runs on. Without --nodelist, srun is +# free to pick ANY $NODES of the allocation's nodes, so the rdzv endpoint (computed here) can +# land on a node that isn't in the subset — rank 0 then never listens there and every worker +# dies with "client socket has timed out while trying to connect". Choosing the subset up front +# makes MASTER provably the node hosting rank 0, and lets stages share one allocation safely. +# By default take the first $NODES nodes of the allocation. Pass NODELIST_OVERRIDE=a,b,c to pin +# an explicit subset — required when several stages share one allocation, otherwise every run +# grabs the same leading nodes and they fight over the same GPUs. +SUBSET="${NODELIST_OVERRIDE:-$(scontrol show hostnames "$NODELIST" | head -n "$NODES" | paste -sd,)}" +MASTER=$(echo "$SUBSET" | cut -d, -f1) +echo "[held] alloc=$ALLOC nodes=$NODELIST" +echo "[held] subset=$SUBSET master=$MASTER using $NODES node(s) / $TOTAL_GPUS GPU(s)" +echo "[held] cfg=$CFG out=$OUTPUT_DIR" +echo "[held] WANTRACK: fixed=$W_FIXED freeze=$W_FREEZE pmask=$W_PMASK chunk=$W_MASK_CHUNK track_drop=$W_TRACK_DROP motion_drop=$W_MOTION_DROP bias=$W_BIAS" + +# Killing this script (or its srun) does NOT kill the torchrun/python ranks out on the compute +# nodes — they survive, keep ~107GB of GPU memory each, and silently keep writing checkpoints. +# That poisons the next run: it contends for GPUs and resurrects deleted output dirs. So sweep +# the subset on any exit. Matches on the config path so we only kill THIS stage's ranks. +# Liveness marker for anything chaining off this stage. Do NOT make chainers use +# `pgrep -f run_wan14b_held.sh`: that also matches any monitoring/shell command whose command +# line merely CONTAINS the string, so a watcher looking for this launcher keeps "seeing" it long +# after it exited (that stall cost ~1h of idle nodes between stages A and B). +RUNFILE="$WORK/logs/${JOB}.running" +echo "$$" > "$RUNFILE" + +cleanup_ranks() { + echo "[held] sweeping stray ranks for $CFG on $SUBSET ..." + timeout 120 srun --overlap --jobid="$ALLOC" --nodelist="$SUBSET" --nodes="$NODES" \ + --ntasks="$NODES" --ntasks-per-node=1 \ + bash -c "pkill -9 -f 'entrypoint/train.py --config $CFG'; exit 0" >/dev/null 2>&1 || true +} +trap 'echo "[held] interrupted — cleaning up"; cleanup_ranks; rm -f "$RUNFILE"; exit 130' INT TERM +trap 'rm -f "$RUNFILE"' EXIT + +attempt=0 +while :; do + [ -n "$(squeue -h -j "$ALLOC" -o '%t' 2>/dev/null)" ] || { echo "[held] alloc $ALLOC vanished"; exit 1; } + attempt=$((attempt + 1)) + clean_bad_ckpt + echo "=== [train] attempt $attempt on alloc $ALLOC ($(date)) ===" + srun --overlap --jobid="$ALLOC" --nodelist="$SUBSET" --nodes="$NODES" --ntasks="$NODES" \ + --ntasks-per-node=1 --chdir="$REPO" bash -lc " + source .venv/bin/activate + export HOME=$WORK HF_HOME=$WORK/.hf TORCH_HOME=$WORK/.torch MPLCONFIGDIR=$WORK/.mpl \ + TRITON_CACHE_DIR=$WORK/.cache/triton_${JOB} TORCHINDUCTOR_CACHE_DIR=$WORK/.cache/inductor_${JOB} \ + TOKENIZERS_PARALLELISM=false NCCL_CUMEM_ENABLE=0 PYTHONPATH=$REPO \ + WANDB_API_KEY=$WANDB_API_KEY WANDB_MODE=online \ + WANDB_RUN_ID=$WANDB_RUN_ID WANDB_RESUME=allow \ + WANTRACK_AUG=$W_AUG WANTRACK_SPARSE=$W_SPARSE \ + WANTRACK_EXTRA_RANDOM=$W_EXTRA_RANDOM WANTRACK_EXTRA_MODE=$W_EXTRA_MODE \ + WANTRACK_PMASK=$W_PMASK WANTRACK_MASK_CHUNK=$W_MASK_CHUNK \ + WANTRACK_TRACK_DROP=$W_TRACK_DROP WANTRACK_MOTION_DROP=$W_MOTION_DROP \ + WANTRACK_TEXT_DROP=$W_TEXT_DROP WANTRACK_FIXED_SAMPLE=$W_FIXED \ + WANTRACK_FREEZE_HEAD=$W_FREEZE TRACKWAN_TRACK_BIAS=$W_BIAS \ + WANTRACK_DEBUG=1 \ + TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=1800 TORCH_NCCL_TRACE_BUFFER_SIZE=2000 \ + NCCL_SOCKET_NTHREADS=4 NCCL_NSOCKS_PERTHREAD=8 + torchrun --nnodes=$NODES --nproc-per-node=$GPUS --node-rank=\$SLURM_PROCID \ + --rdzv-backend=c10d --rdzv-endpoint=$MASTER:$PORT \ + fastvideo/train/entrypoint/train.py --config $CFG \ + --training.distributed.num_gpus $TOTAL_GPUS + " + rc=$? + if [ $rc -eq 0 ]; then echo "[train] finished cleanly (rc=0)"; break; fi + echo "[train] exited rc=$rc — nodes still held, relaunching in 20s ..." + # A crashed rank can leave its siblings alive and holding GPU memory; clear them before retry + # or the relaunch will OOM or hang waiting on a rendezvous the stale ranks are still in. + cleanup_ranks + sleep 20 +done +echo "[held] stage complete; allocation $ALLOC left RUNNING for the next stage." diff --git a/examples/train/scenario/worldmodel/finetune_wantrack_openvid_14b_lwlr.yaml b/examples/train/scenario/worldmodel/finetune_wantrack_openvid_14b_lwlr.yaml new file mode 100644 index 000000000..eef9b4c35 --- /dev/null +++ b/examples/train/scenario/worldmodel/finetune_wantrack_openvid_14b_lwlr.yaml @@ -0,0 +1,100 @@ +# EXPERIMENT 2 (layer-wise LR) — can stage-1 bootstrap a FRESH track pathway on OpenVid, +# with no overfit->merge stages at all? +# +# Same goal as experiment 1, different mechanism: instead of staging, train everything from +# step 0 but give the track pathway a 100x LR (1e-3 vs 1e-5) so it can catch up with an already +# pretrained DiT. Adam is scale-invariant per parameter, so this MUST be a real param group with +# its own lr — scaling the gradient would do nothing. patch_embedding's pretrained channels +# [:, :36] are grad-masked so the 100x LR only ever touches the track slot. +# +# NOTE: an earlier in-repo LWLR attempt reported 'no encoder bootstrap', but it referenced +# 'patch_embedding.proj.weight' incorrectly and so likely never applied the boosted LR to +# anything. This re-tests the idea properly; the param-group builder now RAISES if nothing +# matches, so a silent no-op cannot masquerade as a negative result. +# +# Recipe is otherwise IDENTICAL to the 1.3B stage-1 reference: lr 1e-5, global bs 128, 480x832, +# 121f, 24fps, flow_shift 6, bf16 — only the param-group/warm-up behaviour is new. +# +# Init is the ZERO-gate init: patch_embedding[:, 36:] = 0, so the model starts bit-identical to +# base I2V. (The random-gate init was measured to corrupt the prior from step 0.) +# +# Launch env: WANTRACK_TRACK_GROUP=1 (build named track/base param groups), +# WANTRACK_TRACK_LR_MULT=100 (track pathway at 1e-3), WANTRACK_FREEZE_HEAD=0 +# (the encoder MUST train here — this is the opposite of the merged-init stage 1). +# 4 nodes x 4 GPU = 16 GPUs, grad_accum 8 -> 1 * 8 * 16 = 128 global. +models: + student: + _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel + init_from: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_zero_init_bias + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod + +training: + model_path: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_zero_init_bias + dit_precision: bf16 + + distributed: + num_gpus: 16 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 4 + hsdp_shard_dim: 4 + + data: + data_path: /mnt/lustre/vlm-s4duan/openvid_1m/combined_parquet_dataset + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 31 + num_height: 480 + num_width: 832 + num_frames: 121 + + optimizer: + learning_rate: 1.0e-5 + betas: [0.9, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 1000 # experiment: the bootstrap signal shows well before this + gradient_accumulation_steps: 8 # 1 * 8 * 16 gpus = 128 global + + checkpoint: + output_dir: /mnt/lustre/vlm-s4duan/exp_openvid_14b_lwlr_out + training_state_checkpointing_steps: 100 + checkpoints_total_limit: 100 # keep all — early checkpoints are the whole point + resume_from_checkpoint: latest + + tracker: + trackers: [wandb] + project_name: wantrack-bidir + run_name: exp-openvid-14b-lwlr-100x + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + track_validation: + _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback + every_steps: 250 + num_val_samples: 4 + val_data_path: /mnt/lustre/vlm-s4duan/val_examples_mixed/combined_parquet_dataset + num_inference_steps: 30 + guidance_scale: 3.0 + motion_guidance_scale: 1.5 + fps: 24 + grid_stride: 3 + tail: 12 + validate_at_start: true + seed: 1000 + +pipeline: + flow_shift: 6 diff --git a/examples/train/scenario/worldmodel/finetune_wantrack_openvid_14b_warmup.yaml b/examples/train/scenario/worldmodel/finetune_wantrack_openvid_14b_warmup.yaml new file mode 100644 index 000000000..6b1ce3d56 --- /dev/null +++ b/examples/train/scenario/worldmodel/finetune_wantrack_openvid_14b_warmup.yaml @@ -0,0 +1,100 @@ +# EXPERIMENT 1 (encoder warm-up) — can stage-1 bootstrap a FRESH track pathway on OpenVid, +# with no overfit->merge stages at all? +# +# The overfit+merge pipeline exists only because a from-scratch track encoder would not train +# during stage 1. This tests a more direct route: for the first N steps hold the pretrained DiT +# completely still (base param group at lr 0) and let ONLY the track pathway learn +# (track_encoder.* + patch_embedding.proj.weight[:, 36:]); then ramp the DiT back in. If the +# gate and encoder grow to a healthy scale and the loss behaves, the whole A/B/C detour can go. +# +# Recipe is otherwise IDENTICAL to the 1.3B stage-1 reference: lr 1e-5, global bs 128, 480x832, +# 121f, 24fps, flow_shift 6, bf16 — only the param-group/warm-up behaviour is new. +# +# Init is the ZERO-gate init: patch_embedding[:, 36:] = 0, so the model starts bit-identical to +# base I2V. (The random-gate init was measured to corrupt the prior from step 0.) +# +# Launch env: WANTRACK_TRACK_GROUP=1 (build named track/base param groups), +# WANTRACK_TRACK_LR_MULT=1 (keep the reference LR), WANTRACK_FREEZE_HEAD=0 +# (the encoder MUST train here — this is the opposite of the merged-init stage 1). +# 4 nodes x 4 GPU = 16 GPUs, grad_accum 8 -> 1 * 8 * 16 = 128 global. +models: + student: + _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel + init_from: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_zero_init_bias + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod + +training: + model_path: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_zero_init_bias + dit_precision: bf16 + + distributed: + num_gpus: 16 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 4 + hsdp_shard_dim: 4 + + data: + data_path: /mnt/lustre/vlm-s4duan/openvid_1m/combined_parquet_dataset + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 31 + num_height: 480 + num_width: 832 + num_frames: 121 + + optimizer: + learning_rate: 1.0e-5 + betas: [0.9, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 1000 # experiment: the bootstrap signal shows well before this + gradient_accumulation_steps: 8 # 1 * 8 * 16 gpus = 128 global + + checkpoint: + output_dir: /mnt/lustre/vlm-s4duan/exp_openvid_14b_warmup_out + training_state_checkpointing_steps: 100 + checkpoints_total_limit: 100 # keep all — early checkpoints are the whole point + resume_from_checkpoint: latest + + tracker: + trackers: [wandb] + project_name: wantrack-bidir + run_name: exp-openvid-14b-encoder-warmup + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + track_warmup: + _target_: fastvideo.train.callbacks.track_warmup.TrackWarmupCallback + warmup_steps: 300 # DiT held at lr 0 while the track pathway learns + ramp_steps: 200 # then linearly hand control back to the full model + log_every: 25 + track_validation: + _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback + every_steps: 250 + num_val_samples: 4 + val_data_path: /mnt/lustre/vlm-s4duan/val_examples_mixed/combined_parquet_dataset + num_inference_steps: 30 + guidance_scale: 3.0 + motion_guidance_scale: 1.5 + fps: 24 + grid_stride: 3 + tail: 12 + validate_at_start: true + seed: 1000 + +pipeline: + flow_shift: 6 diff --git a/examples/train/scenario/worldmodel/finetune_wantrack_openvid_sparse_14b_merged_bias.yaml b/examples/train/scenario/worldmodel/finetune_wantrack_openvid_sparse_14b_merged_bias.yaml new file mode 100644 index 000000000..cc1d4b5de --- /dev/null +++ b/examples/train/scenario/worldmodel/finetune_wantrack_openvid_sparse_14b_merged_bias.yaml @@ -0,0 +1,91 @@ +# STEP D — MotionStream stage-1 teacher training on the full OpenVid-1M bidir set, Wan2.1-14B. +# +# Direct 14B analog of finetune_wantrack_openvid_sparse_1p3b_merged_bias.yaml (wandb run +# openvid_bidir_1p3b_merged_bias). Hparams are copied EXACTLY from that run — lr 1e-5, +# global bs 128, 4800 steps, 480x832, 121f, 24fps, flow_shift 6, bf16 DiT — the only changes +# are the ones 14B geometry forces (hsdp dims, grad_accum to keep bs 128 at 32 GPUs). +# +# init_from is the step-C MERGED init: a fresh 14B base whose track_encoder AND +# patch_embedding.weight[:, 36:] both come from the 14B-native overfit (steps A+B), so the two +# are co-adapted. This is the whole point of the pipeline: gradient-SNR analysis showed a +# random 14B init gives per-sample track_encoder.proj grad ~0.011 (the same regime that made +# 1.3B-random FAIL), while 1.3B-merged sat at ~0.502 and succeeded. +# +# Launched with WANTRACK_FREEZE_HEAD=1 (track_encoder frozen) + TRACKWAN_TRACK_BIAS=1. +# 8 nodes x 4 GPU = 32 GPUs; grad_accum 4 -> 1 * 4 * 32 = 128 global. +models: + student: + _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel + init_from: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_merged_from_overfit_bias + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod + +training: + model_path: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_merged_from_overfit_bias + dit_precision: bf16 + + distributed: + num_gpus: 32 # overridden by the launcher to nodes*4 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 8 # replicate across the 8 nodes + hsdp_shard_dim: 4 # shard the 14B within each node over NVLink (8*4 = 32) + + data: + data_path: /mnt/lustre/vlm-s4duan/openvid_1m/combined_parquet_dataset + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 31 + num_height: 480 + num_width: 832 + num_frames: 121 + + optimizer: + learning_rate: 1.0e-5 # MotionStream teacher stage-1 + betas: [0.9, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 4800 + gradient_accumulation_steps: 4 # 1 * 4 * 32 gpus = 128 global + + checkpoint: + output_dir: /mnt/lustre/vlm-s4duan/openvid_bidir_14b_merged_bias_out + training_state_checkpointing_steps: 100 # ~backstop vs the ~80min cordon cycle + checkpoints_total_limit: 3 + resume_from_checkpoint: latest + + tracker: + trackers: [wandb] + project_name: wantrack-bidir + run_name: openvid-bidir-14b-merged-init-bias + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + track_validation: + _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback + every_steps: 250 + num_val_samples: 4 + val_data_path: /mnt/lustre/vlm-s4duan/val_examples_mixed/combined_parquet_dataset # 2 openvid + car + dog + num_inference_steps: 30 + guidance_scale: 3.0 + motion_guidance_scale: 1.5 + fps: 24 + grid_stride: 3 + tail: 12 + validate_at_start: true + seed: 1000 + +pipeline: + flow_shift: 6 diff --git a/examples/train/scenario/worldmodel/finetune_wantrack_synth_sparse_fixed_14b.yaml b/examples/train/scenario/worldmodel/finetune_wantrack_synth_sparse_fixed_14b.yaml new file mode 100644 index 000000000..c1dcd8aad --- /dev/null +++ b/examples/train/scenario/worldmodel/finetune_wantrack_synth_sparse_fixed_14b.yaml @@ -0,0 +1,98 @@ +# STEP A — 14B fixed-track overfit on the 20-clip synth toy set. +# +# Direct 14B analog of finetune_wantrack_synth_sparse_fixed.yaml (1.3B). Purpose is NOT a good +# model: it is to CO-ADAPT track_encoder.{proj,temporal_conv} with patch_embedding.weight[:, 36:] +# for the 14B DiT. Gradient-SNR analysis showed a 14B random init has per-sample +# track_encoder.proj grad ~0.011 (same as the 1.3B random init that FAILED), while the 1.3B +# merged-from-overfit init sat at ~0.502 (50x) and succeeded. Steps A+B produce the 14B-native +# overfit whose encoder + track slot get merged into a fresh init in step C. +# +# Launched with WANTRACK_FIXED_SAMPLE=1 (deterministic per-sample track subset). +# +# GLOBAL BATCH SIZE 1 — this is exact parity with c0xyfg57 and is NOT negotiable for this stage. +# The reference ran 1 GPU, bs 1, accum 1 => global bs 1, 2000 steps => 2000 samples seen. An +# earlier attempt ran this on 16 GPUs at global bs 16 for 1000 steps = 16,000 samples, i.e. 8x +# the reference's ENTIRE run, and the model came out visibly overbaked by step ~1000. Step count +# only means something relative to batch size; matching lr + steps while changing batch silently +# changes exposure by the batch ratio. Large batch also removes gradient noise, so it memorises +# 20 clips faster than the reference even at equal samples. +# +# One node is still needed to hold a 14B: hsdp_shard_dim 4 shards ONE replica across the 4 GPUs +# and hsdp_replicate_dim 1 makes the data-parallel degree 1, so global bs stays 1. +# checkpoints_total_limit is deliberately huge: rotation previously destroyed the early +# checkpoints we wanted to fall back to. +models: + student: + _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel + init_from: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_zero_init_bias + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod + +training: + model_path: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_zero_init_bias + dit_precision: bf16 + + distributed: + num_gpus: 4 # overridden by the launcher to nodes*4 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 # data-parallel degree 1 -> global bs = train_batch_size = 1 + hsdp_shard_dim: 4 # shard ONE 14B replica across the node's 4 GPUs over NVLink + + data: + data_path: /mnt/lustre/vlm-s4duan/data/wantrack_synth_toy/preprocessed_i2v_track_varied/combined_parquet_dataset + dataloader_num_workers: 2 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 31 + num_height: 480 + num_width: 832 + num_frames: 121 + + optimizer: + learning_rate: 1.0e-4 # same as the 1.3B overfit (c0xyfg57 lineage) + betas: [0.9, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 2000 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: /mnt/lustre/vlm-s4duan/wantrack_14b_synth_sparse_fixed_out + training_state_checkpointing_steps: 100 + checkpoints_total_limit: 100 + resume_from_checkpoint: latest + + tracker: + trackers: [wandb] + project_name: wantrack-bidir + run_name: synth-sparse-fixed-14b + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + track_validation: + _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback + every_steps: 250 + num_val_samples: 2 # validates on the train (toy) clips — the point of an overfit + num_inference_steps: 30 + guidance_scale: 3.0 + motion_guidance_scale: 1.5 + fps: 24 + grid_stride: 3 + tail: 12 + validate_at_start: true + seed: 1000 + +pipeline: + flow_shift: 6 diff --git a/examples/train/scenario/worldmodel/finetune_wantrack_synth_sparse_random_14b.yaml b/examples/train/scenario/worldmodel/finetune_wantrack_synth_sparse_random_14b.yaml new file mode 100644 index 000000000..1b581b65b --- /dev/null +++ b/examples/train/scenario/worldmodel/finetune_wantrack_synth_sparse_random_14b.yaml @@ -0,0 +1,91 @@ +# STEP B — 14B RANDOM-track overfit, warm-started from step A checkpoint-1000. +# +# Direct 14B analog of finetune_wantrack_synth_sparse_fixed.yaml (1.3B). Purpose is NOT a good +# model: it is to CO-ADAPT track_encoder.{proj,temporal_conv} with patch_embedding.weight[:, 36:] +# for the 14B DiT. Gradient-SNR analysis showed a 14B random init has per-sample +# track_encoder.proj grad ~0.011 (same as the 1.3B random init that FAILED), while the 1.3B +# merged-from-overfit init sat at ~0.502 (50x) and succeeded. Steps A+B produce the 14B-native +# overfit whose encoder + track slot get merged into a fresh init in step C. +# +# Launched with WANTRACK_FIXED_SAMPLE=1 (deterministic per-sample track subset). +# 4 nodes x 4 GPU = 16 GPUs, per-gpu bs 1, grad_accum 1 -> global bs 16 (20-clip set, so each +# step sees ~80% of the data: effectively full-batch, which is what we want for an overfit). +models: + student: + _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel + init_from: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_zero_init_bias + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod + +training: + model_path: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_zero_init_bias + dit_precision: bf16 + + distributed: + num_gpus: 4 # overridden by the launcher to nodes*4 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 # data-parallel degree 1 -> global bs = train_batch_size = 1 (parity with c0xyfg57) + hsdp_shard_dim: 4 # shard ONE 14B replica across the node's 4 GPUs + + data: + data_path: /mnt/lustre/vlm-s4duan/data/wantrack_synth_toy/preprocessed_i2v_track_varied/combined_parquet_dataset + dataloader_num_workers: 2 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 31 + num_height: 480 + num_width: 832 + num_frames: 121 + + optimizer: + learning_rate: 1.0e-4 # same as the 1.3B overfit (c0xyfg57 lineage) + betas: [0.9, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 2000 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: /mnt/lustre/vlm-s4duan/wantrack_14b_synth_sparse_random_out + training_state_checkpointing_steps: 100 + checkpoints_total_limit: 100 + # NOT a hardcoded path to step A's checkpoint: the launcher auto-restarts on any crash, and a + # fixed path would make EVERY restart resume from A again, silently discarding B's progress. + # run_stepB_seed.sh hardlinks A's checkpoint-1000 into THIS output_dir first, so 'latest' + # resolves to it on the first launch and to B's own newest checkpoint on every restart. + resume_from_checkpoint: latest + + tracker: + trackers: [wandb] + project_name: wantrack-bidir + run_name: synth-sparse-random-14b-warmstart + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + track_validation: + _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback + every_steps: 250 + num_val_samples: 2 # validates on the train (toy) clips — the point of an overfit + num_inference_steps: 30 + guidance_scale: 3.0 + motion_guidance_scale: 1.5 + fps: 24 + grid_stride: 3 + tail: 12 + validate_at_start: true + seed: 1000 + +pipeline: + flow_shift: 6 diff --git a/examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_14b_paperLR.yaml b/examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_14b_paperLR.yaml new file mode 100644 index 000000000..5990f1e9b --- /dev/null +++ b/examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_14b_paperLR.yaml @@ -0,0 +1,92 @@ +# STEP E — MotionStream stage-2 finetune WITH dropping on the combined synth set, Wan2.1-14B. +# +# Direct 14B analog of finetune_wantrack_synth_stage2_paperLR.yaml (wandb run +# synth_stage2_paperLR). Hparams copied EXACTLY from that run — lr 1e-6, global bs 128, +# 600 steps, 480x832, 121f, 24fps, flow_shift 6, bf16 DiT — the only changes are the ones +# 14B geometry forces (hsdp dims, grad_accum to keep bs 128 at 32 GPUs). +# +# init_from is the step-D stage-1 result (OpenVid 4800-step ckpt, exported to diffusers). +# +# Stage 2 is where the stochastic dropping comes in — it is what teaches the model to cope +# with sparse/partial track conditioning at inference. Set at launch, NOT in this file: +# WANTRACK_TRACK_DROP=0.5 WANTRACK_MOTION_DROP=0.3 WANTRACK_PMASK=0.2 WANTRACK_MASK_CHUNK=8 +# WANTRACK_FREEZE_HEAD=1 TRACKWAN_TRACK_BIAS=1 +# Checkpoints every 50 steps and total_limit 100 so NOTHING rotates away: step-400 is the +# MotionStream-paper-parity point and must survive to be evaluated alongside step-600. +# 8 nodes x 4 GPU = 32 GPUs; grad_accum 4 -> 1 * 4 * 32 = 128 global. +models: + student: + _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel + init_from: /mnt/lustre/vlm-s4duan/exports/openvid_14b_merged_bias_ckpt4800 + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod + +training: + model_path: /mnt/lustre/vlm-s4duan/exports/openvid_14b_merged_bias_ckpt4800 + dit_precision: bf16 + + distributed: + num_gpus: 32 # overridden by the launcher to nodes*4 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 8 # replicate across the 8 nodes + hsdp_shard_dim: 4 # shard the 14B within each node over NVLink (8*4 = 32) + + data: + data_path: /mnt/lustre/vlm-s4duan/data/combined_synth_parquets + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 31 + num_height: 480 + num_width: 832 + num_frames: 121 + + optimizer: + learning_rate: 1.0e-6 # MotionStream teacher stage-2 + betas: [0.9, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 600 + gradient_accumulation_steps: 4 # 1 * 4 * 32 gpus = 128 global + + checkpoint: + output_dir: /mnt/lustre/vlm-s4duan/synth_bidir_14b_stage2_paperLR_out + training_state_checkpointing_steps: 50 # every 50 steps -> 12 ckpts, so step-400 (paper parity) is preserved + checkpoints_total_limit: 100 # keep ALL ckpts + resume_from_checkpoint: null + + tracker: + trackers: [wandb] + project_name: wantrack-bidir + run_name: synth-bidir-14b-stage2-paperLR + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + track_validation: + _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback + every_steps: 250 + num_val_samples: 4 + val_data_path: /mnt/lustre/vlm-s4duan/val_examples_mixed/combined_parquet_dataset # 2 openvid + car + dog + num_inference_steps: 30 + guidance_scale: 3.0 + motion_guidance_scale: 1.5 + fps: 24 + grid_stride: 3 + tail: 12 + validate_at_start: true + seed: 1000 + +pipeline: + flow_shift: 6 diff --git a/examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_5x.yaml b/examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_5x.yaml new file mode 100644 index 000000000..f67e91925 --- /dev/null +++ b/examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_5x.yaml @@ -0,0 +1,85 @@ +# MotionStream-style INITIAL teacher training (stage 1) on the full OpenVid-1M bidir set. +# Recipe: sparse point conditioning (1-per-SAM-object + extras) from run c0xyfg57 +# (synth-sparse-random, d64_nobias init) but with MotionStream stage-1 hparams: +# batch size 128, lr 1e-5, 480p (480x832), 24fps, 121f. NO stochastic track masking, +# NO fixed/overfit sampling (set via WANTRACK_* env at launch — see run_openvid_bidir_slurm.sh). +# num_gpus + hsdp dims below assume 4 nodes x 4 GPU = 16 GPUs (grad_accum 8 -> global bs 128). +# Runs CONCURRENTLY with the 14B run (8 nodes) = 12 nodes total <= 14 usable. The small 1.3B +# gets fewer nodes; grad_accum keeps global bs at 128 regardless of GPU count. +# DiT runs bf16 (fp32 full-scale would be ~70h; verified 1-GPU fp32 ~11s/it in smoke). +models: + student: + _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel + init_from: /mnt/lustre/vlm-s4duan/exports/merged_bias_ckpt4800 + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod + +training: + model_path: /mnt/lustre/vlm-s4duan/exports/merged_bias_ckpt4800 + dit_precision: bf16 # fp32 was ~11s/it on 1 GPU (=~70h full); bf16 for throughput + + distributed: + num_gpus: 16 # overridden by run_slurm.sh to nodes*NUM_GPUS + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 4 # replicate across the 4 nodes + hsdp_shard_dim: 4 # shard within each node over NVLink (4*4 = 16) + + data: + data_path: /mnt/lustre/vlm-s4duan/data/combined_synth_parquets + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 31 + num_height: 480 + num_width: 832 + num_frames: 121 + + optimizer: + learning_rate: 5.0e-6 # MotionStream teacher stage-2 + betas: [0.9, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 600 # synth stage-2: ~1.55 epochs on 49k combined synth; step-400 ckpt preserved for paper-parity comparison + gradient_accumulation_steps: 8 # 1 (per-gpu) * 8 * 16 gpus = 128 global + + checkpoint: + output_dir: /mnt/lustre/vlm-s4duan/synth_bidir_1p3b_stage2_5x_out + training_state_checkpointing_steps: 50 # every 50 steps -> 12 ckpts total + checkpoints_total_limit: 100 # keep ALL ckpts (12 << 100) so we can evaluate step-400 too + resume_from_checkpoint: null + + tracker: + trackers: [wandb] + project_name: wantrack-bidir + run_name: synth-bidir-1p3b-stage2-5x + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + track_validation: + _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback + every_steps: 250 + num_val_samples: 4 + val_data_path: /mnt/lustre/vlm-s4duan/val_examples_mixed/combined_parquet_dataset # fixed car+dog synth clips + num_inference_steps: 30 + guidance_scale: 3.0 + motion_guidance_scale: 1.5 + fps: 24 + grid_stride: 3 + tail: 12 + validate_at_start: true + seed: 1000 + +pipeline: + flow_shift: 6 diff --git a/examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_paperLR.yaml b/examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_paperLR.yaml new file mode 100644 index 000000000..d4b603430 --- /dev/null +++ b/examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_paperLR.yaml @@ -0,0 +1,85 @@ +# MotionStream-style INITIAL teacher training (stage 1) on the full OpenVid-1M bidir set. +# Recipe: sparse point conditioning (1-per-SAM-object + extras) from run c0xyfg57 +# (synth-sparse-random, d64_nobias init) but with MotionStream stage-1 hparams: +# batch size 128, lr 1e-5, 480p (480x832), 24fps, 121f. NO stochastic track masking, +# NO fixed/overfit sampling (set via WANTRACK_* env at launch — see run_openvid_bidir_slurm.sh). +# num_gpus + hsdp dims below assume 4 nodes x 4 GPU = 16 GPUs (grad_accum 8 -> global bs 128). +# Runs CONCURRENTLY with the 14B run (8 nodes) = 12 nodes total <= 14 usable. The small 1.3B +# gets fewer nodes; grad_accum keeps global bs at 128 regardless of GPU count. +# DiT runs bf16 (fp32 full-scale would be ~70h; verified 1-GPU fp32 ~11s/it in smoke). +models: + student: + _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel + init_from: /mnt/lustre/vlm-s4duan/exports/merged_bias_ckpt4800 + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod + +training: + model_path: /mnt/lustre/vlm-s4duan/exports/merged_bias_ckpt4800 + dit_precision: bf16 # fp32 was ~11s/it on 1 GPU (=~70h full); bf16 for throughput + + distributed: + num_gpus: 16 # overridden by run_slurm.sh to nodes*NUM_GPUS + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 4 # replicate across the 4 nodes + hsdp_shard_dim: 4 # shard within each node over NVLink (4*4 = 16) + + data: + data_path: /mnt/lustre/vlm-s4duan/data/combined_synth_parquets + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 31 + num_height: 480 + num_width: 832 + num_frames: 121 + + optimizer: + learning_rate: 1.0e-6 # MotionStream teacher stage-2 + betas: [0.9, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 600 # synth stage-2: ~1.55 epochs on 49k combined synth; step-400 ckpt preserved for paper-parity comparison + gradient_accumulation_steps: 8 # 1 (per-gpu) * 8 * 16 gpus = 128 global + + checkpoint: + output_dir: /mnt/lustre/vlm-s4duan/synth_bidir_1p3b_stage2_paperLR_out + training_state_checkpointing_steps: 50 # every 50 steps -> 12 ckpts total + checkpoints_total_limit: 100 # keep ALL ckpts (12 << 100) so we can evaluate step-400 too + resume_from_checkpoint: null + + tracker: + trackers: [wandb] + project_name: wantrack-bidir + run_name: synth-bidir-1p3b-stage2-paperLR + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + track_validation: + _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback + every_steps: 250 + num_val_samples: 4 + val_data_path: /mnt/lustre/vlm-s4duan/val_examples_mixed/combined_parquet_dataset # fixed car+dog synth clips + num_inference_steps: 30 + guidance_scale: 3.0 + motion_guidance_scale: 1.5 + fps: 24 + grid_stride: 3 + tail: 12 + validate_at_start: true + seed: 1000 + +pipeline: + flow_shift: 6 diff --git a/examples/train/watch_gate.sh b/examples/train/watch_gate.sh new file mode 100644 index 000000000..9c289f88a --- /dev/null +++ b/examples/train/watch_gate.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +# Sample the WanTrack gate at each new checkpoint and append to a TSV, so we can see whether it +# is on track for the ~0.011 the successful 1.3B merged init reached. Checkpoints rotate +# (total_limit=3), so poll faster than they are produced or the early trajectory is lost. +set -uo pipefail +WORK=/mnt/lustre/vlm-s4duan; REPO=$WORK/FastVideo +OUT="${OUT:-$WORK/wantrack_14b_synth_sparse_fixed_out}" +LOG="${LOG:-$WORK/logs/gate_growth.tsv}" +cd "$REPO"; source .venv/bin/activate +export HOME=$WORK PYTHONPATH=$REPO NCCL_CUMEM_ENABLE=0 RANK=0 LOCAL_RANK=0 WORLD_SIZE=1 \ + MASTER_ADDR=127.0.0.1 MASTER_PORT=29600 +touch "$LOG" +while :; do + for d in "$OUT"/checkpoint-*; do + [ -f "$d/dcp/.metadata" ] || continue + step=$(basename "$d" | sed 's/checkpoint-//') + grep -q "^${step} " "$LOG" 2>/dev/null && continue + line=$(python data_pipeline/check_gate_growth.py "$d" 2>/dev/null | grep -oE "gate\[:, 36:\] std=[0-9.]+" | grep -oE "[0-9.]+$") + [ -n "$line" ] && { printf "%s\t%s\t%s\n" "$step" "$line" "$(date +%H:%M:%S)" >> "$LOG"; echo "step $step gate=$line"; } + done + pgrep -f "run_wan14b_held.sh" >/dev/null || { echo "[watch] no launcher; exiting"; break; } + sleep 240 +done diff --git a/fastvideo/train/callbacks/track_warmup.py b/fastvideo/train/callbacks/track_warmup.py new file mode 100644 index 000000000..bb76c4b6a --- /dev/null +++ b/fastvideo/train/callbacks/track_warmup.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Staged unfreeze for the WanTrack track pathway. + +Tests whether a RANDOM-init track encoder can be bootstrapped directly on stage-1 data, which +would remove the need for the separate overfit -> merge stages. For the first ``warmup_steps`` +optimizer steps only the track pathway learns; the pretrained DiT is held still. After that the +DiT joins at its normal LR, either instantly or ramped in over ``ramp_steps``. + +Requires ``WANTRACK_TRACK_GROUP=1`` so the optimizer has the named 'track'/'base' param groups +(see FineTuneMethod._build_track_param_groups). + +Why lr=0 rather than requires_grad=False: params are bound to the optimizer at construction, so +flipping requires_grad later would leave the DiT permanently outside the optimizer. Setting the +group's lr to 0 is a true freeze for AdamW — the update is scaled by lr AND so is the decoupled +weight decay (p *= 1 - lr*wd), so at lr=0 the weights genuinely do not move. Zeroing grads +instead would still let weight decay shrink them. + +The LR must be (re)written every step from ``on_before_optimizer_step``, because the LR +scheduler rewrites ``group['lr']`` on its own step; this hook runs after that and immediately +before ``optimizer.step()``, so it wins. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from fastvideo.logger import init_logger +from fastvideo.train.callbacks.callback import Callback + +if TYPE_CHECKING: + from fastvideo.train.methods.base import TrainingMethod + +logger = init_logger(__name__) + + +class TrackWarmupCallback(Callback): + """Hold the pretrained DiT at lr=0 while the track pathway warms up.""" + + def __init__( + self, + *, + warmup_steps: int = 500, + ramp_steps: int = 0, + base_group: str = "base", + track_group: str = "track", + log_every: int = 50, + ) -> None: + self._warmup = int(warmup_steps) + self._ramp = max(int(ramp_steps), 0) + self._base_name = str(base_group) + self._track_name = str(track_group) + self._log_every = max(int(log_every), 1) + self._checked = False + + def _scale_for(self, iteration: int) -> float: + """0 during warmup, then linearly to 1 over ramp_steps (instant if ramp_steps=0).""" + if iteration < self._warmup: + return 0.0 + if self._ramp <= 0: + return 1.0 + return min(1.0, (iteration - self._warmup) / float(self._ramp)) + + def on_before_optimizer_step( + self, + method: TrainingMethod, + iteration: int = 0, + ) -> None: + scale = self._scale_for(iteration) + applied = False + for opt in method.get_optimizers(iteration): + for group in opt.param_groups: + name = group.get("name") + if name == self._base_name: + group["lr"] = group["lr"] * scale + applied = True + + if not self._checked: + self._checked = True + if not applied: + # Fail loud: silently running a plain finetune would look like a successful + # experiment while testing nothing. + raise RuntimeError( + f"TrackWarmupCallback found no param group named {self._base_name!r}. " + "Set WANTRACK_TRACK_GROUP=1 so the optimizer is built with named " + "'track'/'base' groups.") + logger.info( + "[WANTRACK] TrackWarmupCallback active: DiT frozen (lr x0) for %d steps, " + "then ramped over %d steps", self._warmup, self._ramp) + + if iteration % self._log_every == 0: + tracker = getattr(method, "tracker", None) + if tracker is not None: + lrs = {} + for opt in method.get_optimizers(iteration): + for group in opt.param_groups: + if group.get("name"): + lrs[f"lr/{group['name']}"] = float(group["lr"]) + lrs["lr/base_scale"] = scale + if lrs: + tracker.log(lrs, iteration) diff --git a/fastvideo/train/methods/fine_tuning/finetune.py b/fastvideo/train/methods/fine_tuning/finetune.py index 32ae50f4e..ae03dfe29 100644 --- a/fastvideo/train/methods/fine_tuning/finetune.py +++ b/fastvideo/train/methods/fine_tuning/finetune.py @@ -3,16 +3,20 @@ from __future__ import annotations +import os from typing import Any, Literal import torch import torch.nn.functional as F +from fastvideo.logger import init_logger from fastvideo.train.methods.base import TrainingMethod, LogScalar from fastvideo.train.models.base import ModelBase from fastvideo.train.utils.optimizer import ( build_optimizer_and_scheduler, ) +logger = init_logger(__name__) + class FineTuneMethod(TrainingMethod): """Supervised finetuning: only ``student`` participates.""" @@ -153,6 +157,57 @@ class FineTuneMethod(TrainingMethod): del iteration return [self._student_lr_scheduler] + # Names that make up the track pathway. patch_embedding is matched on the CHECKPOINT name + # used by the live module ('patch_embedding.proj.weight'); an earlier experiment silently + # matched nothing by using 'patch_embedding.weight' (the converter-side name) and so never + # actually applied its layer-wise LR. + _TRACK_ENCODER_HINT = "track_encoder" + _PATCH_EMBED_HINTS = ("patch_embedding.proj.weight", "patch_embedding.weight") + _BASE_IN_CHANNELS = 36 + + def _build_track_param_groups(self, base_lr: float) -> list[dict]: + """Split params into a boosted 'track' group and the normal 'base' group.""" + mult = float(os.getenv("WANTRACK_TRACK_LR_MULT", "1")) + track, base = [], [] + pe_masked = 0 + for name, p in self.student.transformer.named_parameters(): + if not p.requires_grad: + continue + is_pe = any(name.endswith(h) for h in self._PATCH_EMBED_HINTS) + if self._TRACK_ENCODER_HINT in name or is_pe: + if is_pe and p.dim() >= 2 and p.shape[1] > self._BASE_IN_CHANNELS: + # Keep the pretrained input channels out of the boosted group entirely. + base_in = self._BASE_IN_CHANNELS + + def _mask_pretrained(grad, _base_in=base_in): + g = grad.clone() + g[:, :_base_in] = 0 + return g + + p.register_hook(_mask_pretrained) + pe_masked += 1 + track.append(p) + else: + base.append(p) + + if not track: + raise ValueError( + "WANTRACK_TRACK_GROUP=1 but no track params matched " + f"({self._TRACK_ENCODER_HINT!r} / {self._PATCH_EMBED_HINTS!r}) — refusing to " + "run an experiment that would silently be a plain finetune") + + n_track = sum(p.numel() for p in track) + logger.info( + "[WANTRACK] track param group: %d tensors (%d params) lr=%.3g (mult %.3g); " + "base group: %d tensors lr=%.3g; patch_embedding grad-masked to [:, %d:] on %d tensor(s)", + len(track), n_track, base_lr * mult, mult, len(base), base_lr, + self._BASE_IN_CHANNELS, pe_masked) + # 'name' is what TrackWarmupCallback looks for to zero the base group during warmup. + return [ + {"params": track, "lr": base_lr * mult, "name": "track"}, + {"params": base, "lr": base_lr, "name": "base"}, + ] + def _init_optimizers_and_schedulers(self) -> None: tc = self.training_config @@ -164,6 +219,20 @@ class FineTuneMethod(TrainingMethod): student_betas = tc.optimizer.betas student_sched = str(tc.optimizer.lr_scheduler) student_params = [p for p in self.student.transformer.parameters() if p.requires_grad] + + # Optional: split the track pathway into its OWN param group so it can be given a + # different LR from the pretrained DiT (WANTRACK_TRACK_GROUP=1). Used to test whether a + # random-init encoder can be bootstrapped directly on stage-1 data — warming up the + # track pathway first, or running it at a higher LR — instead of via the overfit+merge. + # + # Two things make this non-trivial and are handled here: + # * Adam is scale-invariant per parameter, so scaling GRADIENTS does not emulate a + # higher LR. It has to be a real param group with its own lr. + # * patch_embedding.proj.weight is ONE tensor covering both the pretrained input + # channels [:, :36] and the track slot [:, 36:]. Putting it in a boosted group would + # blast the pretrained channels too, so we mask its gradient to the track slot. + if os.getenv("WANTRACK_TRACK_GROUP", "0") not in ("0", "false", "False"): + student_params = self._build_track_param_groups(student_lr) ( self._student_optimizer, self._student_lr_scheduler, diff --git a/gradient_flow_analysis.py b/gradient_flow_analysis.py new file mode 100644 index 000000000..e00c07389 --- /dev/null +++ b/gradient_flow_analysis.py @@ -0,0 +1,470 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +""" +Per-sample gradient / SNR analysis of the WanTrack DiT, comparing 4 init/checkpoint variants +(zero_init, random_init, merged, overfit_full) on real OpenVid parquet samples. + +For each (model, layer) we record: + avg_per_sample_norm = mean_i ||g_i|| (typical bs=1 grad magnitude) + mean_grad_norm = ||mean_i g_i|| (grad magnitude after bs=N averaging) + snr = mean_grad_norm / avg_per_sample_norm + +High SNR => signal survives averaging (universal component). +Low SNR => signal cancels (sample-specific / noise). + +Layers tracked (custom-model names, not safetensors names): + patch_embedding.proj.weight [:, :36] and [:, 36:] + track_encoder.proj.weight + track_encoder.temporal_conv.weight + blocks.0.to_q.weight (early self-attn Q) + blocks.14.ffn.fc_in.weight (mid FFN in-proj) + blocks.29.to_out.weight (late self-attn out) +""" +from __future__ import annotations + +import argparse +import glob +import json +import os +import random +import sys +from dataclasses import dataclass, field +from pathlib import Path + +import numpy as np +import pyarrow.parquet as pq +import torch +import torch.nn.functional as F +from safetensors.torch import load_file as st_load + +# Init distributed BEFORE any fastvideo import that touches SP world. +def _init_distributed_1gpu() -> None: + os.environ.setdefault("RANK", "0") + os.environ.setdefault("LOCAL_RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29511") + os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + from fastvideo.distributed.parallel_state import ( + maybe_init_distributed_environment_and_model_parallel, + ) + maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1) + + +_init_distributed_1gpu() + +# Now safe to import model classes +from fastvideo.configs.models.dits.trackwan import ( # noqa: E402 + TrackWanVideoArchConfig, + TrackWanVideoConfig, +) +from fastvideo.forward_context import set_forward_context # noqa: E402 +from fastvideo.models.dits.trackwan.model import TrackWanTransformer3DModel # noqa: E402 +from fastvideo.models.loader.utils import ( # noqa: E402 + get_param_names_mapping, + hf_to_custom_state_dict, +) + + +# ------------------------------------------------------------------ # +# Model set-up +# ------------------------------------------------------------------ # +MODELS = { + "zero_init": "/mnt/lustre/vlm-s4duan/models/trackwan_1.3b_i2v_d64_nobias_init", + "random_init": "/mnt/lustre/vlm-s4duan/models/trackwan_1.3b_i2v_d64_nobias_random_init", + "merged": "/mnt/lustre/vlm-s4duan/models/trackwan_1.3b_i2v_d64_merged_from_overfit", + "overfit_full": "/mnt/lustre/vlm-s4duan/models/overfit_2000", + "fun_control": "/mnt/lustre/vlm-s4duan/models/trackwan_1.3b_i2v_d64_from_fun_control_init", + # 14B variants + "14b_zero_init": "/mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_nobias_init", + "14b_random_init": "/mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_random_init", + "14b_partial_merged": "/mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_partial_merged", + # zero-gate init WITH bias (the base the step-A overfit started from — gate exactly 0). + "14b_zero_init_bias": "/mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_zero_init_bias", + # fresh 14B base with the FIXED-overfit (step A, ckpt-2000) encoder + gate merged in — this + # is the candidate stage-1 init; its per-sample track grad decides whether it will train. + "14b_merged_fixed2000": "/mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_merged_from_fixed2000_bias", +} + +# Custom-model parameter names we snapshot per sample. +# Special slicing for patch_embedding is handled below. +TRACK_LAYERS = [ + "patch_embedding.proj.weight", # will split into base=[:, :36] / track=[:, 36:] + "track_encoder.proj.weight", + "track_encoder.temporal_conv.weight", + "blocks.0.to_q.weight", + "blocks.14.ffn.fc_in.weight", + "blocks.29.to_out.weight", +] + + +def build_model_from_diffusers(model_dir: str, device: torch.device) -> TrackWanTransformer3DModel: + cfg_path = Path(model_dir) / "transformer" / "config.json" + with open(cfg_path, "r") as f: + hf_config = json.load(f) + + arch = TrackWanVideoArchConfig( + num_attention_heads=int(hf_config["num_attention_heads"]), + attention_head_dim=int(hf_config["attention_head_dim"]), + in_channels=int(hf_config["in_channels"]), + out_channels=int(hf_config["out_channels"]), + text_dim=int(hf_config.get("text_dim", 4096)), + freq_dim=int(hf_config.get("freq_dim", 256)), + ffn_dim=int(hf_config["ffn_dim"]), + num_layers=int(hf_config["num_layers"]), + cross_attn_norm=bool(hf_config.get("cross_attn_norm", True)), + qk_norm=str(hf_config.get("qk_norm", "rms_norm_across_heads")), + eps=float(hf_config.get("eps", 1e-6)), + image_dim=int(hf_config.get("image_dim") or 0) or None, + added_kv_proj_dim=int(hf_config.get("added_kv_proj_dim") or 0) or None, + rope_max_seq_len=int(hf_config.get("rope_max_seq_len", 1024)), + track_config=dict(hf_config.get("track_config", {})), + ) + # The trackwan config has an "in_channels: 52" default, but we still pass hf_config so + # TrackWanTransformer3DModel picks up the correct track_channels from track_config. + cfg = TrackWanVideoConfig(arch_config=arch) + + # Ensure downstream consumers see the right in_channels regardless of default + cfg.arch_config.in_channels = int(hf_config["in_channels"]) + cfg.arch_config.out_channels = int(hf_config["out_channels"]) + + model = TrackWanTransformer3DModel(config=cfg, hf_config=hf_config) + model = model.to(device=device, dtype=torch.bfloat16) + + # --- load weights --- + st_files = sorted(glob.glob(str(Path(model_dir) / "transformer" / "*.safetensors"))) + assert len(st_files) >= 1, f"no safetensors under {model_dir}/transformer" + hf_sd: dict[str, torch.Tensor] = {} + for p in st_files: + hf_sd.update(st_load(p)) + + mapping_fn = get_param_names_mapping(model.param_names_mapping) + custom_sd, _ = hf_to_custom_state_dict(hf_sd, mapping_fn) + + # Load into model + model_sd = model.state_dict() + missing, unexpected = [], [] + for k, v in custom_sd.items(): + if k in model_sd: + model_sd[k].data.copy_(v.to(model_sd[k].dtype)) + else: + unexpected.append(k) + for k in model_sd: + if k not in custom_sd: + missing.append(k) + print(f"[{model_dir}] missing={len(missing)} unexpected={len(unexpected)}") + if unexpected: + print(" unexpected head:", unexpected[:5]) + if missing: + print(" missing head:", missing[:5]) + + # Gradient checkpointing needs the forward context to be re-established during + # backward recompute; we keep both fwd+bwd wrapped in set_forward_context in + # run_one_model, so grad_checkpointing is safe. + model.gradient_checkpointing = True + model._gradient_checkpointing_func = lambda fn, *args, **kwargs: torch.utils.checkpoint.checkpoint( + fn, *args, use_reentrant=False, **kwargs + ) + + # FastVideo's linear layers create Parameters with requires_grad=False (FSDP path). + # For gradient analysis we force all params trainable. + for p in model.parameters(): + p.requires_grad_(True) + + model.train() # ensure gradients flow + return model + + +# ------------------------------------------------------------------ # +# Parquet loader +# ------------------------------------------------------------------ # +@dataclass +class Sample: + vae_latent: torch.Tensor # [16, T, H, W] + first_frame_latent: torch.Tensor # [16, T, H, W] + clip_feature: torch.Tensor # [257, 1280] + text_embedding: torch.Tensor # [L, 4096] + track_points: torch.Tensor # [T_full, N, 2] + track_visibility: torch.Tensor # [T_full, N] + object_ids: torch.Tensor # [N] + track_weights: torch.Tensor # [N] + + +def _decode(row, name) -> torch.Tensor: + shape = row[f"{name}_shape"] + data = row[f"{name}_bytes"] + arr = np.frombuffer(data, dtype=np.float32).reshape(shape).copy() + return torch.from_numpy(arr) + + +def load_samples(pattern: str, n: int, seed: int = 0) -> list[Sample]: + files = sorted(glob.glob(pattern, recursive=True)) + random.Random(seed).shuffle(files) + samples: list[Sample] = [] + for f in files: + if len(samples) >= n: + break + try: + tbl = pq.read_table(f) + except Exception as e: + print("skip", f, e) + continue + for row in tbl.to_pylist(): + try: + s = Sample( + vae_latent=_decode(row, "vae_latent"), + first_frame_latent=_decode(row, "first_frame_latent"), + clip_feature=_decode(row, "clip_feature"), + text_embedding=_decode(row, "text_embedding"), + track_points=_decode(row, "track_points"), + track_visibility=_decode(row, "track_visibility"), + object_ids=_decode(row, "object_ids"), + track_weights=_decode(row, "track_weights"), + ) + except Exception as e: + print("skip row", e) + continue + samples.append(s) + if len(samples) >= n: + break + print(f"loaded {len(samples)} samples from {pattern}") + return samples + + +# ------------------------------------------------------------------ # +# I2V cond concat and flow-matching noise +# ------------------------------------------------------------------ # +def build_i2v_cond_concat(first_frame_latent: torch.Tensor, + vae_temporal_compression: int = 4) -> torch.Tensor: + """[B,16,T,H,W] first-frame latent -> [B,20,T,H,W] (4 mask + 16 latent). + + Mirrors WanTrackModel._build_i2v_cond_concat (mask=1 on first frame slot only). + """ + b, _, num_latent_t, h, w = first_frame_latent.shape + ratio = vae_temporal_compression + num_frames = (num_latent_t - 1) * ratio + 1 + mask = torch.ones(b, 1, num_frames, h, w, + device=first_frame_latent.device, dtype=first_frame_latent.dtype) + mask[:, :, 1:] = 0 + first = torch.repeat_interleave(mask[:, :, :1], dim=2, repeats=ratio) + mask = torch.cat([first, mask[:, :, 1:]], dim=2) + mask = mask.view(b, -1, ratio, h, w).transpose(1, 2) # [B, ratio, num_latent_t, H, W] + return torch.cat([mask, first_frame_latent], dim=1) + + +def flow_matching_noisy(clean: torch.Tensor, noise: torch.Tensor, + sigma: float, flow_shift: float = 6.0) -> tuple[torch.Tensor, float, float]: + """(1-s)*x + s*n with shifted sigma. Returns noisy, shifted_sigma, timestep.""" + shifted = flow_shift * sigma / (1 + (flow_shift - 1) * sigma) + noisy = (1.0 - shifted) * clean + shifted * noise + return noisy, shifted, shifted * 1000.0 + + +# ------------------------------------------------------------------ # +# Per-sample gradient snapshot +# ------------------------------------------------------------------ # +def snapshot_layer_grads(model: torch.nn.Module) -> dict[str, torch.Tensor]: + out: dict[str, torch.Tensor] = {} + param_map = dict(model.named_parameters()) + for name in TRACK_LAYERS: + if name not in param_map: + # Some models pack Q/K/V differently. Try nested attn1 name as fallback. + alt = name.replace("blocks.0.to_q", "blocks.0.attn1.to_q") \ + .replace("blocks.29.to_out", "blocks.29.attn1.to_out") \ + .replace(".ffn.fc_in.", ".ffn.net.0.proj.") + if alt in param_map: + p = param_map[alt] + else: + print(f"WARN param not found: {name}") + continue + else: + p = param_map[name] + if p.grad is None: + print(f"WARN no grad on {name}") + continue + out[name] = p.grad.detach().clone().float().cpu() + return out + + +def run_one_model(model_name: str, model_dir: str, samples: list[Sample], + out_dir: Path, device: torch.device, + flow_shift: float = 6.0) -> dict: + print(f"\n===== {model_name} =====") + torch.cuda.empty_cache() + model = build_model_from_diffusers(model_dir, device) + print("model params:", + sum(p.numel() for p in model.parameters()) / 1e9, "B") + + # Store per-sample grads (may be large: patch_embedding.weight = 1536*52*4 = 320K params). + per_sample_grads: dict[str, list[torch.Tensor]] = {n: [] for n in TRACK_LAYERS} + + for i, s in enumerate(samples): + model.zero_grad(set_to_none=True) + + # Move tensors to device / bf16 dtype + vae_latent = s.vae_latent.unsqueeze(0).to(device, dtype=torch.bfloat16) # [1,16,T,H,W] + first_frame_latent = s.first_frame_latent.unsqueeze(0).to(device, dtype=torch.bfloat16) + clip_feature = s.clip_feature.unsqueeze(0).to(device, dtype=torch.bfloat16) + text_embedding = s.text_embedding.unsqueeze(0).to(device, dtype=torch.bfloat16) + # Note attention mask: use all ones over the raw text seq len + text_mask = torch.ones(1, s.text_embedding.shape[0], device=device, dtype=torch.bfloat16) + + # Track tensors + track_points = s.track_points.unsqueeze(0).to(device, dtype=torch.bfloat16) # [1, T_full, N, 2] + track_visibility = s.track_visibility.unsqueeze(0).to(device, dtype=torch.bfloat16) + + # Truncate track T to (num_latent_t - 1) * ratio + 1 + num_latent_t = vae_latent.shape[2] + expected_T = (num_latent_t - 1) * 4 + 1 + track_points = track_points[:, :expected_T] + track_visibility = track_visibility[:, :expected_T] + + # Sample sigma ~ U(0,1) + sigma_u = float(torch.rand(1).item()) + # Compute noise and noisy + noise = torch.randn_like(vae_latent, dtype=torch.float32) + clean_f = vae_latent.float() + noisy_f, shifted_sigma, timestep_val = flow_matching_noisy(clean_f, noise, sigma_u, + flow_shift=flow_shift) + noisy = noisy_f.to(torch.bfloat16) # [1, 16, T, H, W] + + # Build I2V concat + cond20 = build_i2v_cond_concat(first_frame_latent, vae_temporal_compression=4) # [1,20,T,H,W] + hidden_states = torch.cat([noisy, cond20], dim=1) # [1, 36, T, H, W] + + # Timestep tensor + timestep = torch.tensor([timestep_val], device=device, dtype=torch.bfloat16) + + # Forward + backward BOTH inside set_forward_context so grad-checkpoint + # recompute finds the forward context. + with torch.autocast(device_type="cuda", dtype=torch.bfloat16), \ + set_forward_context(current_timestep=timestep, attn_metadata=None): + pred = model( + hidden_states=hidden_states, + encoder_hidden_states=text_embedding, + encoder_attention_mask=text_mask, + timestep=timestep, + encoder_hidden_states_image=clip_feature, + track_points=track_points, + track_visibility=track_visibility, + return_dict=False, + ) + # pred is [B, C, T, H, W] + # Loss: MSE(pred, noise - clean) computed in fp32 + target = (noise - clean_f) # [B,16,T,H,W] + loss = F.mse_loss(pred.float(), target.float()) + loss.backward() + + # Snapshot + gs = snapshot_layer_grads(model) + for n, g in gs.items(): + per_sample_grads[n].append(g) + + if i < 3 or i % 8 == 0: + print(f" sample {i}: sigma_u={sigma_u:.3f} shifted={shifted_sigma:.3f} " + f"t={timestep_val:.1f} loss={loss.item():.4f} " + f"|g[patch_embed]|={gs.get('patch_embedding.proj.weight', torch.zeros(1)).norm().item():.4f}") + # free bookkeeping + del pred, hidden_states, noisy, noise, cond20, target, loss + torch.cuda.empty_cache() + + # Save all grads + model_out = out_dir / model_name + model_out.mkdir(parents=True, exist_ok=True) + torch.save({n: torch.stack(gs, dim=0) for n, gs in per_sample_grads.items() + if len(gs) > 0}, + model_out / "grads.pt") + + # Compute stats + stats: dict = {} + for n in TRACK_LAYERS: + if n not in per_sample_grads or not per_sample_grads[n]: + continue + g_stack = torch.stack(per_sample_grads[n], dim=0) # [N, *param_shape] + # Special split for patch_embedding + if n == "patch_embedding.proj.weight": + # Shape [N, out, in, kD, kH, kW] with in=52 + base_slice = g_stack[:, :, :36] # [N, out, 36, ...] + track_slice = g_stack[:, :, 36:] # [N, out, 16, ...] + for key, sub in [(f"{n}[:,:36]", base_slice), (f"{n}[:,36:]", track_slice)]: + per_sample_norms = sub.flatten(1).norm(dim=1) # [N] + mean_grad = sub.mean(0) + mean_grad_norm = mean_grad.flatten().norm().item() + avg_per_sample_norm = per_sample_norms.mean().item() + snr = mean_grad_norm / (avg_per_sample_norm + 1e-30) + stats[key] = { + "avg_per_sample_norm": avg_per_sample_norm, + "mean_grad_norm": mean_grad_norm, + "snr": snr, + "n_params": int(sub[0].numel()), + } + else: + per_sample_norms = g_stack.flatten(1).norm(dim=1) + mean_grad = g_stack.mean(0) + mean_grad_norm = mean_grad.flatten().norm().item() + avg_per_sample_norm = per_sample_norms.mean().item() + snr = mean_grad_norm / (avg_per_sample_norm + 1e-30) + stats[n] = { + "avg_per_sample_norm": avg_per_sample_norm, + "mean_grad_norm": mean_grad_norm, + "snr": snr, + "n_params": int(g_stack[0].numel()), + } + + # Free model + del model + torch.cuda.empty_cache() + return stats + + +def print_table(all_stats: dict) -> None: + print("\n\n===== SNR TABLE =====") + print(f"{'model':<14} {'layer':<48} {'avg_per_sample':>16} {'mean_grad':>16} {'snr':>10}") + print("-" * 110) + for model, stats in all_stats.items(): + for layer, s in stats.items(): + print(f"{model:<14} {layer:<48} {s['avg_per_sample_norm']:>16.6f} " + f"{s['mean_grad_norm']:>16.6f} {s['snr']:>10.4f}") + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--n-samples", type=int, default=32) + ap.add_argument("--parquet-glob", type=str, + default="/mnt/lustre/vlm-s4duan/openvid_1m/combined_parquet_dataset/shard_*/**/*.parquet") + ap.add_argument("--out-dir", type=str, default="/mnt/lustre/vlm-s4duan/gradient_analysis") + ap.add_argument("--flow-shift", type=float, default=6.0) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--only", type=str, default="", help="comma list of models to run") + args = ap.parse_args() + + torch.manual_seed(args.seed) + random.seed(args.seed) + np.random.seed(args.seed) + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + device = torch.device("cuda") + + samples = load_samples(args.parquet_glob, args.n_samples, seed=args.seed) + assert len(samples) == args.n_samples, "not enough samples loaded" + + all_stats: dict = {} + models_to_run = args.only.split(",") if args.only else list(MODELS.keys()) + for name in models_to_run: + name = name.strip() + if not name: + continue + stats = run_one_model(name, MODELS[name], samples, out_dir, device, + flow_shift=args.flow_shift) + all_stats[name] = stats + # Save partial after each model + with open(out_dir / "summary.json", "w") as f: + json.dump(all_stats, f, indent=2) + + print_table(all_stats) + print(f"\nsummary -> {out_dir / 'summary.json'}") + + +if __name__ == "__main__": + main() diff --git a/gradient_flow_summarize.py b/gradient_flow_summarize.py new file mode 100644 index 000000000..574c1abea --- /dev/null +++ b/gradient_flow_summarize.py @@ -0,0 +1,52 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Read /mnt/lustre/vlm-s4duan/gradient_analysis/summary.json and print a compact table +plus a "magnitude ratio" analysis (track/base) at bs=N averaging. +""" +from __future__ import annotations +import json +import sys +from pathlib import Path + + +def main() -> None: + path = Path(sys.argv[1] if len(sys.argv) > 1 else "/mnt/lustre/vlm-s4duan/gradient_analysis/summary.json") + with open(path) as f: + stats = json.load(f) + + layers = [ + "patch_embedding.proj.weight[:,:36]", + "patch_embedding.proj.weight[:,36:]", + "track_encoder.temporal_conv.weight", + "track_encoder.proj.weight", + "blocks.0.to_q.weight", + "blocks.14.ffn.fc_in.weight", + "blocks.29.to_out.weight", + ] + + print(f"{'model':<14} {'layer':<44} {'per_sample':>12} {'mean_grad':>12} {'snr':>8}") + print("-" * 96) + for model, s in stats.items(): + for lay in layers: + if lay not in s: + continue + v = s[lay] + print(f"{model:<14} {lay:<44} {v['avg_per_sample_norm']:>12.4e} " + f"{v['mean_grad_norm']:>12.4e} {v['snr']:>8.4f}") + print("-" * 96) + + # Magnitude ratio: how big is bs=N averaged grad on the track slice vs base slice? + print("\nMagnitude ratio (mean_grad_norm(track slice) / mean_grad_norm(base slice))") + print(f"{'model':<14} {'track_slice / base_slice (patch_embed)':>44}") + for model, s in stats.items(): + try: + base = s["patch_embedding.proj.weight[:,:36]"]["mean_grad_norm"] + trk = s["patch_embedding.proj.weight[:,36:]"]["mean_grad_norm"] + ratio = trk / max(base, 1e-30) + except KeyError: + ratio = float("nan") + print(f"{model:<14} {ratio:>44.6f}") + + +if __name__ == "__main__": + main()