[bugfix]: fix modular ops activation checkpoint recomputation

docs/training/train_infra.md

@@ -51,7 +51,7 @@ torchrun --nproc_per_node=8 \
 ## Config Format
 
 Every run is defined by a single YAML file with five top-level sections.
-See `examples/train/example.yaml` for a fully-commented reference.
+See `examples/train/configs/example.yaml` for a fully-commented reference.
 
 ### `models` — Role-based model instances
 
@@ -80,9 +80,17 @@ Common model parameters:
 | `trainable` | `true` | Whether the model's parameters require gradients |
 | `disable_custom_init_weights` | `false` | Skip custom weight initialization (use for teacher/critic) |
 | `flow_shift` | `3.0` | Timestep shifting factor |
-| `enable_gradient_checkpointing_type` | `null` | Gradient checkpointing (`"full"` or `null`) |
+| `enable_gradient_checkpointing_type` | `null` | Gradient checkpointing (`"full"`, `"ops"`, or `null`) |
 | `attention_backend` | `null` | Optional role-local backend for Wan models (for example `ATTN_QAT_TRAIN`); overrides the process default only while this role's transformer is built |
 
+`full` recomputes every transformer-block operation and is the
+memory-conservative choice. `ops` retains outputs from supported,
+dispatcher-visible fused attention operations. It can reduce recompute time at
+the cost of higher activation memory, so use it only when the training shape has
+verified memory headroom. Math SDPA, VMoBA, FA3, and `ATTN_QAT_TRAIN` do not
+expose a retainable dispatcher boundary and therefore still use full attention
+recomputation under `ops`.
+
 Which roles are needed depends on the training method:
 
 | Method | Required roles |

examples/train/configs/example.yaml

@@ -147,7 +147,11 @@ training:
     mode_scale: 1.0                   # default: 1.0
     precondition_outputs: false       # default: false
     moba_config: {}                   # default: {}
-    enable_gradient_checkpointing_type: full  # default: null ("full" or null)
+    # default: null ("full", "ops", or null). "full" minimizes activation
+    # memory. "ops" retains supported dispatcher-visible attention outputs to
+    # reduce recompute time when memory headroom exists. Math SDPA, VMoBA, FA3,
+    # and ATTN_QAT_TRAIN still fully recompute attention under "ops".
+    enable_gradient_checkpointing_type: full
 
   # --- training top-level [TYPED] ---
   dit_precision: fp32                 # default: "fp32" (master weight precision)

fastvideo/tests/train/utils/test_activation_checkpoint.py

@@ -0,0 +1,232 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import copy
+import importlib
+
+import pytest
+import torch
+from torch import nn
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointWrapper
+
+import fastvideo.train.utils.activation_checkpoint as activation_checkpoint
+from fastvideo.train.utils.activation_checkpoint import apply_activation_checkpointing
+
+_TEST_OP_NAME = "fastvideo_activation_checkpoint_test::expensive_op"
+_TEST_OP_CALLS = 0
+_KNOWN_TRAINING_BLOCK_SPARSE_ATTENTION_OP_NAMES = {
+    "fastvideo_kernel::block_sparse_attn_sm90",
+    "fastvideo_kernel::block_sparse_attn_sm100a",
+    "fastvideo_kernel::block_sparse_attn_triton",
+}
+
+
+@torch.library.custom_op(_TEST_OP_NAME, mutates_args=())
+def _expensive_op(value: torch.Tensor) -> torch.Tensor:
+    global _TEST_OP_CALLS
+    _TEST_OP_CALLS += 1
+    return torch.sin(value)
+
+
+@_expensive_op.register_fake
+def _expensive_op_fake(value: torch.Tensor) -> torch.Tensor:
+    return torch.empty_like(value)
+
+
+def _setup_expensive_op_context(ctx, inputs, output) -> None:
+    del output
+    ctx.save_for_backward(inputs[0])
+
+
+def _backward_expensive_op(ctx, grad_output: torch.Tensor) -> torch.Tensor:
+    (value,) = ctx.saved_tensors
+    return grad_output * torch.cos(value)
+
+
+_expensive_op.register_autograd(_backward_expensive_op, setup_context=_setup_expensive_op_context)
+
+
+class _ToyBlock(nn.Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.projection = nn.Linear(4, 4)
+
+    def forward(self, value: torch.Tensor) -> torch.Tensor:
+        return value + _expensive_op(self.projection(value))
+
+
+class _ToyTransformer(nn.Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.blocks = nn.ModuleList([_ToyBlock() for _ in range(4)])
+
+    def forward(self, value: torch.Tensor) -> torch.Tensor:
+        for block in self.blocks:
+            value = block(value)
+        return value
+
+
+def _run_toy_model(
+    state_dict: dict[str, torch.Tensor],
+    checkpointing_type: str | None,
+) -> tuple[int, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]:
+    global _TEST_OP_CALLS
+    _TEST_OP_CALLS = 0
+
+    model = _ToyTransformer()
+    model.load_state_dict(copy.deepcopy(state_dict))
+    if checkpointing_type is not None:
+        apply_activation_checkpointing(model, checkpointing_type)
+
+    value = torch.linspace(-0.5, 0.5, 12, dtype=torch.float64).reshape(3, 4).requires_grad_()
+    model.to(dtype=torch.float64)
+    output = model(value)
+    loss = output.square().mean()
+    loss.backward()
+    parameter_grads = {
+        name.replace("._checkpoint_wrapped_module", ""): parameter.grad.detach().clone()
+        for name, parameter in model.named_parameters()
+    }
+    return _TEST_OP_CALLS, output.detach(), loss.detach(), value.grad.detach().clone(), parameter_grads
+
+
+def _toy_state_dict() -> dict[str, torch.Tensor]:
+    torch.manual_seed(17)
+    return _ToyTransformer().state_dict()
+
+
+@pytest.mark.parametrize(
+    ("checkpointing_type", "retain_test_op", "expected_calls"),
+    [
+        (None, True, 4),
+        ("full", True, 8),
+        ("ops", True, 4),
+        ("ops", False, 8),
+    ],
+)
+def test_checkpoint_policy_controls_expensive_op_recomputation(
+    monkeypatch: pytest.MonkeyPatch,
+    checkpointing_type: str | None,
+    retain_test_op: bool,
+    expected_calls: int,
+) -> None:
+    op_names = activation_checkpoint._SELECTIVE_ACTIVATION_CHECKPOINTING_OP_NAMES
+    if retain_test_op:
+        op_names = op_names | {_TEST_OP_NAME}
+    monkeypatch.setattr(activation_checkpoint, "_SELECTIVE_ACTIVATION_CHECKPOINTING_OP_NAMES", op_names)
+
+    calls, *_ = _run_toy_model(_toy_state_dict(), checkpointing_type)
+
+    assert calls == expected_calls
+
+
+def test_checkpoint_modes_preserve_outputs_and_gradients(monkeypatch: pytest.MonkeyPatch) -> None:
+    op_names = activation_checkpoint._SELECTIVE_ACTIVATION_CHECKPOINTING_OP_NAMES | {_TEST_OP_NAME}
+    monkeypatch.setattr(activation_checkpoint, "_SELECTIVE_ACTIVATION_CHECKPOINTING_OP_NAMES", op_names)
+    state_dict = _toy_state_dict()
+
+    baseline = _run_toy_model(state_dict, None)
+    for checkpointing_type in ("full", "ops"):
+        result = _run_toy_model(state_dict, checkpointing_type)
+        torch.testing.assert_close(result[1], baseline[1], rtol=0, atol=0)
+        torch.testing.assert_close(result[2], baseline[2], rtol=0, atol=0)
+        torch.testing.assert_close(result[3], baseline[3], rtol=0, atol=0)
+        assert result[4].keys() == baseline[4].keys()
+        for name, baseline_grad in baseline[4].items():
+            torch.testing.assert_close(result[4][name], baseline_grad, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("checkpointing_type", ["full", "ops"])
+def test_checkpointing_wraps_each_block_not_transformer_root(checkpointing_type: str) -> None:
+    model = _ToyTransformer()
+
+    result = apply_activation_checkpointing(model, checkpointing_type)
+
+    assert result is model
+    assert not isinstance(model, CheckpointWrapper)
+    assert all(isinstance(block, CheckpointWrapper) for block in model.blocks)
+
+
+@pytest.mark.parametrize("checkpointing_type", ["full", "ops"])
+def test_checkpointing_rejects_transformers_without_known_block_lists(checkpointing_type: str) -> None:
+    with pytest.raises(ValueError, match="Activation checkpointing is not applied successfully"):
+        apply_activation_checkpointing(nn.Linear(4, 4), checkpointing_type)
+
+
+def test_known_training_block_sparse_attention_ops_are_retained() -> None:
+    missing_op_names = (_KNOWN_TRAINING_BLOCK_SPARSE_ATTENTION_OP_NAMES
+                        - activation_checkpoint._SELECTIVE_ACTIVATION_CHECKPOINTING_OP_NAMES)
+
+    assert not missing_op_names, (
+        f"Known training block-sparse attention ops missing from checkpoint save policy: {missing_op_names}")
+
+
+def test_all_available_training_block_sparse_attention_ops_are_retained() -> None:
+    try:
+        block_sparse_attention = importlib.import_module("fastvideo_kernel.block_sparse_attn")
+    except ImportError as exc:
+        pytest.skip(f"fastvideo_kernel is unavailable on this platform: {exc}")
+    training_op_names = {
+        candidate._qualname
+        for candidate in vars(block_sparse_attention).values()
+        if getattr(candidate, "_qualname", "").startswith("fastvideo_kernel::block_sparse_attn_")
+        and getattr(candidate, "_backward_fn", None) is not None
+    }
+
+    assert training_op_names
+    missing_op_names = training_op_names - activation_checkpoint._SELECTIVE_ACTIVATION_CHECKPOINTING_OP_NAMES
+    assert not missing_op_names, (
+        f"Available training block-sparse attention ops missing from checkpoint save policy: {missing_op_names}")
+
+
+@pytest.mark.parametrize(
+    "module_name",
+    [
+        "fastvideo.train.models.kandinsky5.kandinsky5",
+        "fastvideo.train.models.ltx2.ltx2",
+        "fastvideo.train.models.minimax_h3.minimax_h3",
+        "fastvideo.train.models.wan.wan",
+    ],
+)
+def test_modular_model_plugins_use_modular_activation_checkpointing(module_name: str) -> None:
+    model_module = importlib.import_module(module_name)
+
+    assert model_module.apply_activation_checkpointing is apply_activation_checkpointing
+
+
+@pytest.mark.parametrize(
+    ("role_checkpointing_type", "fallback_checkpointing_type", "trainable", "expected_type", "expected_safe"),
+    [
+        ("ops", "full", True, "ops", True),
+        (None, "full", True, "full", True),
+        (None, None, True, None, False),
+        ("ops", None, False, "ops", False),
+    ],
+)
+def test_wan_causal_checkpoint_safe_cache_uses_resolved_role_config(
+    monkeypatch: pytest.MonkeyPatch,
+    role_checkpointing_type: str | None,
+    fallback_checkpointing_type: str | None,
+    trainable: bool,
+    expected_type: str | None,
+    expected_safe: bool,
+) -> None:
+    from fastvideo.train.models.wan.wan import WanModel
+    from fastvideo.train.models.wan.wan_causal import WanCausalModel
+    from fastvideo.train.utils.training_config import TrainingConfig
+
+    training_config = TrainingConfig()
+    training_config.model.enable_gradient_checkpointing_type = fallback_checkpointing_type
+    monkeypatch.setattr(WanModel, "_load_transformer", lambda self, **kwargs: nn.Module())
+
+    model = WanCausalModel(
+        init_from="unused-by-test",
+        training_config=training_config,
+        trainable=trainable,
+        enable_gradient_checkpointing_type=role_checkpointing_type,
+    )
+
+    assert model._resolved_gradient_checkpointing_type == expected_type
+    assert model._should_use_checkpoint_safe_kv_cache() is expected_safe

fastvideo/train/models/kandinsky5/kandinsky5.py

@@ -43,8 +43,6 @@
 from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
     FlowMatchEulerDiscreteScheduler, )
 from fastvideo.pipelines import TrainingBatch
-from fastvideo.training.activation_checkpoint import (
-    apply_activation_checkpointing, )
 from fastvideo.training.training_utils import (
     compute_density_for_timestep_sampling,
     get_sigmas,
@@ -53,6 +51,8 @@
 )
 
 from fastvideo.train.models.base import ModelBase
+from fastvideo.train.utils.activation_checkpoint import (
+    apply_activation_checkpointing, )
 from fastvideo.train.utils.module_state import (
     apply_trainable, )
 from fastvideo.train.utils.moduleloader import (

fastvideo/train/models/ltx2/ltx2.py

@@ -39,10 +39,10 @@
 from fastvideo.models.dits.ltx2 import VideoLatentShape
 from fastvideo.pipelines import ForwardBatch, TrainingBatch
 from fastvideo.platforms import AttentionBackendEnum
-from fastvideo.training.activation_checkpoint import (
-    apply_activation_checkpointing, )
 
 from fastvideo.train.models.wan.wan import WanModel
+from fastvideo.train.utils.activation_checkpoint import (
+    apply_activation_checkpointing, )
 from fastvideo.train.utils.module_state import (
     apply_trainable, )
 from fastvideo.train.utils.moduleloader import (

fastvideo/train/models/wan/wan.py

@@ -19,8 +19,6 @@
     FlowMatchEulerDiscreteScheduler, )
 from fastvideo.pipelines import TrainingBatch
 from fastvideo.platforms import AttentionBackendEnum
-from fastvideo.training.activation_checkpoint import (
-    apply_activation_checkpointing, )
 from fastvideo.training.training_utils import (
     compute_density_for_timestep_sampling,
     get_sigmas,
@@ -33,6 +31,8 @@
 )
 
 from fastvideo.train.models.base import ModelBase
+from fastvideo.train.utils.activation_checkpoint import (
+    apply_activation_checkpointing, )
 from fastvideo.train.utils.module_state import (
     apply_trainable, )
 from fastvideo.train.utils.moduleloader import (
@@ -80,12 +80,17 @@ def __init__(
             attention_backend=attention_backend,
         )
         self._init_from = str(init_from)
+        self._resolved_gradient_checkpointing_type = (enable_gradient_checkpointing_type or getattr(
+            getattr(training_config, "model", None),
+            "enable_gradient_checkpointing_type",
+            None,
+        ))
 
         self.transformer = self._load_transformer(
             init_from=self._init_from,
             trainable=self._trainable,
             disable_custom_init_weights=(disable_custom_init_weights),
-            enable_gradient_checkpointing_type=(enable_gradient_checkpointing_type),
+            enable_gradient_checkpointing_type=(self._resolved_gradient_checkpointing_type),
             training_config=training_config,
             transformer_override_safetensor=(transformer_override_safetensor),
             attention_backend=self.attention_backend,

fastvideo/train/models/wan/wan_causal.py

@@ -426,9 +426,7 @@ def _initialize_kv_cache(
         return kv_cache
 
     def _should_use_checkpoint_safe_kv_cache(self, ) -> bool:
-        tc = getattr(self, "training_config", None)
-        checkpointing_type = tc.model.enable_gradient_checkpointing_type if tc is not None else None
-        return (bool(checkpointing_type) and bool(self._trainable))
+        return (bool(self._resolved_gradient_checkpointing_type) and bool(self._trainable))
 
     def _should_snapshot_streaming_cache(self, ) -> bool:
         return (self._should_use_checkpoint_safe_kv_cache())