perf(trainer): compile MoE blocks with fullgraph by keeping EP experts in the block's FSDP unit

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

@@ -83,7 +83,7 @@ def _normalize_optimizer_in_backward_offload(value: Any) -> Any:
 
 class CompileConfig(BaseConfig):
     fullgraph: bool = False
-    """Compile transformer blocks with ``fullgraph=True``."""
+    """Compile transformer blocks with ``fullgraph=True``. Custom MoE models need ``moe_router_dtype="bfloat16"``: the fp32 router is its own FSDP unit inside the block, and dynamo cannot trace FSDP hooks."""
 
     mode: Literal["reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs", "lite"] | None = None
     """``torch.compile`` mode. ``reduce-overhead`` records CUDA graphs to cut kernel launch overhead; ``max-autotune`` modes trade longer compile times for tuned kernels (``max-autotune`` also records CUDA graphs, ``max-autotune-no-cudagraphs`` does not). CUDA-graphed layers re-record on new input shapes. ``None`` uses PyTorch's default mode."""

src/prime_rl/trainer/activation_checkpointing.py

@@ -42,6 +42,7 @@
         "_c10d_functional",
         "flash_attn",
         "flash_attn_3",
+        "prime_rl_attn",
         "prime_rl_collectives",
         "prime_rl_ring",
     }

src/prime_rl/trainer/model.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import time
+from collections.abc import Callable
 from pathlib import Path
 from typing import cast
 
@@ -17,8 +18,9 @@
 from torch import Tensor
 from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageReader
 from torch.distributed.checkpoint.state_dict_loader import load as dcp_load
-from torch.distributed.device_mesh import DeviceMesh
 from torch.distributed.fsdp import CPUOffloadPolicy, FSDPModule, MixedPrecisionPolicy, OffloadPolicy, fully_shard
+from torch.distributed.fsdp._fully_shard._fsdp_common import FSDPMeshInfo, ShardPlacementResult
+from torch.distributed.fsdp._fully_shard._fsdp_init import _get_mesh_info
 from torch.distributed.tensor import Shard
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig, PretrainedConfig
 from transformers.tokenization_utils import PreTrainedTokenizer
@@ -501,6 +503,23 @@ def setup_processor(config: ModelConfig):
     return processor
 
 
+def _expert_shard_placement_fn(
+    experts: nn.Module,
+    expert_mesh_info: FSDPMeshInfo,
+    shard_placement_fn: Callable[[nn.Parameter], Shard | None],
+) -> Callable[[nn.Parameter], ShardPlacementResult | Shard | None]:
+    """Shards the expert parameters over the EP-complement mesh and everything else over the block's mesh."""
+    expert_params = set(experts.parameters())
+
+    def placement_fn(parameter: nn.Parameter) -> ShardPlacementResult | Shard | None:
+        placement = shard_placement_fn(parameter)
+        if parameter in expert_params:
+            return ShardPlacementResult(placement=placement, mesh_info=expert_mesh_info)
+        return placement
+
+    return placement_fn
+
+
 def setup_fsdp(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDims):
     mp_policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=DTYPE_MAP[config.reduce_dtype])
     offload_policy: OffloadPolicy = CPUOffloadPolicy(pin_memory=True) if config.fsdp_cpu_offload else OffloadPolicy()
@@ -522,14 +541,16 @@ def shard_placement_fn(parameter: nn.Parameter) -> Shard | None:
         "shard_placement_fn": shard_placement_fn,
     }
 
-    dp_mod_ep_mesh: DeviceMesh | None = None
+    expert_mesh_info: FSDPMeshInfo | None = None
     if parallel_dims.ep_enabled:
         dp_mod_ep_mesh_dim_names = []
         if parallel_dims.dp_replicate_enabled:
             dp_mod_ep_mesh_dim_names.append("dp_replicate")
         dp_mod_ep_mesh_dim_names.append("dp_shard_mod_ep")
 
         dp_mod_ep_mesh = parallel_dims.world_mesh[tuple(dp_mod_ep_mesh_dim_names)]
+        expert_mesh_info = _get_mesh_info(dp_mod_ep_mesh)
+        assert isinstance(expert_mesh_info, FSDPMeshInfo)
 
     is_vlm_training = config.vlm is not None
     if is_vlm_training:
@@ -543,14 +564,26 @@ def shard_placement_fn(parameter: nn.Parameter) -> Shard | None:
     language_model = get_language_model(model, override=config.vlm.language_model_attr if is_vlm_training else None)
     transformer_layers = language_model.layers
 
+    fullgraph = config.compile is not None and config.compile.fullgraph
     for transformer_block in transformer_layers:
         block_mlp = getattr(transformer_block, "mlp", None)
-        if parallel_dims.ep_enabled and block_mlp is not None and isinstance(block_mlp, MoE):
-            fully_shard(block_mlp.experts, mesh=dp_mod_ep_mesh, **fsdp_config)
-
-            block_mlp.experts.set_gradient_divide_factor(parallel_dims.fsdp_gradient_divide_factor)
+        block_fsdp_config = fsdp_config
+        if expert_mesh_info is not None and isinstance(block_mlp, MoE):
+            # The experts shard over the EP-complement mesh but stay in the block's FSDP unit: a nested
+            # unit would put dynamo-disabled FSDP hooks inside the compiled block and break fullgraph.
+            block_fsdp_config = {
+                **fsdp_config,
+                "shard_placement_fn": _expert_shard_placement_fn(
+                    block_mlp.experts, expert_mesh_info, shard_placement_fn
+                ),
+            }
 
         if config.moe_router_dtype == "float32" and isinstance(block_mlp, MoE):
+            if fullgraph:
+                raise ValueError(
+                    "model.compile.fullgraph=true requires model.moe_router_dtype='bfloat16': the fp32 router is "
+                    "its own FSDP unit inside the compiled block, and dynamo cannot trace FSDP hooks."
+                )
             # Own FSDP unit with an fp32 policy so the gate weight is not cast to
             # bf16 for forward and its gradients reduce in fp32.
             fully_shard(
@@ -567,8 +600,12 @@ def shard_placement_fn(parameter: nn.Parameter) -> Shard | None:
         fully_shard(
             transformer_block,
             mesh=hsdp_mesh,
-            **fsdp_config,
+            **block_fsdp_config,
         )
+        if expert_mesh_info is not None and isinstance(block_mlp, MoE):
+            # Expert gradients reduce over the EP-complement mesh only, so they divide by the full
+            # data-parallel size. That is already the dense parameters' default divisor.
+            transformer_block.set_gradient_divide_factor(parallel_dims.fsdp_gradient_divide_factor)
 
     shard_norm_and_lm_head = hasattr(model, "config") and not model.config.tie_word_embeddings
 
@@ -618,14 +655,10 @@ def shard_placement_fn(parameter: nn.Parameter) -> Shard | None:
     for transformer_block, next_transformer_block in zip(transformer_blocks, next_transformer_blocks):
         if next_transformer_block is not None:
             next_mlp = getattr(next_transformer_block, "mlp", None)
-            if next_mlp is not None and isinstance(next_mlp, MoE):
-                prefetch_modules = [next_transformer_block]
-                if isinstance(next_mlp.router, FSDPModule):
-                    prefetch_modules.append(next_mlp.router)
-                prefetch_modules.append(next_mlp.experts)
-                transformer_block.set_modules_to_forward_prefetch(prefetch_modules)
-            else:
-                transformer_block.set_modules_to_forward_prefetch([next_transformer_block])
+            prefetch_modules = [next_transformer_block]
+            if isinstance(next_mlp, MoE) and isinstance(next_mlp.router, FSDPModule):
+                prefetch_modules.append(next_mlp.router)
+            transformer_block.set_modules_to_forward_prefetch(prefetch_modules)
         elif language_model.norm is not None and model.lm_head is not None:
             if shard_norm_and_lm_head:
                 transformer_block.set_modules_to_forward_prefetch([language_model.norm, model.lm_head])
@@ -638,10 +671,8 @@ def shard_placement_fn(parameter: nn.Parameter) -> Shard | None:
         last_transformer_block = reversed_transformer_blocks[0]
         prefetch_modules = [last_transformer_block]
         last_mlp = getattr(last_transformer_block, "mlp", None)
-        if last_mlp is not None and isinstance(last_mlp, MoE):
-            prefetch_modules.append(last_mlp.experts)
-            if isinstance(last_mlp.router, FSDPModule):
-                prefetch_modules.append(last_mlp.router)
+        if isinstance(last_mlp, MoE) and isinstance(last_mlp.router, FSDPModule):
+            prefetch_modules.append(last_mlp.router)
 
         if shard_norm_and_lm_head:
             model.lm_head.set_modules_to_backward_prefetch(prefetch_modules)
@@ -651,13 +682,10 @@ def shard_placement_fn(parameter: nn.Parameter) -> Shard | None:
     for transformer_block, prev_transformer_block in zip(reversed_transformer_blocks, prev_transformer_blocks):
         if prev_transformer_block is not None:
             prev_mlp = getattr(prev_transformer_block, "mlp", None)
-            if prev_mlp is not None and isinstance(prev_mlp, MoE):
-                prefetch_modules = [prev_transformer_block, prev_mlp.experts]
-                if isinstance(prev_mlp.router, FSDPModule):
-                    prefetch_modules.append(prev_mlp.router)
-                transformer_block.set_modules_to_backward_prefetch(prefetch_modules)
-            else:
-                transformer_block.set_modules_to_backward_prefetch([prev_transformer_block])
+            prefetch_modules = [prev_transformer_block]
+            if isinstance(prev_mlp, MoE) and isinstance(prev_mlp.router, FSDPModule):
+                prefetch_modules.append(prev_mlp.router)
+            transformer_block.set_modules_to_backward_prefetch(prefetch_modules)
         elif embed_module is not None:
             if shard_norm_and_lm_head:
                 transformer_block.set_modules_to_backward_prefetch([embed_module])

src/prime_rl/trainer/models/layers/attn.py

@@ -24,10 +24,114 @@
 
 try:
     from flash_attn.cute import flash_attn_varlen_func as flash_attn_4_varlen_func
+    from flash_attn.cute.interface import _flash_attn_bwd, _flash_attn_fwd
 except ImportError:
     flash_attn_4_varlen_func = None  # type: ignore
 
 
+# FA4's flash_attn_varlen_func is a Python autograd.Function that Dynamo cannot trace. Inside a
+# checkpointed block that graph break drops the whole block to eager, so FA4 is exposed as an
+# opaque custom op instead, which also lets selective activation checkpointing save its output.
+@torch.library.custom_op("prime_rl_attn::flash_attn_4_varlen", mutates_args=())
+def _flash_attn_4_varlen(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    causal: bool,
+    window_size_left: int | None,
+    window_size_right: int | None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    out, lse, _, _ = _flash_attn_fwd(
+        q,
+        k,
+        v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        causal=causal,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        return_lse=True,
+    )
+    return out, lse
+
+
+@_flash_attn_4_varlen.register_fake
+def _(q, k, v, cu_seqlens_q, cu_seqlens_k, causal, window_size_left, window_size_right):
+    out = q.new_empty((*q.shape[:-1], v.shape[-1]))
+    lse = q.new_empty((q.shape[1], q.shape[0]), dtype=torch.float32)
+    return out, lse
+
+
+@torch.library.custom_op("prime_rl_attn::flash_attn_4_varlen_backward", mutates_args=())
+def _flash_attn_4_varlen_backward(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    lse: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    causal: bool,
+    window_size_left: int | None,
+    window_size_right: int | None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    dq, dk, dv = _flash_attn_bwd(
+        q,
+        k,
+        v,
+        out,
+        dout,
+        lse,
+        causal=causal,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+    )
+    return dq, dk, dv
+
+
+@_flash_attn_4_varlen_backward.register_fake
+def _(dout, q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k, causal, window_size_left, window_size_right):
+    return torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
+
+
+def _flash_attn_4_varlen_setup_context(ctx, inputs, output) -> None:
+    q, k, v, cu_seqlens_q, cu_seqlens_k, causal, window_size_left, window_size_right = inputs
+    out, lse = output
+    ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k)
+    ctx.causal = causal
+    ctx.window_size = (window_size_left, window_size_right)
+
+
+def _flash_attn_4_varlen_autograd(ctx, dout: torch.Tensor, _dlse: torch.Tensor | None):
+    q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensors
+    dq, dk, dv = _flash_attn_4_varlen_backward(
+        dout.contiguous(), q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k, ctx.causal, *ctx.window_size
+    )
+    return dq, dk, dv, None, None, None, None, None
+
+
+_flash_attn_4_varlen.register_autograd(_flash_attn_4_varlen_autograd, setup_context=_flash_attn_4_varlen_setup_context)
+
+
+def flash_attn_4_varlen_op(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    causal: bool = False,
+    window_size: tuple[int | None, int | None] = (None, None),
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Drop-in for FA4's ``flash_attn_varlen_func`` as called by ``FlashAttention``, routed through the custom op."""
+    return _flash_attn_4_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, causal, window_size[0], window_size[1])
+
+
 @dataclass
 class AttentionConfig:
     hidden_size: int
@@ -55,7 +159,7 @@ class FlashAttention(nn.Module):
     _funcs = {
         2: flash_attn_varlen_func,
         3: flash_attn_3_varlen_func,
-        4: flash_attn_4_varlen_func,
+        4: flash_attn_4_varlen_op if flash_attn_4_varlen_func is not None else None,
     }
 
     def __init__(self, config: AttentionConfig, flash_attn_version: int = 2):