scripts/README_predictor.md

# 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 = **A**nchor evolution MLP(e, a, distance) -> backbone input;
**T**ransport: RoPE'd global attention from current tokens to the previous
chunk (`chunk` or `last_frame` scope) aligns the previous-chunk state;
**C**orrect: 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.