[perf]: QAT/QAD add safe regional compile for modular training

[jzhang38]

## Purpose

Add opt-in regional torch.compile support to the YAML-driven modular Wan
trainer, including compiled forward and backward execution with FA4, FSDP, and
activation checkpointing.

Fix an Inductor backward output-assembly issue that produced invalid Wan
modulation gradients, and preserve eager-like BF16 numerical behavior without
giving up steady-state compile performance.

Fixes #

Changes

• Add modular trainer YAML options:
  • training.model.enable_torch_compile
  • training.model.torch_compile_kwargs
• Apply activation checkpointing before FSDP for trainable Wan transformers.
• Compile repeated DiT blocks after FSDP setup while keeping:
  • FSDP hooks outside compiled regions
  • activation-checkpoint control flow eager
  • checkpoint state-dict names unchanged
• Require fullgraph=True to prevent silent partial-graph fallback.
• Compile both gradient-bearing and no-grad block calls.
• Enable attention compilation by default, with
  FASTVIDEO_DISABLE_ATTENTION_COMPILE=1 as a debugging escape hatch.
• Make FA4 forward and backward available through opaque custom operators:
  • explicitly retain LSE for AOTAutograd
  • support fixed-length and variable-length backward
• Add opaque Wan modulation forward/backward operators:
  • materialize all six modulation outputs independently
  • explicitly assemble all six backward gradient slices
  • avoid the incomplete Inductor output buffer that previously corrupted
    scale_shift_table and timestep-embedding gradients
• Enable Inductor emulate_precision_casts by default for regional training
  compile so fused BF16 operators retain eager's intermediate rounding points.
  Users may override it under torch_compile_kwargs.options.
• Add compile policy, config propagation, checkpoint ordering, FA4 backward,
  Wan modulation gradient, and state-dict regression tests.
• Document setup, warmup behavior, FA4 opt-in, numerical policy, and debugging
  options.

Test Plan

.venv/bin/pre-commit run --all-files

PYTHONPATH=/tmp/fastvideo-memory-cap \
FASTVIDEO_BENCH_MEMORY_CAP_GIB=73 \
.venv/bin/pytest -q \
  fastvideo/tests/attention/test_compile_policy.py \
  fastvideo/tests/attention/test_flash_attn_cute_custom_op.py \
  fastvideo/tests/models/test_wan_modulation_custom_op.py \
  fastvideo/tests/train/models/test_wan_compile_setup.py \
  fastvideo/tests/train/utils/test_config.py \
  fastvideo/tests/train/utils/test_moduleloader_attention_backend.py \
  fastvideo/tests/train/utils/test_torch_compile.py

PYTHONPATH=/tmp/fastvideo-memory-cap \
FASTVIDEO_BENCH_MEMORY_CAP_GIB=73 \
.venv/bin/pytest -q \
  fastvideo/tests/train/models \
  fastvideo/tests/train/utils

# Run once with enable_torch_compile=false and once with true.
PYTHONPATH=/tmp/fastvideo-memory-cap \
FASTVIDEO_BENCH_MEMORY_CAP_GIB=73 \
FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN \
FASTVIDEO_FA4=1 \
NUM_GPUS=4 \
WANDB_MODE=offline \
bash examples/train/run.sh \
  examples/train/configs/distribution_matching/wan/dmd2_t2v.yaml \
  --training.loop.max_train_steps 21 \
  --training.data.dataloader_num_workers 0 \
  --training.checkpoint.training_state_checkpointing_steps 0 \
  --callbacks.validation.every_steps 0 \
  --training.model.enable_torch_compile true

TORCH_LOGS=graph_breaks,recompiles \
PYTHONPATH=/tmp/fastvideo-memory-cap \
FASTVIDEO_BENCH_MEMORY_CAP_GIB=73 \
FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN \
FASTVIDEO_FA4=1 \
NUM_GPUS=4 \
WANDB_MODE=offline \
bash examples/train/run.sh \
  examples/train/configs/distribution_matching/wan/dmd2_t2v.yaml \
  --training.loop.max_train_steps 6 \
  --training.data.dataloader_num_workers 0 \
  --training.checkpoint.training_state_checkpointing_steps 0 \
  --callbacks.validation.every_steps 0 \
  --training.model.enable_torch_compile true

## Test Results
<details>
<summary>Test output and performance</summary>

pre-commit run --all-files
All hooks passed.

47 passed, 686 warnings in 63.66s
102 passed, 15 warnings in 422.99s

All GPU runs used an allocator limit of 73 GiB per device.

FA4 was selected explicitly with:

FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
FASTVIDEO_FA4=1

The runtime confirmed:

Using FlashAttention-4 backend

### Performance

Four-GB200, 21-step DMD2 comparison using the same seed and configuration:

 Metric                         Eager    Compile, warm cache             Change
━━━━━━━━━━━━━━━━━━━━━━━━━  ━━━━━━━━━━━  ━━━━━━━━━━━━━━━━━━━━━  ━━━━━━━━━━━━━━━━━
 21-step total              101.606 s               82.207 s    -19.09% latency
─────────────────────────  ───────────  ─────────────────────  ─────────────────
 Steady-state mean            3.877 s                2.812 s    -27.48% latency
─────────────────────────  ───────────  ─────────────────────  ─────────────────
 Steady-state throughput            —                      —            +37.90%
─────────────────────────  ───────────  ─────────────────────  ─────────────────
 Regular critic step          3.108 s                2.297 s            -26.10%
─────────────────────────  ───────────  ─────────────────────  ─────────────────
 Student-update step          6.951 s                4.869 s            -29.95%
─────────────────────────  ───────────  ─────────────────────  ─────────────────
 First step                  24.065 s               25.975 s     compile warmup

A completely cold Inductor cache increased the first step to 43.662 seconds
and the 21-step total to 103.136 seconds. The compile path is therefore intended
for longer training jobs where the steady-state improvement amortizes the
one-time kernel-generation cost.

### Gradient correctness

Before the fix, a one-step compiled critic update produced:

Eager critic grad norm:     0.2645
Compiled critic grad norm:  2870.1

The generated Inductor backward buffer wrote only four of the six Wan
modulation gradient slices. The missing slices left uninitialized data in the
combined timestep/modulation gradient.

After the fix:

- all 21 critic gradients and all 4 student gradients were finite
- no NaN or Inf was observed
- no FSDP collective timeout occurred
- the isolated rank-3 and rank-4 modulation regression tests matched eager
  gradients exactly

Twenty-one-step numerical comparisons after enabling
emulate_precision_casts:

 Metric                              Eager-to-eager variation    Compile-to-eager variation
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━  ━━━━━━━━━━━━━━━━━━━━━━━━━━  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Total loss, mean relative                              1.08%                   0.76%–1.03%
──────────────────────────────────  ──────────────────────────  ────────────────────────────
 Critic grad norm, mean relative                        2.46%                   1.30%–2.38%
──────────────────────────────────  ──────────────────────────  ────────────────────────────
 Student grad norm, mean relative                       1.27%                   0.42%–2.37%

The remaining differences are in the same low-single-digit range as repeated
BF16/FA4 eager runs.

TORCH_LOGS=graph_breaks,recompiles reported:

Graph breaks: 0
Steady-state recompiles: 0

Each rank created the expected initial grad/no-grad and block-type
specializations only.
</details>

## Checklist

- [x] I ran pre-commit run --all-files and fixed all issues
- [x] I added or updated tests for my changes
- [x] I updated documentation if needed
- [x] I considered GPU memory impact of my changes

Model/pipeline quality checks are not applicable: this changes only the
opt-in modular training execution strategy, not model weights, inference
pipeline wiring, or generated-output references. SSIM baselines and the support
matrix do not require updates.

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI

	Protection	Waiting on
🔴	PR merge requirements	👀 reviews and 🤖 CI

🔴 PR merge requirements

Waiting for

• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success=full-suite-passed

This rule is failing.

• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success=full-suite-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[mergify]

This PR has merge conflicts with the base branch. Please rebase:

git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease