shaoxiongduan
92cc893cfc
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.
45 lines
2.0 KiB
Bash
45 lines
2.0 KiB
Bash
#!/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
|