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.

examples/train/acquire_nodes.sh

@@ -0,0 +1,80 @@
+#!/usr/bin/env bash
+# Acquire a single gang allocation of N nodes on the Slinky (k8s-backed) pool.
+#
+# Why this is not just `sbatch -N<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

examples/train/chain_stepAB.sh

@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+# Wait for step A to finish, then seed and launch step B — so the held allocation never sits
+# idle between the two overfit stages.
+#
+# "A is done" means BOTH: no step-A launcher process is alive, AND a COMPLETE checkpoint-1000
+# exists (dcp/.metadata present). Requiring both avoids launching B off a half-written
+# checkpoint if A died mid-save.
+set -uo pipefail
+WORK=/mnt/lustre/vlm-s4duan
+REPO=$WORK/FastVideo
+ALLOC="${ALLOC:-728}"
+A_JOB="${A_JOB:-wan14b_stepA_v2}"
+A_CKPT="$WORK/wantrack_14b_synth_sparse_fixed_out/checkpoint-1000"
+: "${WANDB_API_KEY:?export WANDB_API_KEY}"
+say() { echo "[chain $(date +%H:%M:%S)] $*"; }
+
+say "waiting for step A to reach a complete $A_CKPT ..."
+while :; do
+  if [ -f "$A_CKPT/dcp/.metadata" ]; then
+    # Checkpoint is complete; wait for the launcher to actually exit before reusing the nodes.
+    # Use the launcher's PID file, NOT `pgrep -f run_wan14b_held.sh` — that pattern also matches
+    # any watcher/shell whose command line contains the string, so the check never goes false
+    # and the chain waits forever on nodes that are already idle.
+    A_RUNFILE="$WORK/logs/${A_JOB}.running"
+    if [ ! -f "$A_RUNFILE" ] || ! kill -0 "$(cat "$A_RUNFILE" 2>/dev/null)" 2>/dev/null; then
+      say "step A finished and $A_CKPT is complete"
+      break
+    fi
+    say "checkpoint-1000 complete, waiting for step A launcher (pid $(cat "$A_RUNFILE")) to exit ..."
+  fi
+  [ -n "$(squeue -h -j "$ALLOC" -o '%t' 2>/dev/null)" ] || { say "ALLOC $ALLOC vanished — aborting chain"; exit 1; }
+  sleep 60
+done
+
+say "seeding step B output dir from A's checkpoint-1000"
+bash "$REPO/examples/train/run_stepB_seed.sh" || { say "seed failed"; exit 1; }
+
+say "launching step B (random-track overfit, WANTRACK_FIXED_SAMPLE=0)"
+cd "$REPO"
+exec env ALLOC="$ALLOC" NODES=4 JOB=wan14b_stepB_v2 PORT=30918 \
+  CFG=examples/train/scenario/worldmodel/finetune_wantrack_synth_sparse_random_14b.yaml \
+  WANTRACK_FIXED_SAMPLE=0 WANTRACK_FREEZE_HEAD=0 TRACKWAN_TRACK_BIAS=1 \
+  WANDB_API_KEY="$WANDB_API_KEY" \
+  bash examples/train/run_wan14b_held.sh

examples/train/run_stepB_seed.sh

@@ -0,0 +1,26 @@
+#!/usr/bin/env bash
+# Seed step B's output dir with step A's final checkpoint, so B starts from A while its config
+# can still use resume_from_checkpoint: latest (which is what makes crash-restarts safe).
+#
+# Hardlinked (cp -al), not copied: a 14B training-state checkpoint is hundreds of GB (bf16
+# weights + fp32 master + Adam m/v), and both dirs live on the same Lustre filesystem.
+# Checkpoints are write-once, so sharing inodes is safe here.
+set -euo pipefail
+WORK=/mnt/lustre/vlm-s4duan
+STEP="${STEP:-1000}"
+SRC="${SRC:-$WORK/wantrack_14b_synth_sparse_fixed_out/checkpoint-$STEP}"
+DST_DIR="${DST_DIR:-$WORK/wantrack_14b_synth_sparse_random_out}"
+DST="$DST_DIR/checkpoint-$STEP"
+
+[ -d "$SRC" ] || { echo "[B-seed] missing $SRC"; exit 1; }
+[ -f "$SRC/dcp/.metadata" ] || { echo "[B-seed] $SRC incomplete (no dcp/.metadata) — refusing"; exit 1; }
+
+if [ -d "$DST" ]; then
+  echo "[B-seed] $DST already exists — leaving it alone"
+else
+  mkdir -p "$DST_DIR"
+  echo "[B-seed] hardlinking $SRC -> $DST"
+  cp -al "$SRC" "$DST"
+fi
+echo "[B-seed] checkpoints now in $DST_DIR:"
+ls -d "$DST_DIR"/checkpoint-* 2>/dev/null | sed 's/^/  /'

examples/train/run_stepC_merge.sh

@@ -0,0 +1,74 @@
+#!/usr/bin/env bash
+# STEP C — build the 14B "merged" init from the step-B overfit.
+#
+# Takes the 14B-native overfit (steps A+B) and transplants BOTH halves of the track pathway
+# into a FRESH 14B base:
+#   * track_encoder.{proj,temporal_conv}.{weight,bias}   (--track-src)
+#   * patch_embedding.weight[:, 36:]  — the track slot   (--pe-src)
+# while patch_embedding.weight[:, :36] stays the pretrained I2V weights from --base.
+#
+# Both halves come from the SAME source on purpose: they were co-adapted during the overfit.
+# Lifting the encoder alone (the "partial_merged" experiment) measured WORSE than random init,
+# because the trained encoder ends up feeding a random projection.
+#
+# Usage: [STEP=2000] bash examples/train/run_stepC_merge.sh
+set -euo pipefail
+WORK=/mnt/lustre/vlm-s4duan
+REPO=$WORK/FastVideo
+STEP="${STEP:-2000}"
+SRC_CKPT="${SRC_CKPT:-$WORK/wantrack_14b_synth_sparse_random_out/checkpoint-$STEP}"
+EXPORT_DIR="${EXPORT_DIR:-$WORK/exports/overfit_14b_random_ckpt$STEP}"
+OUT="${OUT:-$WORK/models/trackwan_14b_i2v_d64_merged_from_overfit_bias}"
+BASE="${BASE:-$WORK/models/Wan2.1-I2V-14B-720P-Diffusers}"
+
+cd "$REPO"
+source .venv/bin/activate
+export HOME=$WORK HF_HOME=$WORK/.hf TORCH_HOME=$WORK/.torch MPLCONFIGDIR=$WORK/.mpl \
+       PYTHONPATH=$REPO TOKENIZERS_PARALLELISM=false NCCL_CUMEM_ENABLE=0 \
+       TRACKWAN_TRACK_BIAS=1
+
+[ -d "$SRC_CKPT" ] || { echo "[C] missing $SRC_CKPT"; exit 1; }
+[ -f "$SRC_CKPT/dcp/.metadata" ] || { echo "[C] $SRC_CKPT is an INCOMPLETE checkpoint (no dcp/.metadata)"; exit 1; }
+
+# 1) DCP -> diffusers (single process; DCP reshards automatically)
+if [ ! -d "$EXPORT_DIR/transformer" ]; then
+  echo "[C] exporting $SRC_CKPT -> $EXPORT_DIR"
+  python -m fastvideo.train.entrypoint.dcp_to_diffusers \
+    --checkpoint "$SRC_CKPT" --output-dir "$EXPORT_DIR"
+else
+  echo "[C] reusing existing export $EXPORT_DIR"
+fi
+
+# 2) fresh 14B base + co-adapted encoder AND track slot from that export
+echo "[C] building merged init -> $OUT"
+python data_pipeline/convert_trackwan_init_v2.py \
+  --base "$BASE" --out "$OUT" --id-dim 64 --pe-init random \
+  --track-src "$EXPORT_DIR" --pe-src "$EXPORT_DIR"
+
+# 3) verify: pretrained channels untouched, track slot + encoder match the overfit
+echo "[C] verifying merged init against source ..."
+python - "$OUT" "$EXPORT_DIR" "$BASE" <<'PY'
+import sys, glob
+import torch
+from safetensors.torch import load_file
+
+def load(d):
+    st = {}
+    for f in sorted(glob.glob(f"{d}/transformer/*.safetensors")):
+        st.update(load_file(f))
+    return st
+
+out, src, base = (load(p) for p in sys.argv[1:4])
+pe_o, pe_s, pe_b = out["patch_embedding.weight"], src["patch_embedding.weight"], base["patch_embedding.weight"]
+print(f"  pe shape {tuple(pe_o.shape)}  (base {tuple(pe_b.shape)})")
+print(f"  pe[:, :36] == base           : {torch.equal(pe_o[:, :36], pe_b)}   std={pe_o[:, :36].float().std():.5f}")
+print(f"  pe[:, 36:] == overfit slot   : {torch.equal(pe_o[:, 36:], pe_s[:, 36:])}   std={pe_o[:, 36:].float().std():.5f}")
+ok = torch.equal(pe_o[:, :36], pe_b) and torch.equal(pe_o[:, 36:], pe_s[:, 36:])
+for k in sorted(k for k in out if "track_encoder" in k):
+    same = k in src and torch.equal(out[k], src[k])
+    ok &= same
+    print(f"  {k:42s} lifted={same}  norm={out[k].float().norm():.5f}")
+print("  RESULT:", "OK — encoder and track slot are co-adapted" if ok else "MISMATCH")
+sys.exit(0 if ok else 1)
+PY
+echo "[C] done -> $OUT"

examples/train/run_synth_stage2_slurm.sh

@@ -0,0 +1,43 @@
+#!/usr/bin/env bash
+# Synth stage-2: MotionStream-style fine-tuning on 49k Wan2.2 synth data.
+# Init: merged_bias_ckpt4800 (our stage-1 end on openvid).
+# Recipe: TRACK_DROP=0.5 + MOTION_DROP=0.3 + PMASK=0.2 + freeze head.
+# 600 steps -> ~1.55 epochs on combined 49k (paper claims "~1 epoch" for stage-2).
+# Two VARIANTs: paperLR (1e-6, paper spec) or 5x (5e-6, 5x paper, 2x below stage-1's 1e-5).
+set -euo pipefail
+VARIANT="${VARIANT:-paperLR}"
+NODES="${1:-4}"; shift || true
+: "${WANDB_API_KEY:?export WANDB_API_KEY before launching}"
+
+case "$VARIANT" in
+  paperLR)
+    CFG=examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_paperLR.yaml
+    JOB=synth_stage2_paperLR
+    PORT=30700
+    ;;
+  5x)
+    CFG=examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_5x.yaml
+    JOB=synth_stage2_5x
+    PORT=30800
+    ;;
+  *) echo "unknown VARIANT=$VARIANT"; exit 1 ;;
+esac
+
+export WANTRACK_AUG=1
+export WANTRACK_SPARSE=1
+export WANTRACK_EXTRA_RANDOM=20
+export WANTRACK_EXTRA_MODE=random
+export WANTRACK_TRACK_DROP=0.5
+export WANTRACK_MOTION_DROP=0.3
+export WANTRACK_PMASK=0.2
+export WANTRACK_MASK_CHUNK=8
+export WANTRACK_TEXT_DROP=0
+export WANTRACK_FIXED_SAMPLE=0
+export WANTRACK_DEBUG="${WANTRACK_DEBUG:-1}"
+export WANTRACK_FREEZE_HEAD=1
+export TRACKWAN_TRACK_BIAS=1
+
+export WANDB_MODE=online
+
+CFG="$CFG" JOB="$JOB" PORT="$PORT" \
+  bash "$(dirname "$0")/run_openvid_bidir_held.sh" "$NODES" "$@"

examples/train/run_wan14b_held.sh

@@ -0,0 +1,137 @@
+#!/usr/bin/env bash
+# Run a Wan2.1-14B WanTrack stage inside an ALREADY-HELD Slurm allocation.
+#
+# Same held-alloc + auto-restart pattern as run_openvid_bidir_held.sh, but every WANTRACK_*
+# knob is parameterized instead of hardcoded, because the four stages of the 14B teacher
+# pipeline need DIFFERENT sampling/masking settings:
+#
+#   A (fixed overfit)  : WANTRACK_FIXED_SAMPLE=1 FREEZE_HEAD=0
+#   B (random overfit) : WANTRACK_FIXED_SAMPLE=0 FREEZE_HEAD=0
+#   D (stage-1 openvid): FREEZE_HEAD=1, no dropping
+#   E (stage-2 synth)  : FREEZE_HEAD=1 TRACK_DROP=0.5 MOTION_DROP=0.3 PMASK=0.2 MASK_CHUNK=8
+#
+# FREEZE_HEAD is the important one: the overfit stages MUST train track_encoder (that is the
+# whole point — co-adapting it with the patch_embedding track slot), while stages D/E freeze it.
+#
+# Usage:
+#   ALLOC=728 NODES=4 JOB=wan14b_stepA PORT=30910 \
+#   CFG=examples/train/scenario/worldmodel/finetune_wantrack_synth_sparse_fixed_14b.yaml \
+#   WANTRACK_FIXED_SAMPLE=1 WANTRACK_FREEZE_HEAD=0 \
+#   WANDB_API_KEY=... bash examples/train/run_wan14b_held.sh
+set -uo pipefail
+WORK=/mnt/lustre/vlm-s4duan
+REPO=$WORK/FastVideo
+CFG="${CFG:?set CFG}"
+NODES="${NODES:-4}"; GPUS=4
+JOB="${JOB:-wan14b}"
+PORT="${PORT:-30900}"
+ALLOC="${ALLOC:?set ALLOC to a running held allocation jobid}"
+WANDB_RUN_ID="${WANDB_RUN_ID:-${JOB}}"
+: "${WANDB_API_KEY:?export WANDB_API_KEY}"
+TOTAL_GPUS=$(( NODES * GPUS ))
+mkdir -p "$WORK/logs"
+
+OUTPUT_DIR=$(grep -oE 'output_dir:[^#]*' "$REPO/$CFG" | head -1 | sed 's/.*output_dir:[[:space:]]*//; s/"//g' | xargs)
+
+# Per-stage WANTRACK defaults (overridable from the environment).
+W_AUG="${WANTRACK_AUG:-1}"
+W_SPARSE="${WANTRACK_SPARSE:-1}"
+W_EXTRA_RANDOM="${WANTRACK_EXTRA_RANDOM:-20}"
+W_EXTRA_MODE="${WANTRACK_EXTRA_MODE:-random}"
+W_PMASK="${WANTRACK_PMASK:-0}"
+W_MASK_CHUNK="${WANTRACK_MASK_CHUNK:-0}"
+W_TRACK_DROP="${WANTRACK_TRACK_DROP:-0}"
+W_MOTION_DROP="${WANTRACK_MOTION_DROP:-0}"
+W_TEXT_DROP="${WANTRACK_TEXT_DROP:-0}"
+W_FIXED="${WANTRACK_FIXED_SAMPLE:-0}"
+W_FREEZE="${WANTRACK_FREEZE_HEAD:-0}"
+W_BIAS="${TRACKWAN_TRACK_BIAS:-1}"
+
+# A crash mid-save leaves checkpoint-N/dcp without .metadata; drop it so resume falls back to
+# the last COMPLETE checkpoint instead of crash-looping on "metadata is None".
+clean_bad_ckpt() {
+  [ -n "$OUTPUT_DIR" ] || return 0
+  local latest
+  latest=$(ls -d "$OUTPUT_DIR"/checkpoint-* 2>/dev/null | sed 's/.*checkpoint-//' | sort -n | tail -1)
+  [ -n "$latest" ] || return 0
+  local d="$OUTPUT_DIR/checkpoint-$latest"
+  if [ ! -f "$d/dcp/.metadata" ]; then
+    echo "[held] checkpoint-$latest incomplete (no dcp/.metadata) — removing"
+    rm -rf "$d"
+  fi
+}
+
+[ "$(squeue -h -j "$ALLOC" -o '%t' 2>/dev/null)" = R ] || { echo "[held] alloc $ALLOC not running"; exit 1; }
+NODELIST=$(squeue -h -j "$ALLOC" -o '%N')
+# Pin the exact subset of the held allocation this stage runs on. Without --nodelist, srun is
+# free to pick ANY $NODES of the allocation's nodes, so the rdzv endpoint (computed here) can
+# land on a node that isn't in the subset — rank 0 then never listens there and every worker
+# dies with "client socket has timed out while trying to connect". Choosing the subset up front
+# makes MASTER provably the node hosting rank 0, and lets stages share one allocation safely.
+# By default take the first $NODES nodes of the allocation. Pass NODELIST_OVERRIDE=a,b,c to pin
+# an explicit subset — required when several stages share one allocation, otherwise every run
+# grabs the same leading nodes and they fight over the same GPUs.
+SUBSET="${NODELIST_OVERRIDE:-$(scontrol show hostnames "$NODELIST" | head -n "$NODES" | paste -sd,)}"
+MASTER=$(echo "$SUBSET" | cut -d, -f1)
+echo "[held] alloc=$ALLOC nodes=$NODELIST"
+echo "[held] subset=$SUBSET master=$MASTER using $NODES node(s) / $TOTAL_GPUS GPU(s)"
+echo "[held] cfg=$CFG out=$OUTPUT_DIR"
+echo "[held] WANTRACK: fixed=$W_FIXED freeze=$W_FREEZE pmask=$W_PMASK chunk=$W_MASK_CHUNK track_drop=$W_TRACK_DROP motion_drop=$W_MOTION_DROP bias=$W_BIAS"
+
+# Killing this script (or its srun) does NOT kill the torchrun/python ranks out on the compute
+# nodes — they survive, keep ~107GB of GPU memory each, and silently keep writing checkpoints.
+# That poisons the next run: it contends for GPUs and resurrects deleted output dirs. So sweep
+# the subset on any exit. Matches on the config path so we only kill THIS stage's ranks.
+# Liveness marker for anything chaining off this stage. Do NOT make chainers use
+# `pgrep -f run_wan14b_held.sh`: that also matches any monitoring/shell command whose command
+# line merely CONTAINS the string, so a watcher looking for this launcher keeps "seeing" it long
+# after it exited (that stall cost ~1h of idle nodes between stages A and B).
+RUNFILE="$WORK/logs/${JOB}.running"
+echo "$$" > "$RUNFILE"
+
+cleanup_ranks() {
+  echo "[held] sweeping stray ranks for $CFG on $SUBSET ..."
+  timeout 120 srun --overlap --jobid="$ALLOC" --nodelist="$SUBSET" --nodes="$NODES" \
+    --ntasks="$NODES" --ntasks-per-node=1 \
+    bash -c "pkill -9 -f 'entrypoint/train.py --config $CFG'; exit 0" >/dev/null 2>&1 || true
+}
+trap 'echo "[held] interrupted — cleaning up"; cleanup_ranks; rm -f "$RUNFILE"; exit 130' INT TERM
+trap 'rm -f "$RUNFILE"' EXIT
+
+attempt=0
+while :; do
+  [ -n "$(squeue -h -j "$ALLOC" -o '%t' 2>/dev/null)" ] || { echo "[held] alloc $ALLOC vanished"; exit 1; }
+  attempt=$((attempt + 1))
+  clean_bad_ckpt
+  echo "=== [train] attempt $attempt on alloc $ALLOC ($(date)) ==="
+  srun --overlap --jobid="$ALLOC" --nodelist="$SUBSET" --nodes="$NODES" --ntasks="$NODES" \
+    --ntasks-per-node=1 --chdir="$REPO" bash -lc "
+      source .venv/bin/activate
+      export HOME=$WORK HF_HOME=$WORK/.hf TORCH_HOME=$WORK/.torch MPLCONFIGDIR=$WORK/.mpl \
+        TRITON_CACHE_DIR=$WORK/.cache/triton_${JOB} TORCHINDUCTOR_CACHE_DIR=$WORK/.cache/inductor_${JOB} \
+        TOKENIZERS_PARALLELISM=false NCCL_CUMEM_ENABLE=0 PYTHONPATH=$REPO \
+        WANDB_API_KEY=$WANDB_API_KEY WANDB_MODE=online \
+        WANDB_RUN_ID=$WANDB_RUN_ID WANDB_RESUME=allow \
+        WANTRACK_AUG=$W_AUG WANTRACK_SPARSE=$W_SPARSE \
+        WANTRACK_EXTRA_RANDOM=$W_EXTRA_RANDOM WANTRACK_EXTRA_MODE=$W_EXTRA_MODE \
+        WANTRACK_PMASK=$W_PMASK WANTRACK_MASK_CHUNK=$W_MASK_CHUNK \
+        WANTRACK_TRACK_DROP=$W_TRACK_DROP WANTRACK_MOTION_DROP=$W_MOTION_DROP \
+        WANTRACK_TEXT_DROP=$W_TEXT_DROP WANTRACK_FIXED_SAMPLE=$W_FIXED \
+        WANTRACK_FREEZE_HEAD=$W_FREEZE TRACKWAN_TRACK_BIAS=$W_BIAS \
+        WANTRACK_DEBUG=1 \
+        TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=1800 TORCH_NCCL_TRACE_BUFFER_SIZE=2000 \
+        NCCL_SOCKET_NTHREADS=4 NCCL_NSOCKS_PERTHREAD=8
+      torchrun --nnodes=$NODES --nproc-per-node=$GPUS --node-rank=\$SLURM_PROCID \
+        --rdzv-backend=c10d --rdzv-endpoint=$MASTER:$PORT \
+        fastvideo/train/entrypoint/train.py --config $CFG \
+        --training.distributed.num_gpus $TOTAL_GPUS
+    "
+  rc=$?
+  if [ $rc -eq 0 ]; then echo "[train] finished cleanly (rc=0)"; break; fi
+  echo "[train] exited rc=$rc — nodes still held, relaunching in 20s ..."
+  # A crashed rank can leave its siblings alive and holding GPU memory; clear them before retry
+  # or the relaunch will OOM or hang waiting on a rendezvous the stale ranks are still in.
+  cleanup_ranks
+  sleep 20
+done
+echo "[held] stage complete; allocation $ALLOC left RUNNING for the next stage."

examples/train/scenario/worldmodel/finetune_wantrack_openvid_14b_lwlr.yaml

@@ -0,0 +1,100 @@
+# EXPERIMENT 2 (layer-wise LR) — can stage-1 bootstrap a FRESH track pathway on OpenVid,
+# with no overfit->merge stages at all?
+#
+# Same goal as experiment 1, different mechanism: instead of staging, train everything from
+# step 0 but give the track pathway a 100x LR (1e-3 vs 1e-5) so it can catch up with an already
+# pretrained DiT. Adam is scale-invariant per parameter, so this MUST be a real param group with
+# its own lr — scaling the gradient would do nothing. patch_embedding's pretrained channels
+# [:, :36] are grad-masked so the 100x LR only ever touches the track slot.
+#
+# NOTE: an earlier in-repo LWLR attempt reported 'no encoder bootstrap', but it referenced
+# 'patch_embedding.proj.weight' incorrectly and so likely never applied the boosted LR to
+# anything. This re-tests the idea properly; the param-group builder now RAISES if nothing
+# matches, so a silent no-op cannot masquerade as a negative result.
+#
+# Recipe is otherwise IDENTICAL to the 1.3B stage-1 reference: lr 1e-5, global bs 128, 480x832,
+# 121f, 24fps, flow_shift 6, bf16 — only the param-group/warm-up behaviour is new.
+#
+# Init is the ZERO-gate init: patch_embedding[:, 36:] = 0, so the model starts bit-identical to
+# base I2V. (The random-gate init was measured to corrupt the prior from step 0.)
+#
+# Launch env: WANTRACK_TRACK_GROUP=1 (build named track/base param groups),
+#             WANTRACK_TRACK_LR_MULT=100 (track pathway at 1e-3), WANTRACK_FREEZE_HEAD=0
+#             (the encoder MUST train here — this is the opposite of the merged-init stage 1).
+# 4 nodes x 4 GPU = 16 GPUs, grad_accum 8 -> 1 * 8 * 16 = 128 global.
+models:
+  student:
+    _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel
+    init_from: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_zero_init_bias
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  model_path: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_zero_init_bias
+  dit_precision: bf16
+
+  distributed:
+    num_gpus: 16
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 4
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: /mnt/lustre/vlm-s4duan/openvid_1m/combined_parquet_dataset
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 31
+    num_height: 480
+    num_width: 832
+    num_frames: 121
+
+  optimizer:
+    learning_rate: 1.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 1000            # experiment: the bootstrap signal shows well before this
+    gradient_accumulation_steps: 8   # 1 * 8 * 16 gpus = 128 global
+
+  checkpoint:
+    output_dir: /mnt/lustre/vlm-s4duan/exp_openvid_14b_lwlr_out
+    training_state_checkpointing_steps: 100
+    checkpoints_total_limit: 100     # keep all — early checkpoints are the whole point
+    resume_from_checkpoint: latest
+
+  tracker:
+    trackers: [wandb]
+    project_name: wantrack-bidir
+    run_name: exp-openvid-14b-lwlr-100x
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  track_validation:
+    _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback
+    every_steps: 250
+    num_val_samples: 4
+    val_data_path: /mnt/lustre/vlm-s4duan/val_examples_mixed/combined_parquet_dataset
+    num_inference_steps: 30
+    guidance_scale: 3.0
+    motion_guidance_scale: 1.5
+    fps: 24
+    grid_stride: 3
+    tail: 12
+    validate_at_start: true
+    seed: 1000
+
+pipeline:
+  flow_shift: 6

examples/train/scenario/worldmodel/finetune_wantrack_openvid_14b_warmup.yaml

@@ -0,0 +1,100 @@
+# EXPERIMENT 1 (encoder warm-up) — can stage-1 bootstrap a FRESH track pathway on OpenVid,
+# with no overfit->merge stages at all?
+#
+# The overfit+merge pipeline exists only because a from-scratch track encoder would not train
+# during stage 1. This tests a more direct route: for the first N steps hold the pretrained DiT
+# completely still (base param group at lr 0) and let ONLY the track pathway learn
+# (track_encoder.* + patch_embedding.proj.weight[:, 36:]); then ramp the DiT back in. If the
+# gate and encoder grow to a healthy scale and the loss behaves, the whole A/B/C detour can go.
+#
+# Recipe is otherwise IDENTICAL to the 1.3B stage-1 reference: lr 1e-5, global bs 128, 480x832,
+# 121f, 24fps, flow_shift 6, bf16 — only the param-group/warm-up behaviour is new.
+#
+# Init is the ZERO-gate init: patch_embedding[:, 36:] = 0, so the model starts bit-identical to
+# base I2V. (The random-gate init was measured to corrupt the prior from step 0.)
+#
+# Launch env: WANTRACK_TRACK_GROUP=1 (build named track/base param groups),
+#             WANTRACK_TRACK_LR_MULT=1 (keep the reference LR), WANTRACK_FREEZE_HEAD=0
+#             (the encoder MUST train here — this is the opposite of the merged-init stage 1).
+# 4 nodes x 4 GPU = 16 GPUs, grad_accum 8 -> 1 * 8 * 16 = 128 global.
+models:
+  student:
+    _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel
+    init_from: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_zero_init_bias
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  model_path: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_zero_init_bias
+  dit_precision: bf16
+
+  distributed:
+    num_gpus: 16
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 4
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: /mnt/lustre/vlm-s4duan/openvid_1m/combined_parquet_dataset
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 31
+    num_height: 480
+    num_width: 832
+    num_frames: 121
+
+  optimizer:
+    learning_rate: 1.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 1000            # experiment: the bootstrap signal shows well before this
+    gradient_accumulation_steps: 8   # 1 * 8 * 16 gpus = 128 global
+
+  checkpoint:
+    output_dir: /mnt/lustre/vlm-s4duan/exp_openvid_14b_warmup_out
+    training_state_checkpointing_steps: 100
+    checkpoints_total_limit: 100     # keep all — early checkpoints are the whole point
+    resume_from_checkpoint: latest
+
+  tracker:
+    trackers: [wandb]
+    project_name: wantrack-bidir
+    run_name: exp-openvid-14b-encoder-warmup
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  track_warmup:
+    _target_: fastvideo.train.callbacks.track_warmup.TrackWarmupCallback
+    warmup_steps: 300     # DiT held at lr 0 while the track pathway learns
+    ramp_steps: 200       # then linearly hand control back to the full model
+    log_every: 25
+  track_validation:
+    _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback
+    every_steps: 250
+    num_val_samples: 4
+    val_data_path: /mnt/lustre/vlm-s4duan/val_examples_mixed/combined_parquet_dataset
+    num_inference_steps: 30
+    guidance_scale: 3.0
+    motion_guidance_scale: 1.5
+    fps: 24
+    grid_stride: 3
+    tail: 12
+    validate_at_start: true
+    seed: 1000
+
+pipeline:
+  flow_shift: 6

examples/train/scenario/worldmodel/finetune_wantrack_openvid_sparse_14b_merged_bias.yaml

@@ -0,0 +1,91 @@
+# STEP D — MotionStream stage-1 teacher training on the full OpenVid-1M bidir set, Wan2.1-14B.
+#
+# Direct 14B analog of finetune_wantrack_openvid_sparse_1p3b_merged_bias.yaml (wandb run
+# openvid_bidir_1p3b_merged_bias). Hparams are copied EXACTLY from that run — lr 1e-5,
+# global bs 128, 4800 steps, 480x832, 121f, 24fps, flow_shift 6, bf16 DiT — the only changes
+# are the ones 14B geometry forces (hsdp dims, grad_accum to keep bs 128 at 32 GPUs).
+#
+# init_from is the step-C MERGED init: a fresh 14B base whose track_encoder AND
+# patch_embedding.weight[:, 36:] both come from the 14B-native overfit (steps A+B), so the two
+# are co-adapted. This is the whole point of the pipeline: gradient-SNR analysis showed a
+# random 14B init gives per-sample track_encoder.proj grad ~0.011 (the same regime that made
+# 1.3B-random FAIL), while 1.3B-merged sat at ~0.502 and succeeded.
+#
+# Launched with WANTRACK_FREEZE_HEAD=1 (track_encoder frozen) + TRACKWAN_TRACK_BIAS=1.
+# 8 nodes x 4 GPU = 32 GPUs; grad_accum 4 -> 1 * 4 * 32 = 128 global.
+models:
+  student:
+    _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel
+    init_from: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_merged_from_overfit_bias
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  model_path: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_merged_from_overfit_bias
+  dit_precision: bf16
+
+  distributed:
+    num_gpus: 32          # overridden by the launcher to nodes*4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 8 # replicate across the 8 nodes
+    hsdp_shard_dim: 4     # shard the 14B within each node over NVLink (8*4 = 32)
+
+  data:
+    data_path: /mnt/lustre/vlm-s4duan/openvid_1m/combined_parquet_dataset
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 31
+    num_height: 480
+    num_width: 832
+    num_frames: 121
+
+  optimizer:
+    learning_rate: 1.0e-5     # MotionStream teacher stage-1
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 4800
+    gradient_accumulation_steps: 4   # 1 * 4 * 32 gpus = 128 global
+
+  checkpoint:
+    output_dir: /mnt/lustre/vlm-s4duan/openvid_bidir_14b_merged_bias_out
+    training_state_checkpointing_steps: 100   # ~backstop vs the ~80min cordon cycle
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: latest
+
+  tracker:
+    trackers: [wandb]
+    project_name: wantrack-bidir
+    run_name: openvid-bidir-14b-merged-init-bias
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  track_validation:
+    _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback
+    every_steps: 250
+    num_val_samples: 4
+    val_data_path: /mnt/lustre/vlm-s4duan/val_examples_mixed/combined_parquet_dataset   # 2 openvid + car + dog
+    num_inference_steps: 30
+    guidance_scale: 3.0
+    motion_guidance_scale: 1.5
+    fps: 24
+    grid_stride: 3
+    tail: 12
+    validate_at_start: true
+    seed: 1000
+
+pipeline:
+  flow_shift: 6

examples/train/scenario/worldmodel/finetune_wantrack_synth_sparse_fixed_14b.yaml

@@ -0,0 +1,98 @@
+# STEP A — 14B fixed-track overfit on the 20-clip synth toy set.
+#
+# Direct 14B analog of finetune_wantrack_synth_sparse_fixed.yaml (1.3B). Purpose is NOT a good
+# model: it is to CO-ADAPT track_encoder.{proj,temporal_conv} with patch_embedding.weight[:, 36:]
+# for the 14B DiT. Gradient-SNR analysis showed a 14B random init has per-sample
+# track_encoder.proj grad ~0.011 (same as the 1.3B random init that FAILED), while the 1.3B
+# merged-from-overfit init sat at ~0.502 (50x) and succeeded. Steps A+B produce the 14B-native
+# overfit whose encoder + track slot get merged into a fresh init in step C.
+#
+# Launched with WANTRACK_FIXED_SAMPLE=1 (deterministic per-sample track subset).
+#
+# GLOBAL BATCH SIZE 1 — this is exact parity with c0xyfg57 and is NOT negotiable for this stage.
+# The reference ran 1 GPU, bs 1, accum 1 => global bs 1, 2000 steps => 2000 samples seen. An
+# earlier attempt ran this on 16 GPUs at global bs 16 for 1000 steps = 16,000 samples, i.e. 8x
+# the reference's ENTIRE run, and the model came out visibly overbaked by step ~1000. Step count
+# only means something relative to batch size; matching lr + steps while changing batch silently
+# changes exposure by the batch ratio. Large batch also removes gradient noise, so it memorises
+# 20 clips faster than the reference even at equal samples.
+#
+# One node is still needed to hold a 14B: hsdp_shard_dim 4 shards ONE replica across the 4 GPUs
+# and hsdp_replicate_dim 1 makes the data-parallel degree 1, so global bs stays 1.
+# checkpoints_total_limit is deliberately huge: rotation previously destroyed the early
+# checkpoints we wanted to fall back to.
+models:
+  student:
+    _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel
+    init_from: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_zero_init_bias
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  model_path: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_zero_init_bias
+  dit_precision: bf16
+
+  distributed:
+    num_gpus: 4           # overridden by the launcher to nodes*4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1 # data-parallel degree 1 -> global bs = train_batch_size = 1
+    hsdp_shard_dim: 4     # shard ONE 14B replica across the node's 4 GPUs over NVLink
+
+  data:
+    data_path: /mnt/lustre/vlm-s4duan/data/wantrack_synth_toy/preprocessed_i2v_track_varied/combined_parquet_dataset
+    dataloader_num_workers: 2
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 31
+    num_height: 480
+    num_width: 832
+    num_frames: 121
+
+  optimizer:
+    learning_rate: 1.0e-4   # same as the 1.3B overfit (c0xyfg57 lineage)
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 2000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: /mnt/lustre/vlm-s4duan/wantrack_14b_synth_sparse_fixed_out
+    training_state_checkpointing_steps: 100
+    checkpoints_total_limit: 100
+    resume_from_checkpoint: latest
+
+  tracker:
+    trackers: [wandb]
+    project_name: wantrack-bidir
+    run_name: synth-sparse-fixed-14b
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  track_validation:
+    _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback
+    every_steps: 250
+    num_val_samples: 2      # validates on the train (toy) clips — the point of an overfit
+    num_inference_steps: 30
+    guidance_scale: 3.0
+    motion_guidance_scale: 1.5
+    fps: 24
+    grid_stride: 3
+    tail: 12
+    validate_at_start: true
+    seed: 1000
+
+pipeline:
+  flow_shift: 6

examples/train/scenario/worldmodel/finetune_wantrack_synth_sparse_random_14b.yaml

@@ -0,0 +1,91 @@
+# STEP B — 14B RANDOM-track overfit, warm-started from step A checkpoint-1000.
+#
+# Direct 14B analog of finetune_wantrack_synth_sparse_fixed.yaml (1.3B). Purpose is NOT a good
+# model: it is to CO-ADAPT track_encoder.{proj,temporal_conv} with patch_embedding.weight[:, 36:]
+# for the 14B DiT. Gradient-SNR analysis showed a 14B random init has per-sample
+# track_encoder.proj grad ~0.011 (same as the 1.3B random init that FAILED), while the 1.3B
+# merged-from-overfit init sat at ~0.502 (50x) and succeeded. Steps A+B produce the 14B-native
+# overfit whose encoder + track slot get merged into a fresh init in step C.
+#
+# Launched with WANTRACK_FIXED_SAMPLE=1 (deterministic per-sample track subset).
+# 4 nodes x 4 GPU = 16 GPUs, per-gpu bs 1, grad_accum 1 -> global bs 16 (20-clip set, so each
+# step sees ~80% of the data: effectively full-batch, which is what we want for an overfit).
+models:
+  student:
+    _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel
+    init_from: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_zero_init_bias
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  model_path: /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_zero_init_bias
+  dit_precision: bf16
+
+  distributed:
+    num_gpus: 4           # overridden by the launcher to nodes*4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1 # data-parallel degree 1 -> global bs = train_batch_size = 1 (parity with c0xyfg57)
+    hsdp_shard_dim: 4     # shard ONE 14B replica across the node's 4 GPUs
+
+  data:
+    data_path: /mnt/lustre/vlm-s4duan/data/wantrack_synth_toy/preprocessed_i2v_track_varied/combined_parquet_dataset
+    dataloader_num_workers: 2
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 31
+    num_height: 480
+    num_width: 832
+    num_frames: 121
+
+  optimizer:
+    learning_rate: 1.0e-4   # same as the 1.3B overfit (c0xyfg57 lineage)
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 2000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: /mnt/lustre/vlm-s4duan/wantrack_14b_synth_sparse_random_out
+    training_state_checkpointing_steps: 100
+    checkpoints_total_limit: 100
+    # NOT a hardcoded path to step A's checkpoint: the launcher auto-restarts on any crash, and a
+    # fixed path would make EVERY restart resume from A again, silently discarding B's progress.
+    # run_stepB_seed.sh hardlinks A's checkpoint-1000 into THIS output_dir first, so 'latest'
+    # resolves to it on the first launch and to B's own newest checkpoint on every restart.
+    resume_from_checkpoint: latest
+
+  tracker:
+    trackers: [wandb]
+    project_name: wantrack-bidir
+    run_name: synth-sparse-random-14b-warmstart
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  track_validation:
+    _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback
+    every_steps: 250
+    num_val_samples: 2      # validates on the train (toy) clips — the point of an overfit
+    num_inference_steps: 30
+    guidance_scale: 3.0
+    motion_guidance_scale: 1.5
+    fps: 24
+    grid_stride: 3
+    tail: 12
+    validate_at_start: true
+    seed: 1000
+
+pipeline:
+  flow_shift: 6

examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_14b_paperLR.yaml

@@ -0,0 +1,92 @@
+# STEP E — MotionStream stage-2 finetune WITH dropping on the combined synth set, Wan2.1-14B.
+#
+# Direct 14B analog of finetune_wantrack_synth_stage2_paperLR.yaml (wandb run
+# synth_stage2_paperLR). Hparams copied EXACTLY from that run — lr 1e-6, global bs 128,
+# 600 steps, 480x832, 121f, 24fps, flow_shift 6, bf16 DiT — the only changes are the ones
+# 14B geometry forces (hsdp dims, grad_accum to keep bs 128 at 32 GPUs).
+#
+# init_from is the step-D stage-1 result (OpenVid 4800-step ckpt, exported to diffusers).
+#
+# Stage 2 is where the stochastic dropping comes in — it is what teaches the model to cope
+# with sparse/partial track conditioning at inference. Set at launch, NOT in this file:
+#   WANTRACK_TRACK_DROP=0.5 WANTRACK_MOTION_DROP=0.3 WANTRACK_PMASK=0.2 WANTRACK_MASK_CHUNK=8
+#   WANTRACK_FREEZE_HEAD=1 TRACKWAN_TRACK_BIAS=1
+# Checkpoints every 50 steps and total_limit 100 so NOTHING rotates away: step-400 is the
+# MotionStream-paper-parity point and must survive to be evaluated alongside step-600.
+# 8 nodes x 4 GPU = 32 GPUs; grad_accum 4 -> 1 * 4 * 32 = 128 global.
+models:
+  student:
+    _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel
+    init_from: /mnt/lustre/vlm-s4duan/exports/openvid_14b_merged_bias_ckpt4800
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  model_path: /mnt/lustre/vlm-s4duan/exports/openvid_14b_merged_bias_ckpt4800
+  dit_precision: bf16
+
+  distributed:
+    num_gpus: 32          # overridden by the launcher to nodes*4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 8 # replicate across the 8 nodes
+    hsdp_shard_dim: 4     # shard the 14B within each node over NVLink (8*4 = 32)
+
+  data:
+    data_path: /mnt/lustre/vlm-s4duan/data/combined_synth_parquets
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 31
+    num_height: 480
+    num_width: 832
+    num_frames: 121
+
+  optimizer:
+    learning_rate: 1.0e-6     # MotionStream teacher stage-2
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 600
+    gradient_accumulation_steps: 4   # 1 * 4 * 32 gpus = 128 global
+
+  checkpoint:
+    output_dir: /mnt/lustre/vlm-s4duan/synth_bidir_14b_stage2_paperLR_out
+    training_state_checkpointing_steps: 50    # every 50 steps -> 12 ckpts, so step-400 (paper parity) is preserved
+    checkpoints_total_limit: 100             # keep ALL ckpts
+    resume_from_checkpoint: null
+
+  tracker:
+    trackers: [wandb]
+    project_name: wantrack-bidir
+    run_name: synth-bidir-14b-stage2-paperLR
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  track_validation:
+    _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback
+    every_steps: 250
+    num_val_samples: 4
+    val_data_path: /mnt/lustre/vlm-s4duan/val_examples_mixed/combined_parquet_dataset   # 2 openvid + car + dog
+    num_inference_steps: 30
+    guidance_scale: 3.0
+    motion_guidance_scale: 1.5
+    fps: 24
+    grid_stride: 3
+    tail: 12
+    validate_at_start: true
+    seed: 1000
+
+pipeline:
+  flow_shift: 6

examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_5x.yaml

@@ -0,0 +1,85 @@
+# MotionStream-style INITIAL teacher training (stage 1) on the full OpenVid-1M bidir set.
+# Recipe: sparse point conditioning (1-per-SAM-object + extras) from run c0xyfg57
+# (synth-sparse-random, d64_nobias init) but with MotionStream stage-1 hparams:
+# batch size 128, lr 1e-5, 480p (480x832), 24fps, 121f. NO stochastic track masking,
+# NO fixed/overfit sampling (set via WANTRACK_* env at launch — see run_openvid_bidir_slurm.sh).
+# num_gpus + hsdp dims below assume 4 nodes x 4 GPU = 16 GPUs (grad_accum 8 -> global bs 128).
+# Runs CONCURRENTLY with the 14B run (8 nodes) = 12 nodes total <= 14 usable. The small 1.3B
+# gets fewer nodes; grad_accum keeps global bs at 128 regardless of GPU count.
+# DiT runs bf16 (fp32 full-scale would be ~70h; verified 1-GPU fp32 ~11s/it in smoke).
+models:
+  student:
+    _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel
+    init_from: /mnt/lustre/vlm-s4duan/exports/merged_bias_ckpt4800
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  model_path: /mnt/lustre/vlm-s4duan/exports/merged_bias_ckpt4800
+  dit_precision: bf16     # fp32 was ~11s/it on 1 GPU (=~70h full); bf16 for throughput
+
+  distributed:
+    num_gpus: 16          # overridden by run_slurm.sh to nodes*NUM_GPUS
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 4 # replicate across the 4 nodes
+    hsdp_shard_dim: 4     # shard within each node over NVLink (4*4 = 16)
+
+  data:
+    data_path: /mnt/lustre/vlm-s4duan/data/combined_synth_parquets
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 31
+    num_height: 480
+    num_width: 832
+    num_frames: 121
+
+  optimizer:
+    learning_rate: 5.0e-6     # MotionStream teacher stage-2
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 600             # synth stage-2: ~1.55 epochs on 49k combined synth; step-400 ckpt preserved for paper-parity comparison
+    gradient_accumulation_steps: 8   # 1 (per-gpu) * 8 * 16 gpus = 128 global
+
+  checkpoint:
+    output_dir: /mnt/lustre/vlm-s4duan/synth_bidir_1p3b_stage2_5x_out
+    training_state_checkpointing_steps: 50   # every 50 steps -> 12 ckpts total
+    checkpoints_total_limit: 100             # keep ALL ckpts (12 << 100) so we can evaluate step-400 too
+    resume_from_checkpoint: null
+
+  tracker:
+    trackers: [wandb]
+    project_name: wantrack-bidir
+    run_name: synth-bidir-1p3b-stage2-5x
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  track_validation:
+    _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback
+    every_steps: 250
+    num_val_samples: 4
+    val_data_path: /mnt/lustre/vlm-s4duan/val_examples_mixed/combined_parquet_dataset   # fixed car+dog synth clips
+    num_inference_steps: 30
+    guidance_scale: 3.0
+    motion_guidance_scale: 1.5
+    fps: 24
+    grid_stride: 3
+    tail: 12
+    validate_at_start: true
+    seed: 1000
+
+pipeline:
+  flow_shift: 6

examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_paperLR.yaml

@@ -0,0 +1,85 @@
+# MotionStream-style INITIAL teacher training (stage 1) on the full OpenVid-1M bidir set.
+# Recipe: sparse point conditioning (1-per-SAM-object + extras) from run c0xyfg57
+# (synth-sparse-random, d64_nobias init) but with MotionStream stage-1 hparams:
+# batch size 128, lr 1e-5, 480p (480x832), 24fps, 121f. NO stochastic track masking,
+# NO fixed/overfit sampling (set via WANTRACK_* env at launch — see run_openvid_bidir_slurm.sh).
+# num_gpus + hsdp dims below assume 4 nodes x 4 GPU = 16 GPUs (grad_accum 8 -> global bs 128).
+# Runs CONCURRENTLY with the 14B run (8 nodes) = 12 nodes total <= 14 usable. The small 1.3B
+# gets fewer nodes; grad_accum keeps global bs at 128 regardless of GPU count.
+# DiT runs bf16 (fp32 full-scale would be ~70h; verified 1-GPU fp32 ~11s/it in smoke).
+models:
+  student:
+    _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel
+    init_from: /mnt/lustre/vlm-s4duan/exports/merged_bias_ckpt4800
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  model_path: /mnt/lustre/vlm-s4duan/exports/merged_bias_ckpt4800
+  dit_precision: bf16     # fp32 was ~11s/it on 1 GPU (=~70h full); bf16 for throughput
+
+  distributed:
+    num_gpus: 16          # overridden by run_slurm.sh to nodes*NUM_GPUS
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 4 # replicate across the 4 nodes
+    hsdp_shard_dim: 4     # shard within each node over NVLink (4*4 = 16)
+
+  data:
+    data_path: /mnt/lustre/vlm-s4duan/data/combined_synth_parquets
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 31
+    num_height: 480
+    num_width: 832
+    num_frames: 121
+
+  optimizer:
+    learning_rate: 1.0e-6     # MotionStream teacher stage-2
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 600             # synth stage-2: ~1.55 epochs on 49k combined synth; step-400 ckpt preserved for paper-parity comparison
+    gradient_accumulation_steps: 8   # 1 (per-gpu) * 8 * 16 gpus = 128 global
+
+  checkpoint:
+    output_dir: /mnt/lustre/vlm-s4duan/synth_bidir_1p3b_stage2_paperLR_out
+    training_state_checkpointing_steps: 50   # every 50 steps -> 12 ckpts total
+    checkpoints_total_limit: 100             # keep ALL ckpts (12 << 100) so we can evaluate step-400 too
+    resume_from_checkpoint: null
+
+  tracker:
+    trackers: [wandb]
+    project_name: wantrack-bidir
+    run_name: synth-bidir-1p3b-stage2-paperLR
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  track_validation:
+    _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback
+    every_steps: 250
+    num_val_samples: 4
+    val_data_path: /mnt/lustre/vlm-s4duan/val_examples_mixed/combined_parquet_dataset   # fixed car+dog synth clips
+    num_inference_steps: 30
+    guidance_scale: 3.0
+    motion_guidance_scale: 1.5
+    fps: 24
+    grid_stride: 3
+    tail: 12
+    validate_at_start: true
+    seed: 1000
+
+pipeline:
+  flow_shift: 6

examples/train/watch_gate.sh

@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+# Sample the WanTrack gate at each new checkpoint and append to a TSV, so we can see whether it
+# is on track for the ~0.011 the successful 1.3B merged init reached. Checkpoints rotate
+# (total_limit=3), so poll faster than they are produced or the early trajectory is lost.
+set -uo pipefail
+WORK=/mnt/lustre/vlm-s4duan; REPO=$WORK/FastVideo
+OUT="${OUT:-$WORK/wantrack_14b_synth_sparse_fixed_out}"
+LOG="${LOG:-$WORK/logs/gate_growth.tsv}"
+cd "$REPO"; source .venv/bin/activate
+export HOME=$WORK PYTHONPATH=$REPO NCCL_CUMEM_ENABLE=0 RANK=0 LOCAL_RANK=0 WORLD_SIZE=1 \
+       MASTER_ADDR=127.0.0.1 MASTER_PORT=29600
+touch "$LOG"
+while :; do
+  for d in "$OUT"/checkpoint-*; do
+    [ -f "$d/dcp/.metadata" ] || continue
+    step=$(basename "$d" | sed 's/checkpoint-//')
+    grep -q "^${step}	" "$LOG" 2>/dev/null && continue
+    line=$(python data_pipeline/check_gate_growth.py "$d" 2>/dev/null | grep -oE "gate\[:, 36:\] std=[0-9.]+" | grep -oE "[0-9.]+$")
+    [ -n "$line" ] && { printf "%s\t%s\t%s\n" "$step" "$line" "$(date +%H:%M:%S)" >> "$LOG"; echo "step $step gate=$line"; }
+  done
+  pgrep -f "run_wan14b_held.sh" >/dev/null || { echo "[watch] no launcher; exiting"; break; }
+  sleep 240
+done