feat: support fullgraph Qwen3.5 training

configs/README.md

@@ -12,7 +12,8 @@ Configurations for running prime-rl.
   [`examples/advanced/`](../examples/advanced).
 - **`ci/`** — integration and nightly configs used by CI.
 - **`debug/`** — throwaway configs for developing the framework itself: `algo/`
-  (per-algorithm smokes), `fake/` (fake-data trainer/SFT smokes), `multi-env/`
+  (per-algorithm smokes), `fake/` (fake-data trainer/SFT smokes, including a one-node
+  Qwen3.5-9B full-graph compile and partitioned CUDA Graph check), `multi-env/`
   (two reverse-text train sources + one eval source, one env server per source), and
   `eval/` (`uv run eval` smokes against Prime Inference, one per shape: `single-turn`
   (gsm8k), `multi-turn` (terminal-bench-2 fix-git in sandboxes against a local vLLM deployment with the adaptive band), `resume`, `multi-env`).

configs/debug/fake/qwen3_5_9b_fullgraph_sft.toml

@@ -0,0 +1,41 @@
+max_steps = 5
+dashboard = false
+
+[model]
+name = "Qwen/Qwen3.5-9B"
+impl = "custom"
+seq_len = 8192
+optim_cpu_offload = false
+fsdp_cpu_offload = false
+reshard_after_forward = false
+ac_offloading = "None"
+
+[model.debug]
+random_init = true
+
+[model.compile]
+fullgraph = true
+mode = "reduce-overhead"
+cudagraph_copy_static_inputs = true
+
+[data]
+type = "fake"
+batch_size = 8
+micro_batch_size = 1
+seq_len = 8192
+length = "fixed"
+input_ids = "random"
+
+[deployment]
+type = "single_node"
+gpus_per_node = 8
+num_train_gpus = 8
+
+[env_vars]
+TORCH_LOGS = "perf_hints"
+FLA_DISABLE_TENSOR_CACHE = "1"
+
+[slurm]
+job_name = "qwen3-5-9b-fullgraph"
+partition = "all"
+time = "00:45:00"

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

@@ -85,6 +85,26 @@ class CompileConfig(BaseConfig):
     fullgraph: bool = False
     """Compile transformer blocks with ``fullgraph=True``."""
 
+    mode: Literal["default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"] | None = None
+    """Optimization mode passed to ``torch.compile``."""
+
+    cudagraph_partition_ops: list[str] = []
+    """Custom operators that run outside CUDA Graph capture while the surrounding compiled partitions remain captured."""
+
+    cudagraph_copy_static_inputs: bool = False
+    """Copy lifted parameters into stable CUDA Graph input buffers instead of requiring their addresses to remain fixed."""
+
+    @model_validator(mode="after")
+    def cudagraph_options_require_cudagraphs(self):
+        cudagraph_modes = {"reduce-overhead", "max-autotune"}
+        if self.cudagraph_partition_ops and self.mode not in cudagraph_modes:
+            raise ValueError("cudagraph_partition_ops requires mode='reduce-overhead' or mode='max-autotune'")
+        if self.cudagraph_copy_static_inputs and self.mode not in cudagraph_modes:
+            raise ValueError("cudagraph_copy_static_inputs requires mode='reduce-overhead' or mode='max-autotune'")
+        if self.cudagraph_copy_static_inputs and self.cudagraph_partition_ops:
+            raise ValueError("cudagraph_copy_static_inputs cannot be combined with cudagraph_partition_ops")
+        return self
+
 
 class FusionsConfig(BaseConfig):
     enabled: list[Literal["gate_up", "qkv"]] = ["gate_up", "qkv"]

src/prime_rl/trainer/model.py

@@ -14,6 +14,7 @@
 from huggingface_hub import snapshot_download
 from jaxtyping import Int
 from torch import Tensor
+from torch._inductor.cudagraph_utils import CUDAGraphPolicy
 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
@@ -828,13 +829,40 @@ def apply_ac(model: nn.Module, ac_config: ActivationCheckpointConfig):
     )
 
 
+class _CopyStaticInputsCUDAGraphPolicy(CUDAGraphPolicy):
+    def cudagraphify(self, model, example_inputs, static_input_idxs, **kwargs):
+        return super().cudagraphify(model, example_inputs, (), **kwargs)
+
+
 def apply_compile(model: nn.Module, compile_config: CompileConfig):
     torch._dynamo.config.capture_scalar_outputs = True
+    if compile_config.cudagraph_copy_static_inputs:
+        torch._inductor.config.cudagraph_policy = _CopyStaticInputsCUDAGraphPolicy()
     language_model = get_language_model(model)
+    mode = compile_config.mode
+    options = None
+    if compile_config.cudagraph_partition_ops:
+        mode_options = torch._inductor.list_mode_options().get(mode or "default", {})
+        options = {**mode_options, "custom_should_partition_ops": compile_config.cudagraph_partition_ops}
+        mode = None
     for layer_id in range(len(language_model.layers)):
         # Doing it in-place avoids mangled fqn which can break checkpoint loading
-        language_model.layers[layer_id].compile(fullgraph=compile_config.fullgraph)
-    get_logger().info(f"Compiled {len(language_model.layers)} layers (fullgraph={compile_config.fullgraph})")
+        language_model.layers[layer_id].compile(
+            fullgraph=compile_config.fullgraph,
+            mode=mode,
+            options=options,
+        )
+    get_logger().info(
+        f"Compiled {len(language_model.layers)} layers "
+        f"(fullgraph={compile_config.fullgraph}, mode={compile_config.mode}, "
+        f"cudagraph_partition_ops={compile_config.cudagraph_partition_ops}, "
+        f"cudagraph_copy_static_inputs={compile_config.cudagraph_copy_static_inputs})"
+    )
+
+
+def mark_cudagraph_step_begin(compile_config: CompileConfig | None) -> None:
+    if compile_config is not None and compile_config.mode in {"reduce-overhead", "max-autotune"}:
+        torch.compiler.cudagraph_mark_step_begin()
 
 
 def apply_quantization(model: nn.Module, config: ModelConfig) -> None:

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

@@ -24,8 +24,11 @@
 
 try:
     from flash_attn.cute import flash_attn_varlen_func as flash_attn_4_varlen_func
+
+    from .flash_attn_4_ops import flash_attn_4_varlen
 except ImportError:
     flash_attn_4_varlen_func = None  # type: ignore
+    flash_attn_4_varlen = None  # type: ignore
 
 
 @dataclass
@@ -101,17 +104,15 @@ def project_qkv(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.
     def _compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens, max_seqlen):
         """Run the flash attention kernel. q/k/v are [total_tokens, heads, dim]."""
         kwargs: dict = {"causal": True}
+        cu_seqlens_k = cu_seqlens.clone()
         sliding_window = getattr(self, "sliding_window", None)
         if sliding_window is not None:
             kwargs["window_size"] = (sliding_window - 1, 0)
         if self._flash_attn_version == 4:
-            # FA4's flash_attn_varlen_func has qv as the 4th positional arg,
-            # so cu_seqlens must be passed as keyword args to avoid misalignment.
-            kwargs["cu_seqlens_q"] = cu_seqlens
-            kwargs["cu_seqlens_k"] = cu_seqlens
-            out, _ = self.func(q, k, v, **kwargs)
+            # Custom-op wrapper: FA4's public wrapper is not traceable under fullgraph compile.
+            out = flash_attn_4_varlen(q, k, v, cu_seqlens, cu_seqlens_k, **kwargs)
         else:
-            out = self.func(q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, **kwargs)
+            out = self.func(q, k, v, cu_seqlens, cu_seqlens_k, max_seqlen, max_seqlen, **kwargs)
         return out
 
     def forward(

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

@@ -0,0 +1,143 @@
+import torch
+from flash_attn.cute.interface import _flash_attn_bwd, _flash_attn_fwd
+
+# FA4's public wrapper calls torch._guards.active_fake_mode(), which Dynamo refuses to trace, so
+# fullgraph compile breaks inside it. Wrapping the kernels in custom ops keeps them opaque to Dynamo.
+
+
+@torch.library.custom_op("prime_rl::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 _flash_attn_4_varlen_fake(
+    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]:
+    total_q, num_heads, _ = q.shape
+    out = q.new_empty((total_q, num_heads, v.shape[-1]))
+    lse = q.new_empty((num_heads, total_q), dtype=torch.float32)
+    return out, lse
+
+
+@torch.library.custom_op("prime_rl::flash_attn_4_varlen_backward", mutates_args=())
+def _flash_attn_4_varlen_backward(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    grad_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]:
+    grad_q, grad_k, grad_v = _flash_attn_bwd(
+        q,
+        k,
+        v,
+        out,
+        grad_out,
+        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 grad_q, grad_k, grad_v
+
+
+@_flash_attn_4_varlen_backward.register_fake
+def _flash_attn_4_varlen_backward_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    grad_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]:
+    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_left = window_size_left
+    ctx.window_size_right = window_size_right
+
+
+def _flash_attn_4_varlen_autograd_backward(ctx, grad_out: torch.Tensor, grad_lse: torch.Tensor):
+    q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensors
+    grad_q, grad_k, grad_v = _flash_attn_4_varlen_backward(
+        q,
+        k,
+        v,
+        out,
+        grad_out,
+        lse,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        ctx.causal,
+        ctx.window_size_left,
+        ctx.window_size_right,
+    )
+    return grad_q, grad_k, grad_v, None, None, None, None, None
+
+
+_flash_attn_4_varlen.register_autograd(
+    _flash_attn_4_varlen_autograd_backward,
+    setup_context=_flash_attn_4_varlen_setup_context,
+)
+
+
+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 = True,
+    window_size: tuple[int | None, int | None] = (None, None),
+) -> torch.Tensor:
+    out, _ = _flash_attn_4_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, causal, window_size[0], window_size[1])
+    return out
+
+
+__all__ = ["flash_attn_4_varlen"]