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 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,31 @@
+max_steps = 1
+dashboard = false
+
+[model]
+name = "Qwen/Qwen3.5-9B"
+impl = "custom"
+optim_cpu_offload = false
+ac_offloading = "None"
+
+[model.debug]
+random_init = true
+
+[model.compile]
+fullgraph = true
+
+[data]
+type = "fake"
+batch_size = 8
+micro_batch_size = 1
+seq_len = 128
+length = "fixed"
+
+[deployment]
+type = "single_node"
+gpus_per_node = 8
+num_train_gpus = 8
+
+[slurm]
+job_name = "qwen3-5-9b-fullgraph"
+partition = "all"
+time = "00:45:00"

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

@@ -101,17 +101,18 @@ 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
+            kwargs["cu_seqlens_k"] = cu_seqlens_k
             out, _ = self.func(q, k, v, **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/qwen3_5/gated_delta_net.py

@@ -1,18 +1,19 @@
 import torch
 import torch.nn.functional as F
-from fla.modules import FusedRMSNormGated
-from fla.modules.conv import causal_conv1d
+from fla.modules.conv import causal_conv1d as fla_causal_conv1d
 from fla.ops.cp import build_cp_context
-from fla.ops.gated_delta_rule import chunk_gated_delta_rule
+from fla.ops.gated_delta_rule import chunk_gated_delta_rule as fla_chunk_gated_delta_rule
 from torch import nn
 
 from prime_rl.trainer.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig
+from prime_rl.trainer.models.qwen3_5.norm import Qwen3_5RMSNormGated
+from prime_rl.trainer.models.qwen3_5.ops import causal_conv1d, chunk_gated_delta_rule
 from prime_rl.utils.cp import CPContext
 
 # Dynamo lowers all-gather to concatenation, then fails to copy the result into
 # FLA's stacked output buffer. Keep CP convolution eager until this is fixed:
 # https://github.com/pytorch/pytorch/issues/155632
-causal_conv1d_with_context_parallelism = torch.compiler.disable(causal_conv1d)
+causal_conv1d_with_context_parallelism = torch.compiler.disable(fla_causal_conv1d)
 
 
 class Qwen3_5GatedDeltaNet(nn.Module):
@@ -38,7 +39,7 @@ def __init__(self, config: Qwen3_5TextConfig) -> None:
         )
         self.dt_bias = nn.Parameter(torch.ones(self.num_value_heads))
         self.A_log = nn.Parameter(torch.empty(self.num_value_heads).uniform_(0, 16).log_())
-        self.norm = FusedRMSNormGated(
+        self.norm = Qwen3_5RMSNormGated(
             self.value_head_dim,
             eps=config.rms_norm_eps,
             activation=config.output_gate_type,
@@ -71,15 +72,22 @@ def forward(
                 conv1d_kernel_size=self.conv_kernel_size,
             )
 
-        convolution = causal_conv1d_with_context_parallelism if context is not None else causal_conv1d
-        mixed_qkv, _ = convolution(
-            x=mixed_qkv,
-            weight=self.conv1d.weight.squeeze(1),
-            bias=self.conv1d.bias,
-            activation=self.activation,
-            cu_seqlens=cu_seqlens,
-            cp_context=context,
-        )
+        if context is None:
+            mixed_qkv = causal_conv1d(
+                mixed_qkv,
+                self.conv1d.weight.squeeze(1),
+                cu_seqlens,
+                self.activation,
+            )
+        else:
+            mixed_qkv, _ = causal_conv1d_with_context_parallelism(
+                x=mixed_qkv,
+                weight=self.conv1d.weight.squeeze(1),
+                bias=self.conv1d.bias,
+                activation=self.activation,
+                cu_seqlens=cu_seqlens,
+                cp_context=context,
+            )
 
         query, key, value = torch.split(mixed_qkv, [self.key_dim, self.key_dim, self.value_dim], dim=-1)
         query = query.reshape(batch_size, sequence_length, self.num_key_heads, self.key_head_dim)
@@ -91,16 +99,19 @@ def forward(
             query = query.repeat_interleave(heads_per_key, dim=2)
             key = key.repeat_interleave(heads_per_key, dim=2)
 
-        core_output, _ = chunk_gated_delta_rule(
-            q=query,
-            k=key,
-            v=value,
-            g=decay,
-            beta=beta,
-            use_qk_l2norm_in_kernel=True,
-            cu_seqlens=context.cu_seqlens if context is not None else cu_seqlens,
-            cp_context=context,
-        )
+        if context is None:
+            core_output = chunk_gated_delta_rule(query, key, value, decay, beta, cu_seqlens)
+        else:
+            core_output, _ = fla_chunk_gated_delta_rule(
+                q=query,
+                k=key,
+                v=value,
+                g=decay,
+                beta=beta,
+                use_qk_l2norm_in_kernel=True,
+                cu_seqlens=context.cu_seqlens,
+                cp_context=context,
+            )
 
         core_output = self.norm(
             core_output.reshape(-1, self.value_head_dim),

src/prime_rl/trainer/models/qwen3_5/norm.py

@@ -1,4 +1,5 @@
 import torch
+import torch.nn.functional as F
 from torch import nn
 
 
@@ -22,4 +23,25 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         return ((1.0 + self.weight.float()) * hidden_states).to(input_dtype)
 
 
-__all__ = ["Qwen3_5RMSNorm"]
+class Qwen3_5RMSNormGated(nn.Module):
+    def __init__(self, hidden_size: int, eps: float, activation: str) -> None:
+        super().__init__()
+        if activation not in ("silu", "swish", "sigmoid"):
+            raise ValueError(f"Unsupported activation: {activation}")
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+        self.activation = activation
+
+    def forward(self, hidden_states: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.float()
+        variance = hidden_states.square().mean(dim=-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        hidden_states = hidden_states * self.weight.float()
+
+        gate = gate.float()
+        gate = torch.sigmoid(gate) if self.activation == "sigmoid" else F.silu(gate)
+        return (hidden_states * gate).to(input_dtype)
+
+
+__all__ = ["Qwen3_5RMSNorm", "Qwen3_5RMSNormGated"]