nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

trackwan: 14B teacher pipeline, layer-wise track LR, and overfit tooling

Browse Source 
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
2026-07-23 03:52:56 +00:00
parent ac919b1c87
commit 92cc893cfc
31 changed files with 3458 additions and 174 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/inference/gradio/trackwan/app_multi_trace.py
+124 -174
View File
@@ -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)
examples/train/acquire_nodes.sh
+80
View File
@@ -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
+44
View File
@@ -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
+26
View File
@@ -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
+74
View File
@@ -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
Executable
+43
View File
@@ -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
+137
View File
@@ -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
+100
View File
@@ -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
+100
View File
@@ -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
+91
View File
@@ -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
+98
View File
@@ -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
+91
View File
@@ -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
+92
View File
@@ -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
+85
View File
@@ -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
+85
View File
@@ -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
+23
View File
@@ -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