trackwan: 14B teacher pipeline, layer-wise track LR, and overfit tooling
Training:
- finetune.py: optional track/base param groups (WANTRACK_TRACK_GROUP) so the
track pathway can take a different LR (WANTRACK_TRACK_LR_MULT); patch_embedding
gradient masked to the track slot [:, 36:]. Raises if nothing matches.
- track_warmup.py callback: hold the pretrained DiT at lr 0 while the track
pathway warms up, then ramp it back in.
14B configs + held-alloc launchers:
- finetune_wantrack_{synth_sparse_fixed,synth_sparse_random,openvid_sparse_merged_bias,
synth_stage2}_14b + openvid_14b_{warmup,lwlr} + 1.3B stage2 configs.
- run_wan14b_held.sh (held alloc + auto-restart + rank cleanup, per-stage WANTRACK
env), acquire_nodes.sh, chain_stepAB.sh, run_stepB_seed.sh, run_stepC_merge.sh,
watch_gate.sh, run_synth_stage2_slurm.sh.
Init / diagnostics:
- convert_trackwan_init_v2.py: --pe-src to lift patch_embedding[:, 36:] and
--track-src to lift track_encoder together (co-adapted merge); sharded-source support.
- check_gate_growth.py, gradient_flow_analysis.py, gradient_flow_summarize.py.
Data pipeline: segment_tracks.py sharding; decode/segment/tracks slurm scripts;
precompute_null_text, two_track_val, cfg_ablation, lwlr_experiment.
Inference: app_multi_trace.py updates.
This commit is contained in:
shaoxiongduan
@@ -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()
|
||||
@@ -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 <checkpoint-dir> [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()
|
||||
@@ -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 <ckpt.safetensors> : 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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Executable
+25
@@ -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'"
|
||||
Executable
+25
@@ -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'"
|
||||
Executable
+25
@@ -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'"
|
||||
@@ -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"]
|
||||
|
||||
@@ -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()
|
||||
@@ -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 = ("<div style='user-select:none'>"
|
||||
"<canvas id='cvs' style='border:1px solid #888;cursor:crosshair;max-width:100%'></canvas>"
|
||||
"<div id='recstatus' style='font-weight:bold;color:#0a0;height:1.4em;margin-top:6px'>"
|
||||
"pick a clip, click <b>Load frame</b>, then hover and press <b>Space</b> to record trace #1</div>"
|
||||
"pick a clip, click <b>Load frame</b>, then choose a mode below</div>"
|
||||
"</div>")
|
||||
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)
|
||||
|
||||
@@ -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<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<N>` until Slurm says it *could* schedule it
|
||||
# 3. submit the real -N<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<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
|
||||
@@ -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
|
||||
@@ -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/^/ /'
|
||||
@@ -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"
|
||||
Executable
+43
@@ -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" "$@"
|
||||
@@ -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."
|
||||
@@ -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
|
||||
@@ -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
|
||||
+91
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user