[RL][qwen3_5] GDN: reset conv/recurrent state at packed sample boundaries

torchtitan/experiments/rl/tests/test_qwen3_5_gdn_packing.py

@@ -0,0 +1,97 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""GatedDeltaNet packed-sequence correctness (the cu_seqlens boundary fix).
+
+The RL batcher packs several training samples into one row (positions restart at 0
+per sample). Linear-attention (GatedDeltaNet) must reset its recurrent state AND
+causal conv at those boundaries; otherwise state/conv bleed across samples and the
+trainer computes wrong logprobs for every non-first packed sample (softmax layers
+avoid this via a block-diagonal attention mask). This test pins:
+
+1. cu_seqlens derivation: a single sample -> None (unchanged path, no regression);
+   a packed sequence -> the boundary offsets.
+2. Packed + positions matches processing each sample separately (the fix).
+3. Packed WITHOUT positions is contaminated at the non-first sample (the bug the
+   fix removes) -- guards against silently dropping the boundary handling.
+
+Test 1 (cu_seqlens derivation) runs on CPU; test 2 needs a GPU because the FLA
+gated-delta kernels are Triton/CUDA only.
+"""
+
+import pytest
+import torch
+
+from torchtitan.models.qwen3_5 import _qwen35_deltanet_config
+from torchtitan.models.qwen3_5.model import _cu_seqlens_from_positions
+
+
+def _max_abs_diff(x: torch.Tensor, y: torch.Tensor) -> float:
+    return (x - y).abs().max().item()
+
+
+def test_cu_seqlens_from_positions():
+    # Single sample (positions [0..L-1]) has one boundary -> None -> old path.
+    single = torch.arange(50).unsqueeze(0)
+    assert _cu_seqlens_from_positions(single) is None
+
+    # Packed [0..29, 0..39] -> boundaries at 0, 30 plus the total length 70.
+    packed = torch.cat([torch.arange(30), torch.arange(40)]).unsqueeze(0)
+    cu = _cu_seqlens_from_positions(packed)
+    assert cu.tolist() == [0, 30, 70]
+
+    # 3D MRoPE positions [B, L, 3]: boundary is the temporal-component reset.
+    mrope = packed.unsqueeze(-1).repeat(1, 1, 3)
+    cu_mrope = _cu_seqlens_from_positions(mrope)
+    assert cu_mrope.tolist() == [0, 30, 70]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="FLA GDN kernels need CUDA")
+def test_gated_delta_net_packing_matches_separate():
+    torch.manual_seed(0)
+    dev, dt, dim = "cuda", torch.bfloat16, 256
+    cfg = _qwen35_deltanet_config(
+        dim=dim,
+        n_key_heads=2,
+        n_value_heads=4,
+        key_head_dim=64,
+        value_head_dim=64,
+        layer_id=0,
+    )
+    gdn = cfg.build().to(dev)
+    with torch.no_grad():
+        for name, param in gdn.named_parameters():
+            if "A_log" in name or "dt_bias" in name:
+                param.zero_()
+            elif "norm" in name and "weight" in name:
+                param.fill_(1.0)
+            else:
+                param.normal_(0, 0.02)
+    gdn = gdn.to(dt)
+
+    len_a, len_b = 30, 40
+    x_a = torch.randn(1, len_a, dim, device=dev, dtype=dt)
+    x_b = torch.randn(1, len_b, dim, device=dev, dtype=dt)
+    pos_a = torch.arange(len_a, device=dev).unsqueeze(0)
+    pos_b = torch.arange(len_b, device=dev).unsqueeze(0)
+    x_packed = torch.cat([x_a, x_b], dim=1)
+    pos_packed = torch.cat([pos_a, pos_b], dim=1)  # restart at 0 for B
+
+    with torch.no_grad():
+        out_a = gdn(x_a, pos_a)
+        out_b = gdn(x_b, pos_b)
+        out_fix = gdn(x_packed, pos_packed)  # packed WITH positions (the fix)
+        out_bug = gdn(x_packed, None)  # packed WITHOUT positions (contaminated)
+
+    # The fix: the packed second sample matches its standalone output to bf16.
+    fix_b = _max_abs_diff(out_fix[:, len_a:], out_b)
+    bug_b = _max_abs_diff(out_bug[:, len_a:], out_b)
+    assert fix_b < 1e-3, f"packed+positions should match separate, got {fix_b}"
+    # Without positions the recurrent state / conv bleed -> clearly contaminated.
+    assert bug_b > 10 * fix_b, f"expected contamination without positions, got {bug_b}"
+
+    # The first packed sample is unaffected either way (nothing precedes it).
+    assert _max_abs_diff(out_bug[:, :len_a], out_a) == 0.0

torchtitan/models/qwen3_5/model.py

@@ -11,6 +11,7 @@
 import torch
 import torch.nn.functional as F
 
+from fla.modules.convolution import causal_conv1d as _fla_causal_conv1d
 from fla.ops.gated_delta_rule import (
     chunk_gated_delta_rule as _fla_chunk_gated_delta_rule,
     fused_recurrent_gated_delta_rule as _fla_fused_recurrent_gated_delta_rule,
@@ -35,12 +36,41 @@ def _l2norm(x: torch.Tensor, dim: int = -1, eps: float = 1e-6) -> torch.Tensor:
     return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
 
 
+def _cu_seqlens_from_positions(positions: torch.Tensor) -> torch.Tensor | None:
+    """Build FLA ``cu_seqlens`` from packed ``positions`` ([B, L], restart at 0 per sample).
+
+    The RL batcher packs several training samples into one row and restarts
+    ``positions`` at 0 at each sample (and each pad region). For linear-attention
+    (GatedDeltaNet) the recurrent state and causal conv must RESET at those
+    boundaries; softmax layers get a block-diagonal mask, but the FLA kernels need
+    ``cu_seqlens`` instead. We flatten [B, L] row-major to one [1, B*L] varlen
+    sequence, so every boundary (row start, within-row sample start, pad start) is a
+    ``positions == 0`` index. Mirrors open-instruct's ``_compute_packing_kwargs``.
+
+    Returns None when there is a single sample (no interior boundary), so the
+    non-packed / generator path is unchanged (bitwise identical).
+    """
+    # MRoPE passes 3D positions [B, L, 3]; the sample boundary is the reset of the
+    # temporal component. Text (RL) positions are 2D [B, L].
+    if positions.dim() == 3:
+        positions = positions[..., 0]
+    flat = positions.reshape(-1)
+    total = flat.numel()
+    starts = torch.nonzero(flat == 0, as_tuple=False).squeeze(-1).to(torch.int32)
+    # No interior boundary (single sample starting at 0) -> let callers skip varlen.
+    if starts.numel() <= 1:
+        return None
+    end = torch.tensor([total], device=positions.device, dtype=torch.int32)
+    return torch.cat([starts, end])
+
+
 def _torch_native_gated_delta(
     q: torch.Tensor,
     k: torch.Tensor,
     v: torch.Tensor,
     g: torch.Tensor,
     beta: torch.Tensor,
+    cu_seqlens: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Standalone math reference for the gated delta rule recurrence.
 
@@ -51,10 +81,36 @@ def _torch_native_gated_delta(
         v: (bs, seqlen, n_heads, value_head_dim)
         g: (bs, seqlen, n_heads) — log-space decay, always negative
         beta: (bs, seqlen, n_heads) — update gate ∈ (0, 1)
+        cu_seqlens: optional varlen boundaries over a flattened [1, total] input;
+            when set, the recurrent state resets at each boundary (packed samples).
 
     Returns:
         output: (bs, seqlen, n_heads, value_head_dim)
     """
+    # Packed: run the recurrence per segment with a fresh state, then restore shape.
+    if cu_seqlens is not None:
+        B, L = q.shape[0], q.shape[1]
+
+        def _flatten(t: torch.Tensor) -> torch.Tensor:
+            return t.reshape(1, B * L, *t.shape[2:])
+
+        qf, kf, vf, gf, bf = (
+            _flatten(q),
+            _flatten(k),
+            _flatten(v),
+            _flatten(g),
+            _flatten(beta),
+        )
+        bounds = cu_seqlens.tolist()
+        segs = [
+            _torch_native_gated_delta(
+                qf[:, s:e], kf[:, s:e], vf[:, s:e], gf[:, s:e], bf[:, s:e]
+            )
+            for s, e in zip(bounds[:-1], bounds[1:])
+        ]
+        out = torch.cat(segs, dim=1)
+        return out.reshape(B, L, *out.shape[2:])
+
     B, L, H, D_k = q.shape
     D_v = v.shape[-1]
     dtype = q.dtype
@@ -185,6 +241,7 @@ def forward(
         v: torch.Tensor,
         g: torch.Tensor,
         beta: torch.Tensor,
+        cu_seqlens: torch.Tensor | None = None,
     ) -> torch.Tensor:
         # Expand Q/K heads to match V when n_value_heads > n_key_heads
         if q.shape[2] != v.shape[2]:
@@ -194,7 +251,18 @@ def forward(
             k = k.repeat_interleave(repeat, dim=2)
 
         if self.backend == "torch_native":
-            return _torch_native_gated_delta(q, k, v, g, beta)
+            return _torch_native_gated_delta(q, k, v, g, beta, cu_seqlens=cu_seqlens)
+
+        # The FLA varlen path takes ONE [1, total] sequence with cu_seqlens marking
+        # the sample boundaries, so the recurrent state resets per packed sample.
+        # Flatten [B, L, ...] -> [1, B*L, ...]; cu_seqlens is None when not packed.
+        bs, seqlen = q.shape[0], q.shape[1]
+        if cu_seqlens is not None:
+            q = q.reshape(1, bs * seqlen, *q.shape[2:])
+            k = k.reshape(1, bs * seqlen, *k.shape[2:])
+            v = v.reshape(1, bs * seqlen, *v.shape[2:])
+            g = g.reshape(1, bs * seqlen, *g.shape[2:])
+            beta = beta.reshape(1, bs * seqlen, *beta.shape[2:])
 
         if self.backend == "fla_chunked":
             result = _fla_chunk_gated_delta_rule(
@@ -204,6 +272,7 @@ def forward(
                 g,
                 beta,
                 use_qk_l2norm_in_kernel=True,
+                cu_seqlens=cu_seqlens,
             )
         elif self.backend == "fla_fused_recurrent":
             result = _fla_fused_recurrent_gated_delta_rule(
@@ -213,6 +282,7 @@ def forward(
                 g,
                 beta=beta,
                 use_qk_l2norm_in_kernel=True,
+                cu_seqlens=cu_seqlens,
             )
         else:
             raise ValueError(
@@ -221,7 +291,10 @@ def forward(
             )
 
         # FLA kernels return (output, final_state); we only need output
-        return result[0]
+        out = result[0]
+        if cu_seqlens is not None:
+            out = out.reshape(bs, seqlen, *out.shape[2:])
+        return out
 
 
 class GatedDeltaNet(Module):
@@ -279,7 +352,37 @@ def __init__(self, config: Config):
         self.norm = config.norm.build()
         self.out_proj = config.out_proj.build()
 
-    def _causal_conv(self, x: torch.Tensor, conv: nn.Module) -> torch.Tensor:
+    def _causal_conv(
+        self, x: torch.Tensor, conv: nn.Module, cu_seqlens: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        # Packed samples: the causal conv window must not cross sample boundaries
+        # (else the first conv_kernel_size-1 tokens of each sample see the previous
+        # one). fla's causal_conv1d resets at cu_seqlens. Flatten [B, L, C] channels-
+        # last -> [1, B*L, C]; depthwise weight [C, 1, k] -> [C, k]. cu_seqlens is only
+        # built under packing (FSDP, non-TP activations), so x is a plain tensor here.
+        if cu_seqlens is not None:
+            assert not isinstance(
+                x, DTensor
+            ), "packed GDN conv (cu_seqlens) assumes non-TP (FSDP) activations"
+            bs, seqlen, channels = x.shape
+            weight = conv.weight
+            bias = conv.bias
+            if isinstance(weight, DTensor):
+                weight = weight.full_tensor()
+            if isinstance(bias, DTensor):
+                bias = bias.full_tensor()
+            # fla's causal_conv1d is untyped (pyrefly reads it as a Tensor).
+            y = _fla_causal_conv1d(
+                x.reshape(1, bs * seqlen, channels),
+                weight=weight.squeeze(1),  # pyrefly: ignore [not-callable]
+                bias=bias,
+                activation="silu",
+                cu_seqlens=cu_seqlens,
+            )
+            if isinstance(y, tuple):
+                y = y[0]
+            return y.reshape(bs, seqlen, channels)
+
         x = F.pad(x.transpose(1, 2), [self.conv_kernel_size - 1, 0])
         if isinstance(x, DTensor):
             # TODO: Remove once the DTensor Conv1d dispatch fix for sharded
@@ -315,16 +418,25 @@ def _conv(x_local: torch.Tensor, w_local: torch.Tensor) -> torch.Tensor:
             x = conv(x)
         return F.silu(x).transpose(1, 2)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, x: torch.Tensor, positions: torch.Tensor | None = None
+    ) -> torch.Tensor:
         bs, seqlen, _ = x.shape
 
+        # When several samples are packed into a row, `positions` restarts at 0 per
+        # sample; derive cu_seqlens so the conv AND recurrent state reset at each
+        # boundary. None (single sample / generator decode) -> unchanged path.
+        cu_seqlens = (
+            _cu_seqlens_from_positions(positions) if positions is not None else None
+        )
+
         # Shapes:
         #   xq, xk: (bs, seqlen, n_key_heads * key_head_dim)
         #   xv, xz: (bs, seqlen, n_value_heads * value_head_dim)
         #   xa, xb: (bs, seqlen, n_value_heads)
-        xq = self._causal_conv(self.in_proj_q(x), self.conv_q)
-        xk = self._causal_conv(self.in_proj_k(x), self.conv_k)
-        xv = self._causal_conv(self.in_proj_v(x), self.conv_v)
+        xq = self._causal_conv(self.in_proj_q(x), self.conv_q, cu_seqlens)
+        xk = self._causal_conv(self.in_proj_k(x), self.conv_k, cu_seqlens)
+        xv = self._causal_conv(self.in_proj_v(x), self.conv_v, cu_seqlens)
         xz = self.in_proj_z(x)
         xa = self.in_proj_a(x)
         xb = self.in_proj_b(x)
@@ -339,7 +451,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         g = -torch.exp(self.A_log.float()) * F.softplus(xa.float() + self.dt_bias)
         beta = torch.sigmoid(xb)
 
-        output = self.kernel(xq, xk, xv, g, beta)
+        output = self.kernel(xq, xk, xv, g, beta, cu_seqlens)
 
         xz = xz.view(bs, seqlen, -1, self.value_head_dim)
         output = self.norm(output, xz)
@@ -488,7 +600,9 @@ def forward(
         if self.full_attn:
             h = self.attn(h, attention_masks, positions)
         else:
-            h = self.attn(h)
+            # GatedDeltaNet needs positions to reset conv/recurrent state at packed
+            # sample boundaries (softmax uses attention_masks instead).
+            h = self.attn(h, positions)
         x = x + h
 
         h = self.ffn_norm(x)