Hugging Face's logo

Cccccz
/
Causal-Forcing-a

Model card Files
xet
 Community
Causal-Forcing-a / scripts /README_predictor.md
Cccccz's picture
Cccccz
Add files using upload-large-folder tool
0fe7113 verified
preview code
|
Raw History Blame Contribute Delete
6.89 kB

Predictor pipeline on this host

All Predictor code lives in predictor_training/ and scripts/. Shell launchers share scripts/env.sh; override any value from the environment:

Variable Default Meaning
PYTHON_BIN /local/zoubin/cz/envs/self_forcing/bin/python torch 2.5.1 + flash-attn env
GPUS 4 5 6 7 physical GPU indices (0-3 are shared with other jobs)
DATA_ROOT <repo>/data offline datasets
OUTPUT_ROOT <repo>/output training / evaluation runs
SELF_FORCING_ROOT <repo>/../Self-Forcing-a if it has the FPPF evaluator, else ../Self-Forcing sibling checkout used by the evaluators
INPUT_VARIANT self_forcing Predictor input fusion: self_forcing, disca, atc

Assets are symlinked into the repo: checkpoints/chunkwise and wan_models point into ../Causal-Forcing, prompts/* into ../Self-Forcing/prompts.

Pipeline

  1. Offline dataset (scripts/build_predictor_offline_dataset.sh)

    • sweep dataset: bash scripts/run_stage1_offline_four_gpu_and_restore.sh (100 prompts, all 30 blocks, ~3 GB per prompt)
    • layer-17 dataset: bash scripts/run_layer17_1000p_offline_four_gpu_and_restore.sh (1000 prompts, block 17 only, ~0.5 GB per prompt)
    • smoke test: DATASET_NAME=_smoke_test GPUS=4 LAYERS=17 MAX_NEW_PROMPTS=1 bash scripts/build_predictor_offline_dataset.sh

    Each GPU builds a disjoint strided subset of prompt IDs directly into the dataset directory (atomic prompt_NNNN.partial -> prompt_NNNN), so a rerun resumes and never rewrites finished prompts.

  2. Stage-1 layer sweep: bash scripts/run_single_block_stage1_four_gpu_and_restore.sh -> output/single_block_stage1_layer_sweep_100p_21f/summary.csv

  3. Stage-1 layer-17 training: bash scripts/run_layer17_stage1_4gpu_acc2_and_restore.sh (per-GPU batch 16, accumulation 1, as in Self-Forcing-a; effective batch 64 on 4 GPUs)

  4. Stage-2 random-exit DMD: bash scripts/run_layer17_predictor_stage2_dmd_4gpu_and_restore.sh (formal setup 4 GPUs, wan_models/Wan2.1-T2V-14B; PREDICTOR_INIT=... SWANLAB=0 to pick the Stage-1 checkpoint and silence SwanLab). The trainer mirrors Self-Forcing-a: seeded prompt-pool permutation, --expected_world_size, metadata-validated resume; resume with scripts/run_layer17_predictor_stage2_dmd_resume1000_4gpu_restore_helios0123.sh

  5. Evaluation: run_layer17_step2000_val_four_gpu_and_restore.sh, run_moviebench100_step2000_vbench_and_restore.sh, run_layer17_stage2_step2000_ema_moviebench100_4gpu_restore.sh

    These import scripts/evaluate_single_block_fppf.py from the Self-Forcing checkout (FinalHiddenCapture, generate_rollout, frame_metrics, discover_experiments, ...); env.sh picks ../Self-Forcing-a, which ships it. scripts/evaluate_long_video_vbench.py (VBench) is still absent from both checkouts, so the VBench step is skipped with a log line.

Input variants and ATC

SingleBlockPredictor(input_variant=...) selects how the three token streams (current noisy tokens, same-chunk previous-step hidden = anchor, previous-chunk same-step hidden) are fused before the single Teacher block:

variant fusion extra inputs
self_forcing (default) LayerNorm + concat MLP (TripleFeatureFusion) optional previous-feature gate
disca same MLP without the previous-chunk channel (DualFeatureFusion) none
atc Anchor-Transport-Correct (predictor_training/atc_fusion.py) target-timestep condition tokens, anchor distance

ATC = Anchor evolution MLP(e, a, distance) -> backbone input; Transport: RoPE'd global attention from current tokens to the previous chunk (chunk or last_frame scope) aligns the previous-chunk state; Correct: zero-initialised MLP(e, a, transported, condition, distance) gated per token by a sigmoid TokenGate (initial p=0.3) and added as a bypass to anchor + residual_out(block(...)). At initialisation the Predictor therefore returns the anchor exactly. Training logs atc_* diagnostics (gate statistics, transport entropy/displacement, delta norms).

Every Stage-1 checkpoint now stores predictor_config in its safetensors metadata (predictor_training/metadata.py); Stage-2 and the evaluators rebuild the module from it, so old metadata-free checkpoints are read as the concat Predictor on layer 17.

Launchers take INPUT_VARIANT and STAGE1_EXTRA_ARGS, e.g.

INPUT_VARIANT=atc STAGE1_EXTRA_ARGS="--atc_previous_scope last_frame" \
    bash scripts/run_layer17_stage1_4gpu_acc2_and_restore.sh
INPUT_VARIANT=atc bash scripts/run_single_block_stage1_four_gpu_and_restore.sh

Online rollouts (Stage-2, evaluators) call predictor_training/online.py, which slices the KV history at chunk * TOKENS_PER_CHUNK so the Predictor sees exactly the clean earlier-chunk history it was trained on. Tests: python -m unittest tests/test_atc_fusion.py (needs one visible GPU because wan/modules/t5.py touches CUDA at import).

Dataset layout (dataset_version 2) 

prompt_NNNN/
  trajectory.safetensors        chunks 1..6: noisy latent, timestep, flow,
                                final hidden per denoising step, clean latent
  cross_attention.safetensors   text K/V per cached Teacher block
  clean_prefeatures/block_XX.safetensors   clean self-attn K input per chunk
  chunk0_context/               chunk-0 final hidden + prefeatures (context only)
  metadata.json, _SUCCESS
manifest.json, prompt_selection.json, progress.json

Note: Teacher-derived block must be unfrozen

initialize_predictor_block deep-copies a Teacher block; when the Teacher is frozen (always, at Stage-2 and evaluation time) the copy inherits requires_grad=False. Stage-1 trainers always called set_block_trainable(True), but the Stage-2 trainer inherited from Self-Forcing-a did not, so its block optimizer group was empty and only the fusion trained. Since 2026-09-04 initialize_predictor_block returns a trainable block and Stage-2 refuses to start with an empty block group (it logs the trainable parameter counts).

Stage-2 fake-score critic: LoRA by default

train_layer17_predictor_stage2_dmd.py trains the Wan2.1-1.3B fake score with LoRA (predictor_training/lora.py): rank 128 on all 10 Linears of every block (self/cross attention q,k,v,o and both FFN layers), alpha = rank, B zero-init so the critic starts identical to the base model. --fake_score_lora_rank 0 restores full fine-tuning. The critic learning rate defaults to 1e-5 with LoRA and 4e-7 for full fine-tuning (--critic_lr overrides). training_latest.pt stores only the adapters (fake_score_lora); launchers expose FAKE_SCORE_LORA_RANK, BATCH_SIZE (default 2 rollouts per rank) and GRAD_ACCUM (default 1); global batch = GPUs x BATCH_SIZE x GRAD_ACCUM.