[bugfix]: fix modular ops activation checkpoint recomputation

[shaoxiongduan]

Purpose

Fix the modular trainer's ops activation-checkpoint policy. It previously
wrapped the transformer root and retained only a small set of PyTorch op
objects, so custom attention backends silently fell back to full attention
recomputation. Several modular model plugins also still imported the separate
legacy-training helper, which meant the corrected policy reached MiniMax H3
but not the other modular families.

Changes

• Wrap each discovered transformer block in its own checkpoint region.
• Match supported fused attention operations by exact dispatcher name, avoiding
  import-time dependencies on optional backend registrations.
• Retain fused SDPA, FastVideo FlashAttention, and all autograd-enabled generic
  block-sparse attention forwards, including the current SM90, SM100a, and
  Triton VSA ops.
• Do not retain parameter collectives or inexpensive matrix multiplies.
• Remove the ineffective math-SDPA entry: PyTorch decomposes that path before
  the selective policy sees a retainable attention boundary.
• Route Wan and its descendants, LTX-2, Kandinsky5, and MiniMax H3 through the
  modular trainer's checkpoint helper; the legacy training stack is unchanged.
• Resolve Wan checkpoint configuration once so role-level overrides also enable
  checkpoint-safe causal KV-cache snapshots.
• Keep kernel registration drift coverage portable when fastvideo_kernel is
  unavailable on a supported platform.
• Document the full versus ops memory/runtime tradeoff and correct the
  canonical example-config path.

Supported and fallback paths

ops retains supported dispatcher-visible fused attention outputs and can
reduce backward recomputation when memory headroom exists. It uses more
activation memory than full. Math SDPA, VMoBA, FA3, and ATTN_QAT_TRAIN do
not expose a retainable dispatcher boundary today, so their attention still
fully recomputes under ops. Memory-constrained jobs, including tight GB10
configurations, should retain full.

Regression coverage

The CPU behavioral suite uses a real autograd-enabled custom op across four
transformer blocks. It verifies:

• forward counts: 4 without checkpointing, 8 under full, 4 when ops
  retains the custom op, and 8 for an unsupported op;
• exact output, loss, input-gradient, and parameter-gradient parity;
• per-block wrapper granularity and invalid-layout failure behavior;
• known and dynamically registered block-sparse op coverage;
• modular model-plugin wiring; and
• role-level versus global Wan causal checkpoint/cache resolution.

Fresh verification on rebased head cce7fdf819a164b616cf946ac0e1889033f66f1e:

pytest fastvideo/tests/train/utils/test_activation_checkpoint.py -q
19 passed, 14 warnings in 0.26s

pytest fastvideo/tests/train/utils -q
98 passed, 14 warnings in 0.46s

pytest fastvideo/tests/train \
  --ignore=fastvideo/tests/train/models \
  --ignore=fastvideo/tests/train/methods -q
186 passed, 14 warnings in 0.49s

pytest fastvideo/tests/train/methods/test_minimax_h3_finetune.py::\
test_h3_model_uses_modular_activation_checkpointing -q
1 passed, 14 warnings in 0.25s

PYTHONPATH="$PWD" bash .buildkite/scripts/unit_test.sh
1012 passed, 7 skipped, 19 warnings in 23.32s

pre-commit run --files <all changed paths>
Passed

A real CUDA probe on the local NVIDIA GB10 confirmed that PyTorch 2.12
dispatches the tested shape through
aten::_scaled_dot_product_flash_attention. The ops policy executed four
underlying fused forward kernels across four blocks versus eight under full,
with exact output, loss, input-gradient, and parameter-gradient parity.

Performance evidence scope

The earlier author measurements were single-sequence-parallel-rank GB200
experiments on an earlier/private H3 setup. They show the intended
speed-for-memory direction but are not presented here as a reproducible
current-tree SP=8 benchmark. This change's merge contract is functional
checkpoint behavior and numerical parity; production throughput should be
validated with each job's actual backend, sharding, and memory headroom.

SSIM is not required: this changes which activations are retained versus
recomputed during modular training, not inference output or model weights.

Checklist

• Rebased onto current main after PR [perf] Disable every offload path on unified memory, unblocking MiniMax H3 generation on one GB10 #1715
• Added behavioral, numerical, wrapper, failure, op-drift, model-wiring,
  and causal-cache tests
• Updated modular training documentation and example guidance
• Ran the exact Buildkite unit-lane script locally
• Ran repository pre-commit hooks on all applicable changed paths

[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=full-suite-passed

This rule is failing.

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

[shaoxiongduan]

/merge

[alexzms]

LGTM. Verified on GB200: ops now wraps per block and retains attention (6 flash_fwd launches across 6 blocks, vs 12 under full), with loss and input grad bit-identical to no checkpointing, and func.name() never raises on the ops a real step dispatches.

The red checks look stale — builds #4878/#4905 fall in a window where unit-tests were red repo-wide (#1726 and #1739 too), and ./fastvideo/tests/train/ passes on this branch locally (161 tests). Worth a CI re-run rather than a code change.

The branch was rebased and materially updated after this review (SM100a coverage, regression tests, and documentation). A fresh exact-head approval is required after CI.

[SolitaryThinker]

/test full

[SolitaryThinker]

/merge

[SolitaryThinker]

/test fastcheck

[SolitaryThinker]

/test full