refactor(trainer): make activation checkpointing operator-based

benchmarks/scripts/run_single_benchmark.py

@@ -70,11 +70,6 @@ class BenchmarkConfig(BaseConfig):
         Field(description="Activation checkpointing type"),
     ] = "Recompute"
 
-    selective_targets: Annotated[
-        list[str] | None,
-        Field(description="Selective activation checkpoint targets when ac=Selective"),
-    ] = None
-
     attention: Annotated[
         Literal["sdpa", "flash_attention_2", "flash_attention_3", "flash_attention_4"],
         Field(description="Attention implementation"),
@@ -160,8 +155,6 @@ def build_command(config: BenchmarkConfig, output_dir: Path) -> list[str]:
         cmd.append("--model.ac")
     elif config.ac == "Selective":
         cmd.extend(["--model.ac", "--model.ac.mode", "selective"])
-        if config.selective_targets:
-            cmd.extend(["--model.ac.targets", json.dumps(config.selective_targets)])
     elif config.ac == "Offload":
         cmd.append("--model.ac-offloading")
 

docs/advanced.md

@@ -35,7 +35,7 @@ impl = "custom"        # or "hf" to force the HF path
 | GLM-4 / GLM-4.5 / INTELLECT-3 | `THUDM/GLM-4-9B-0414`, `zai-org/GLM-4.5`, `PrimeIntellect/INTELLECT-3`, … | ✅ | ✅ |
 | GPT-OSS (HF MoE) | `openai/gpt-oss-20b`, `openai/gpt-oss-120b` | ❌ | ✅ |
 
-The custom path enables EP, CP, selective activation checkpointing, low-precision training, and grouped MoE kernels. Forcing `impl = "hf"` is mostly useful when debugging — it's slower and disables most MoE-specific knobs.
+Selective activation checkpointing works with either implementation. The custom path additionally enables EP, CP, low-precision training, and grouped MoE kernels. Forcing `impl = "hf"` is mostly useful when debugging and disables those model-specific runtime features.
 
 ### Low-precision training
 

docs/scaling.md

@@ -130,15 +130,18 @@ cp_style = "ulysses"         # "ring"
 | `trainer.model.ac.mode = "selective"` | medium | small | 
 | `trainer.model.ac_offloading` | extra | a bit more |
 
-AC and AC offloading are enabled by default (full mode). For the best memory/throughput tradeoff, switch to selective AC (custom impl only):
+AC and AC offloading are enabled by default (full mode). Both AC modes retain stateful MoE routing updates and DeepEP communication that cannot safely replay. Selective AC uses the same transformer-block boundaries and additionally retains distributed communication, expensive matrix multiplications, grouped GEMMs, and supported attention kernels:
 
 ```toml
 [trainer.model.ac]
 mode = "selective"
-targets = ["norm", "attn_proj"]  # see Reference for the full list per architecture
 ```
 
-`ac_offloading` is also on by default with `max_inflight_activations = 5`. We've observed this feature to be very effective, lowering the peak memory usage by 30-40% in some cases, while only lossing ~3-5% of throughput. To disable either, set `model.ac = "None"` or `model.ac_offloading = "None"`.
+Set `targets` to operator names (for example, `"aten::mm"`) or namespaces (for example, `"prime_rl_collectives"`) to replace the default selective targets. Correctness-required operations remain retained.
+
+Activation offloading still applies to tensors saved by autograd, but tensors retained by the checkpoint policy remain on the accelerator.
+
+`ac_offloading` is also on by default with `max_inflight_activations = 5`. We've observed this feature to be very effective, lowering the peak memory usage by 30-40% in some cases, while only losing ~3-5% of throughput. To disable either, set `model.ac = "None"` or `model.ac_offloading = "None"`.
 
 ### Optimizer Offloading
 

examples/vlm_sft_moe/sft.toml

@@ -23,11 +23,8 @@ rank = 32
 alpha = 32
 target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "experts"]
 
-# MoE + mode="full" recomputes the router, whose non-deterministic top-k routing
-# raises CheckpointError; selective AC excluding routed_experts avoids that.
 [model.ac]
 mode = "selective"
-targets = ["norm", "attn_proj", "linear_attn"]
 
 [data]
 type = "sft"

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -30,20 +30,13 @@ class GCConfig(BaseConfig):
 
 class ActivationCheckpointConfig(BaseConfig):
     mode: Literal["full", "selective"] = "full"
-    """``full`` checkpoints whole transformer blocks; ``selective`` checkpoints only the subcomponents listed in ``targets`` inside supported custom decoder layers."""
+    """Both modes checkpoint whole transformer blocks. ``selective`` additionally retains selected operations."""
 
     freq: int = Field(1, ge=1)
     """Apply activation checkpointing to every N layers."""
 
-    targets: list[str] = ["norm"]
-    """Selective checkpoint targets. ``norm`` checkpoints every norm module inside selected layers. ``attn_proj`` checkpoints projection-side attention work outside the kernel (input/output projections, attention-local norms, RoPE, gating, model-specific MLA projection helpers). ``mlp`` checkpoints the entire dense MLP forward (not for MoE). ``mla_up_proj`` checkpoints MLA Q/KV up-projection where supported. ``linear_attn`` checkpoints non-softmax token mixers (NemotronH Mamba, Qwen3.5-MoE GatedDeltaNet, AFMoE sliding-window attention)."""
-
-    @model_validator(mode="after")
-    def validate_selective_targets(self):
-        self.targets = list(dict.fromkeys(self.targets))
-        if self.mode == "selective" and not self.targets:
-            raise ValueError("Selective activation checkpointing requires at least one target.")
-        return self
+    targets: list[str] | None = None
+    """Operator names or namespaces retained in selective mode. ``None`` uses the default targets; an explicit list replaces them."""
 
 
 class ActivationOffloadingConfig(BaseConfig):
@@ -347,12 +340,6 @@ def ac_offloading_requires_ac(self):
             self.ac = ActivationCheckpointConfig()
         return self
 
-    @model_validator(mode="after")
-    def selective_ac_only_with_custom_impl(self):
-        if self.ac is not None and self.ac.mode == "selective" and self.impl not in ("custom", "auto"):
-            raise ValueError("Selective activation checkpointing requires model.impl='custom' or 'auto'")
-        return self
-
     @model_validator(mode="after")
     def cpu_offload_mutual_exclusion(self):
         if self.fsdp_cpu_offload and (self.optim_cpu_offload or self.full_offload):

pyproject.toml

@@ -139,6 +139,8 @@ environments = [
 # Override torch's pinned cuDNN to fix the Conv3d performance regression.
 # See: https://github.com/pytorch/pytorch/issues/166122
 override-dependencies = [
+    # Shared by vLLM, TileLang, and the pinned FlashAttention 4 runtime.
+    "apache-tvm-ffi==0.1.11",
     "nvidia-cudnn-cu13>=9.15 ; sys_platform == 'linux'",
     # vLLM 0.28.0 requires quack-kernels 0.6.4. It requires this exact CUTLASS version.
     # Keep this pin aligned with quack-kernels and the flash-attn-4 revision below.
@@ -250,7 +252,7 @@ torchao = [
     { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/torchao-0.18.0+git5f2baf9d5-cp310-abi3-linux_aarch64.whl", marker = "platform_machine == 'aarch64'" },
 ]
 dion = { git = "https://github.com/samsja/dion.git", rev = "d891eeb" }
-flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", subdirectory = "flash_attn/cute", rev = "82d6441" }
+flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", subdirectory = "flash_attn/cute", rev = "0251105" }
 vllm-router = [
     { url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.0/vllm_router-0.2.0-cp38-abi3-manylinux_2_28_x86_64.whl", marker = "platform_machine == 'x86_64'" },
     { url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.0/vllm_router-0.2.0-cp38-abi3-manylinux_2_28_aarch64.whl", marker = "platform_machine == 'aarch64'" },

src/prime_rl/trainer/activation_checkpointing.py

@@ -0,0 +1,117 @@
+"""Whole-block activation checkpointing with an operator-based policy."""
+
+from collections.abc import Callable
+from functools import partial
+
+import torch
+from torch import nn
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointImpl, checkpoint_wrapper
+from torch.utils.checkpoint import (
+    CheckpointPolicy,
+    SelectiveCheckpointContext,
+    create_selective_checkpoint_contexts,
+)
+
+from prime_rl.configs.trainer import ActivationCheckpointConfig
+
+# Adapted from TorchTitan's whole-block selective activation checkpointing policy.
+# These targets and the CUDA-to-CPU copy rule are correctness requirements. They
+# must remain active when custom targets replace the default selective targets.
+MANDATORY_SAVE_NAMESPACES = frozenset({"deepep"})
+MANDATORY_SAVE_OPERATIONS = frozenset(
+    {
+        "aten::topk",
+        "prime_rl::record_moe_routing_statistics",
+    }
+)
+
+DEFAULT_SELECTIVE_SAVE_NAMESPACES = frozenset(
+    {
+        "_c10d_functional",
+        "flash_attn",
+        "flash_attn_3",
+        "prime_rl_collectives",
+        "prime_rl_ring",
+    }
+)
+DEFAULT_SELECTIVE_SAVE_OPERATIONS = frozenset(
+    {
+        "aten::_efficient_attention_forward",
+        "aten::_flash_attention_forward",
+        "aten::_scaled_dot_product_attention_math",
+        "aten::_scaled_dot_product_cudnn_attention",
+        "aten::_scaled_dot_product_efficient_attention",
+        "aten::_scaled_dot_product_flash_attention",
+        "aten::_scaled_dot_product_flash_attention_for_cpu",
+        "aten::_scaled_dot_product_fused_attention_overrideable",
+        "aten::_scaled_grouped_mm",
+        "aten::_scaled_mm",
+        "aten::_scaled_mm_v2",
+        "aten::_grouped_mm",
+        "aten::addmm",
+        "aten::bmm",
+        "aten::convolution",
+        "aten::linear",
+        "aten::mm",
+        "prime_rl::fp8_blockwise_mm",
+        "prime_rl::grouped_fp8_gemm",
+        "prime_rl::sparse_mla",
+    }
+)
+DEFAULT_SELECTIVE_TARGETS = DEFAULT_SELECTIVE_SAVE_NAMESPACES | DEFAULT_SELECTIVE_SAVE_OPERATIONS
+
+
+def _mandatory_checkpoint_policy(
+    _context: SelectiveCheckpointContext,
+    operation: torch._ops.OpOverload | torch._ops.HigherOrderOperator,
+    *args,
+    **kwargs,
+) -> CheckpointPolicy:
+    if operation.namespace in MANDATORY_SAVE_NAMESPACES or operation.name() in MANDATORY_SAVE_OPERATIONS:
+        return CheckpointPolicy.MUST_SAVE
+
+    if operation.name() == "aten::_to_copy":
+        device = kwargs.get("device")
+        if isinstance(device, torch.device) and device.type == "cpu":
+            return CheckpointPolicy.MUST_SAVE
+
+    return CheckpointPolicy.PREFER_RECOMPUTE
+
+
+def _selective_checkpoint_policy(
+    context: SelectiveCheckpointContext,
+    operation: torch._ops.OpOverload | torch._ops.HigherOrderOperator,
+    *args,
+    targets: frozenset[str] = DEFAULT_SELECTIVE_TARGETS,
+    **kwargs,
+) -> CheckpointPolicy:
+    runtime_policy = _mandatory_checkpoint_policy(context, operation, *args, **kwargs)
+    if runtime_policy is CheckpointPolicy.MUST_SAVE:
+        return runtime_policy
+    if operation.namespace in targets or operation.name() in targets:
+        return CheckpointPolicy.MUST_SAVE
+    return CheckpointPolicy.PREFER_RECOMPUTE
+
+
+def get_activation_checkpoint_wrapper(config: ActivationCheckpointConfig) -> Callable[[nn.Module], nn.Module]:
+    if config.mode == "full":
+        return partial(checkpoint_wrapper, checkpoint_impl=CheckpointImpl.NO_REENTRANT)
+
+    targets = DEFAULT_SELECTIVE_TARGETS if config.targets is None else frozenset(config.targets)
+    policy = partial(_selective_checkpoint_policy, targets=targets)
+
+    return partial(
+        checkpoint_wrapper,
+        checkpoint_impl=CheckpointImpl.NO_REENTRANT,
+        context_fn=partial(create_selective_checkpoint_contexts, policy),
+    )
+
+
+__all__ = [
+    "DEFAULT_SELECTIVE_SAVE_NAMESPACES",
+    "DEFAULT_SELECTIVE_SAVE_OPERATIONS",
+    "DEFAULT_SELECTIVE_TARGETS",
+    "MANDATORY_SAVE_NAMESPACES",
+    "MANDATORY_SAVE_OPERATIONS",
+    "get_activation_checkpoint_wrapper",
+]