From b690abb9b97bff8b519756ade0826917edc271e1 Mon Sep 17 00:00:00 2001 From: yichuan-w Date: Tue, 7 Jul 2026 20:27:27 -0700 Subject: [PATCH 1/4] [qwen3_5] GDN: reset conv/recurrent state at packed sample boundaries The RL trainer next-fits several training samples into one row and restarts `positions` at 0 at each sample. Full-attention layers isolate the packed samples with a block-diagonal mask, but GatedDeltaNet was called as `self.attn(h)` with no boundary information, so its causal conv1d window and recurrent state bled across sample boundaries. Every non-first packed sample then got wrong logprobs from the trainer while the (unpacked) vLLM generator did not, producing garbage importance ratios and flat RL reward. At long sequence lengths with ~3 samples per row this corrupts the majority of trained samples. Fix: derive FLA `cu_seqlens` from the packed `positions` (a `positions == 0` index marks a sample start, matching the existing document-mask convention in `attention.py`) and thread it through `Qwen35TransformerBlock` -> `GatedDeltaNet` -> `GatedDeltaKernel`, passing it to the FLA chunk/recurrent kernels and to `fla.causal_conv1d` so both reset per segment. A single, unpacked sequence yields `cu_seqlens=None`, leaving the generator/eval path bitwise unchanged. Adds `test_qwen3_5_gdn_packing.py`: - `cu_seqlens` derivation (CPU): single sample -> None; packed -> boundaries. - GDN forward (GPU): packed + positions matches per-sample outputs to bf16, while packing without positions is clearly contaminated at the second sample. Both tests pass on H100 (packed-vs-separate max abs diff ~1e-4). --- .../rl/tests/test_qwen3_5_gdn_packing.py | 97 +++++++++++++ torchtitan/models/qwen3_5/model.py | 132 ++++++++++++++++-- 2 files changed, 220 insertions(+), 9 deletions(-) create mode 100644 torchtitan/experiments/rl/tests/test_qwen3_5_gdn_packing.py diff --git a/torchtitan/experiments/rl/tests/test_qwen3_5_gdn_packing.py b/torchtitan/experiments/rl/tests/test_qwen3_5_gdn_packing.py new file mode 100644 index 0000000000..00c6494423 --- /dev/null +++ b/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 diff --git a/torchtitan/models/qwen3_5/model.py b/torchtitan/models/qwen3_5/model.py index 9b22eb4090..43b1f1498e 100644 --- a/torchtitan/models/qwen3_5/model.py +++ b/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) From 9d09888737f8b7b61b19129538432366d19e6714 Mon Sep 17 00:00:00 2001 From: yichuan-w Date: Thu, 9 Jul 2026 02:02:44 -0700 Subject: [PATCH 2/4] [qwen3_5] GDN: reuse create_varlen_metadata_for_document for packing boundaries Address review: instead of a bespoke _cu_seqlens_from_positions derived per GDN layer, reuse the shared create_varlen_metadata_for_document (the same `positions == 0` document convention used to build the full-attention masks) and compute the cu_seqlens once in Qwen35Model.forward, threading it to the layers. The boundaries are derived from `positions` rather than read off `attention_masks` on purpose: GatedDeltaNet uses the FLA kernels and always needs sample boundaries, independent of the full-attention backend. `attention_masks` only carries cu_seqlens when the full-attention layers use the varlen backend; for the flex backend it is a BlockMask, and for a pipeline stage holding only GatedDeltaNet layers it is None (decoder.get_attention_masks returns None with no full-attention layers). Deriving from positions keeps GDN correct in all of these cases while computing the metadata a single time per forward. _gdn_cu_seqlens returns None for a single unpacked sequence so the generator / eval path stays on the non-varlen kernel (bitwise unchanged). Verified on H100: unit tests pass (packed-vs-separate ~1e-4; contamination without cu_seqlens) and a debug-model end-to-end forward with packed positions is finite. --- .../rl/tests/test_qwen3_5_gdn_packing.py | 48 ++++++------- torchtitan/models/qwen3_5/model.py | 69 +++++++++++-------- 2 files changed, 64 insertions(+), 53 deletions(-) diff --git a/torchtitan/experiments/rl/tests/test_qwen3_5_gdn_packing.py b/torchtitan/experiments/rl/tests/test_qwen3_5_gdn_packing.py index 00c6494423..5a4b5bf1d7 100644 --- a/torchtitan/experiments/rl/tests/test_qwen3_5_gdn_packing.py +++ b/torchtitan/experiments/rl/tests/test_qwen3_5_gdn_packing.py @@ -9,13 +9,14 @@ 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 +trainer computes wrong logprobs for every non-first packed sample (full attention +avoids this via a block-diagonal mask). This test pins: + +1. cu_seqlens derivation (``_gdn_cu_seqlens``, which reuses the shared + ``create_varlen_metadata_for_document``): a single sequence -> None (unchanged + path, no regression); a packed sequence -> the boundary offsets. +2. Packed + cu_seqlens matches processing each sample separately (the fix). +3. Packed WITHOUT cu_seqlens 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 @@ -26,26 +27,26 @@ import torch from torchtitan.models.qwen3_5 import _qwen35_deltanet_config -from torchtitan.models.qwen3_5.model import _cu_seqlens_from_positions +from torchtitan.models.qwen3_5.model import _gdn_cu_seqlens 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. +def test_gdn_cu_seqlens(): + # Single sequence (positions [0..L-1]) has one boundary -> None -> old path. single = torch.arange(50).unsqueeze(0) - assert _cu_seqlens_from_positions(single) is None + assert _gdn_cu_seqlens(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) + cu = _gdn_cu_seqlens(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) + cu_mrope = _gdn_cu_seqlens(mrope) assert cu_mrope.tolist() == [0, 30, 70] @@ -75,23 +76,24 @@ def test_gated_delta_net_packing_matches_separate(): 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 + pos_packed = torch.cat( # positions restart at 0 for the second sample + [torch.arange(len_a, device=dev), torch.arange(len_b, device=dev)] + ).unsqueeze(0) + cu_seqlens = _gdn_cu_seqlens(pos_packed) 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) + out_a = gdn(x_a) # single sequence -> cu_seqlens None + out_b = gdn(x_b) + out_fix = gdn(x_packed, cu_seqlens) # packed WITH cu_seqlens (the fix) + out_bug = gdn(x_packed) # packed WITHOUT cu_seqlens (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}" + assert fix_b < 1e-3, f"packed+cu_seqlens should match separate, got {fix_b}" + # Without cu_seqlens the recurrent state / conv bleed -> clearly contaminated. + assert bug_b > 10 * fix_b, f"expected contamination without cu_seqlens, 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 diff --git a/torchtitan/models/qwen3_5/model.py b/torchtitan/models/qwen3_5/model.py index 43b1f1498e..285864fe48 100644 --- a/torchtitan/models/qwen3_5/model.py +++ b/torchtitan/models/qwen3_5/model.py @@ -21,7 +21,11 @@ from torch.distributed.tensor.experimental import local_map from torchtitan.models.common import Conv1d, FeedForward, Linear -from torchtitan.models.common.attention import AttentionMasksType, BaseAttention +from torchtitan.models.common.attention import ( + AttentionMasksType, + BaseAttention, + create_varlen_metadata_for_document, +) from torchtitan.models.common.decoder import Decoder from torchtitan.models.utils import get_moe_model_nparams_and_flops from torchtitan.protocols.module import Module @@ -36,32 +40,32 @@ 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). +def _gdn_cu_seqlens(positions: torch.Tensor) -> torch.Tensor | None: + """FLA ``cu_seqlens`` for packed GatedDeltaNet, or None for a single sequence. 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). + ``positions`` at 0 at each sample. GatedDeltaNet must reset its recurrent state + and causal conv at those boundaries (full attention gets a block-diagonal mask + instead). We reuse :func:`create_varlen_metadata_for_document` -- the same + ``positions == 0`` document convention used to build the full-attention masks -- + so both attention types agree on the boundaries. + + Unlike full attention, GatedDeltaNet always needs these boundaries regardless of + the full-attention backend, so the caller derives them from ``positions`` rather + than reading them off ``attention_masks`` (which is a flex ``BlockMask`` -- or + ``None`` for a pipeline stage with no full-attention layers -- and carries no + ``cu_seqlens``). Returns None for a single unpacked sequence so the generator / + eval path stays on the non-varlen kernel (bitwise unchanged). """ # 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: + cu_seqlens = create_varlen_metadata_for_document(positions).cu_seq_q + # A single unpacked sequence yields [0, total]; skip the varlen path. + if cu_seqlens.numel() <= 2: return None - end = torch.tensor([total], device=positions.device, dtype=torch.int32) - return torch.cat([starts, end]) + return cu_seqlens def _torch_native_gated_delta( @@ -419,16 +423,13 @@ def _conv(x_local: torch.Tensor, w_local: torch.Tensor) -> torch.Tensor: return F.silu(x).transpose(1, 2) def forward( - self, x: torch.Tensor, positions: torch.Tensor | None = None + self, x: torch.Tensor, cu_seqlens: 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 - ) + # cu_seqlens holds the packed sample boundaries (computed once at the model + # level); it resets the causal conv AND the recurrent state at each boundary. + # None (single unpacked sequence / generator decode) -> unchanged path. # Shapes: # xq, xk: (bs, seqlen, n_key_heads * key_head_dim) @@ -595,14 +596,15 @@ def forward( x: torch.Tensor, attention_masks: AttentionMasksType | None, positions: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, ) -> torch.Tensor: h = self.attention_norm(x) if self.full_attn: h = self.attn(h, attention_masks, positions) else: - # GatedDeltaNet needs positions to reset conv/recurrent state at packed - # sample boundaries (softmax uses attention_masks instead). - h = self.attn(h, positions) + # GatedDeltaNet resets its conv/recurrent state at the packed sample + # boundaries in cu_seqlens (full attention uses attention_masks instead). + h = self.attn(h, cu_seqlens) x = x + h h = self.ffn_norm(x) @@ -883,8 +885,15 @@ def forward( # pyrefly: ignore [bad-override] # 3D MRoPE positions for multimodal batches, else 2D text positions. rope_positions = mrope_positions if mrope_positions is not None else positions assert rope_positions is not None + # GatedDeltaNet packing boundaries, computed once and shared across layers + # (see _gdn_cu_seqlens). Derived from the same document positions as the + # full-attention masks, but independent of the attention backend so it is + # correct for a pipeline stage that holds only GatedDeltaNet layers. + cu_seqlens = _gdn_cu_seqlens( + positions if positions is not None else rope_positions + ) for layer in self.layers.values(): - x = layer(x, attention_masks, rope_positions) + x = layer(x, attention_masks, rope_positions, cu_seqlens) x = self.norm(x) if self.norm is not None else x if self._skip_lm_head: From c21fceac8a0352b5c035f7741ff5fc844d630b80 Mon Sep 17 00:00:00 2001 From: yichuan-w Date: Sun, 12 Jul 2026 17:26:09 -0700 Subject: [PATCH 3/4] [qwen3_5] GDN packing: address review (inline cu_seqlens, local_map conv, tests) - Inline the cu_seqlens derivation into Qwen35Model.forward; drop the thin _gdn_cu_seqlens wrapper and the MRoPE special-case (derive from positions). - Reset the packed causal conv via fla.causal_conv1d under local_map so it works under TP (channel-sharded), replacing the non-TP assert + full_tensor gather. - Add a CP NOTE: the cu_seqlens boundaries are not context-parallel aware. - Tests: add tests/unit_tests/test_attention.py covering create_varlen_metadata_for_document; reframe the GDN test as a standard packed-vs-separate correctness check (drop the bug-demo half). Lint clean (pre-commit; the 2 residual pyrefly errors are pre-existing env torch-version noise in untouched distributed/utils.py + deepep/hybridep.py). --- tests/unit_tests/test_attention.py | 35 ++++++ .../rl/tests/test_qwen3_5_gdn_packing.py | 61 +++------- torchtitan/models/qwen3_5/model.py | 114 ++++++++---------- 3 files changed, 98 insertions(+), 112 deletions(-) create mode 100644 tests/unit_tests/test_attention.py diff --git a/tests/unit_tests/test_attention.py b/tests/unit_tests/test_attention.py new file mode 100644 index 0000000000..f18871ea3d --- /dev/null +++ b/tests/unit_tests/test_attention.py @@ -0,0 +1,35 @@ +# 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. + +import torch + +from torchtitan.models.common.attention import create_varlen_metadata_for_document + + +def test_single_document(): + # One document per row: positions never reset after 0, so the only boundary + # is the sequence end. + positions = torch.arange(50).unsqueeze(0) + meta = create_varlen_metadata_for_document(positions) + assert meta.cu_seq_q.tolist() == [0, 50] + assert meta.cu_seq_k.tolist() == [0, 50] + assert meta.max_q == 50 + + +def test_packed_documents(): + # Two documents packed in one row (positions restart at 0 at the boundary). + positions = torch.cat([torch.arange(30), torch.arange(40)]).unsqueeze(0) + meta = create_varlen_metadata_for_document(positions) + assert meta.cu_seq_q.tolist() == [0, 30, 70] + assert meta.max_q == 40 + + +def test_batched_rows_are_flattened(): + # Boundaries are cumulative over the flattened [batch * seq_len] layout, so + # each row start is also a boundary. + positions = torch.arange(20).unsqueeze(0).repeat(2, 1) + meta = create_varlen_metadata_for_document(positions) + assert meta.cu_seq_q.tolist() == [0, 20, 40] diff --git a/torchtitan/experiments/rl/tests/test_qwen3_5_gdn_packing.py b/torchtitan/experiments/rl/tests/test_qwen3_5_gdn_packing.py index 5a4b5bf1d7..f27e57ce05 100644 --- a/torchtitan/experiments/rl/tests/test_qwen3_5_gdn_packing.py +++ b/torchtitan/experiments/rl/tests/test_qwen3_5_gdn_packing.py @@ -4,52 +4,27 @@ # 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). +"""GatedDeltaNet packed-sequence correctness. -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 (full attention -avoids this via a block-diagonal mask). This test pins: +When several samples are packed into one row (``positions`` restart at 0 per +sample), GatedDeltaNet is given ``cu_seqlens`` marking the boundaries so its +recurrent state and causal conv reset per sample. This test checks that a packed +row processed with ``cu_seqlens`` matches processing each sample on its own. -1. cu_seqlens derivation (``_gdn_cu_seqlens``, which reuses the shared - ``create_varlen_metadata_for_document``): a single sequence -> None (unchanged - path, no regression); a packed sequence -> the boundary offsets. -2. Packed + cu_seqlens matches processing each sample separately (the fix). -3. Packed WITHOUT cu_seqlens 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. +Needs a GPU: the FLA gated-delta kernels are Triton/CUDA only. """ import pytest import torch +from torchtitan.models.common.attention import create_varlen_metadata_for_document from torchtitan.models.qwen3_5 import _qwen35_deltanet_config -from torchtitan.models.qwen3_5.model import _gdn_cu_seqlens def _max_abs_diff(x: torch.Tensor, y: torch.Tensor) -> float: return (x - y).abs().max().item() -def test_gdn_cu_seqlens(): - # Single sequence (positions [0..L-1]) has one boundary -> None -> old path. - single = torch.arange(50).unsqueeze(0) - assert _gdn_cu_seqlens(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 = _gdn_cu_seqlens(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 = _gdn_cu_seqlens(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) @@ -77,23 +52,17 @@ def test_gated_delta_net_packing_matches_separate(): x_a = torch.randn(1, len_a, dim, device=dev, dtype=dt) x_b = torch.randn(1, len_b, dim, device=dev, dtype=dt) x_packed = torch.cat([x_a, x_b], dim=1) - pos_packed = torch.cat( # positions restart at 0 for the second sample + # Packed positions restart at 0 for the second sample. + positions = torch.cat( [torch.arange(len_a, device=dev), torch.arange(len_b, device=dev)] ).unsqueeze(0) - cu_seqlens = _gdn_cu_seqlens(pos_packed) + cu_seqlens = create_varlen_metadata_for_document(positions).cu_seq_q with torch.no_grad(): - out_a = gdn(x_a) # single sequence -> cu_seqlens None + out_a = gdn(x_a) # each sample on its own (no packing -> cu_seqlens None) out_b = gdn(x_b) - out_fix = gdn(x_packed, cu_seqlens) # packed WITH cu_seqlens (the fix) - out_bug = gdn(x_packed) # packed WITHOUT cu_seqlens (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+cu_seqlens should match separate, got {fix_b}" - # Without cu_seqlens the recurrent state / conv bleed -> clearly contaminated. - assert bug_b > 10 * fix_b, f"expected contamination without cu_seqlens, got {bug_b}" + out_packed = gdn(x_packed, cu_seqlens) - # The first packed sample is unaffected either way (nothing precedes it). - assert _max_abs_diff(out_bug[:, :len_a], out_a) == 0.0 + # Each packed sample matches its standalone output (to bf16 tolerance). + assert _max_abs_diff(out_packed[:, :len_a], out_a) < 1e-3 + assert _max_abs_diff(out_packed[:, len_a:], out_b) < 1e-3 diff --git a/torchtitan/models/qwen3_5/model.py b/torchtitan/models/qwen3_5/model.py index 285864fe48..6cc3801724 100644 --- a/torchtitan/models/qwen3_5/model.py +++ b/torchtitan/models/qwen3_5/model.py @@ -40,34 +40,6 @@ 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 _gdn_cu_seqlens(positions: torch.Tensor) -> torch.Tensor | None: - """FLA ``cu_seqlens`` for packed GatedDeltaNet, or None for a single sequence. - - The RL batcher packs several training samples into one row and restarts - ``positions`` at 0 at each sample. GatedDeltaNet must reset its recurrent state - and causal conv at those boundaries (full attention gets a block-diagonal mask - instead). We reuse :func:`create_varlen_metadata_for_document` -- the same - ``positions == 0`` document convention used to build the full-attention masks -- - so both attention types agree on the boundaries. - - Unlike full attention, GatedDeltaNet always needs these boundaries regardless of - the full-attention backend, so the caller derives them from ``positions`` rather - than reading them off ``attention_masks`` (which is a flex ``BlockMask`` -- or - ``None`` for a pipeline stage with no full-attention layers -- and carries no - ``cu_seqlens``). Returns None for a single unpacked sequence so the generator / - eval path stays on the non-varlen kernel (bitwise unchanged). - """ - # 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] - cu_seqlens = create_varlen_metadata_for_document(positions).cu_seq_q - # A single unpacked sequence yields [0, total]; skip the varlen path. - if cu_seqlens.numel() <= 2: - return None - return cu_seqlens - - def _torch_native_gated_delta( q: torch.Tensor, k: torch.Tensor, @@ -360,32 +332,44 @@ 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. + # (else the first conv_kernel_size-1 tokens of each sample would see the + # previous one). F.conv1d cannot express a per-segment reset, so the packed + # path uses fla.causal_conv1d, which resets at cu_seqlens. The conv is + # depthwise (per-channel), so under TP (channel-sharded) we run it on local + # shards via local_map -- same pattern as the F.conv1d path below. GDN conv + # is bias-free. 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) + bs, seqlen, _ = x.shape + + def _fla_conv(x_local: torch.Tensor, w_local: torch.Tensor) -> torch.Tensor: + # x_local: [bs, seqlen, C_local]; depthwise weight [C_local, 1, k]. + # fla wants one [1, total, C] sequence; cu_seqlens spans the + # flattened [bs*seqlen] layout (offsets already include batch rows). + # fla's causal_conv1d is untyped (pyrefly reads it as a Tensor). + y = _fla_causal_conv1d( + x_local.reshape(1, bs * seqlen, -1), + weight=w_local.squeeze(1), + bias=None, + activation="silu", + cu_seqlens=cu_seqlens, + ) + if isinstance(y, tuple): + y = y[0] + return y.reshape(bs, seqlen, -1) + + if isinstance(x, DTensor): + # Channel-sharded depthwise conv: run per shard, restore DTensor-ness. + x_plc = x.placements + w_plc = conv.weight.placements # pyrefly: ignore [missing-attribute] + conv_dt = local_map( + _fla_conv, + out_placements=(x_plc,), + in_placements=(x_plc, w_plc), + in_grad_placements=(x_plc, w_plc), + device_mesh=x.device_mesh, + ) + return conv_dt(x, conv.weight) # pyrefly: ignore + return _fla_conv(x, conv.weight) # pyrefly: ignore [bad-argument-type] x = F.pad(x.transpose(1, 2), [self.conv_kernel_size - 1, 0]) if isinstance(x, DTensor): @@ -427,10 +411,6 @@ def forward( ) -> torch.Tensor: bs, seqlen, _ = x.shape - # cu_seqlens holds the packed sample boundaries (computed once at the model - # level); it resets the causal conv AND the recurrent state at each boundary. - # None (single unpacked sequence / generator decode) -> unchanged path. - # Shapes: # xq, xk: (bs, seqlen, n_key_heads * key_head_dim) # xv, xz: (bs, seqlen, n_value_heads * value_head_dim) @@ -602,8 +582,6 @@ def forward( if self.full_attn: h = self.attn(h, attention_masks, positions) else: - # GatedDeltaNet resets its conv/recurrent state at the packed sample - # boundaries in cu_seqlens (full attention uses attention_masks instead). h = self.attn(h, cu_seqlens) x = x + h @@ -885,13 +863,17 @@ def forward( # pyrefly: ignore [bad-override] # 3D MRoPE positions for multimodal batches, else 2D text positions. rope_positions = mrope_positions if mrope_positions is not None else positions assert rope_positions is not None - # GatedDeltaNet packing boundaries, computed once and shared across layers - # (see _gdn_cu_seqlens). Derived from the same document positions as the - # full-attention masks, but independent of the attention backend so it is - # correct for a pipeline stage that holds only GatedDeltaNet layers. - cu_seqlens = _gdn_cu_seqlens( - positions if positions is not None else rope_positions - ) + # cu_seqlens holds the packed sample boundaries (positions == 0), computed + # once and shared across GatedDeltaNet layers so their recurrent state and + # causal conv reset per sample. Reuses the same document-varlen metadata as + # the full-attention masks. None for a single unpacked sequence (unchanged + # path). NOTE: not context-parallel aware (CP would split the sequence and + # break these boundaries); CP is unsupported for GatedDeltaNet. + cu_seqlens = None + if positions is not None: + cu_seqlens = create_varlen_metadata_for_document(positions).cu_seq_q + if cu_seqlens.numel() <= 2: # single unpacked sample -> non-varlen path + cu_seqlens = None for layer in self.layers.values(): x = layer(x, attention_masks, rope_positions, cu_seqlens) From af65968dbe1cc7d0318a4df5462d1a45a62baf98 Mon Sep 17 00:00:00 2001 From: yichuan-w Date: Sun, 12 Jul 2026 17:34:21 -0700 Subject: [PATCH 4/4] [qwen3_5] GDN: make kernel cu_seqlens keyword-only so packing works under TP The declarative TP sharding validates that every positional forward arg has an in_dst_shardings entry when local_map is set. cu_seqlens is replicated metadata (not a sharded input), so making it keyword-only excludes it from that check and lets local_map pass it through -- same treatment as attention_masks/positions for full attention. Verified on 2 GPUs: packed GatedDeltaNet forward under TP=2 is bitwise-identical to the non-TP forward (max_abs_diff 0). --- torchtitan/models/qwen3_5/model.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/torchtitan/models/qwen3_5/model.py b/torchtitan/models/qwen3_5/model.py index 6cc3801724..53238b994f 100644 --- a/torchtitan/models/qwen3_5/model.py +++ b/torchtitan/models/qwen3_5/model.py @@ -217,6 +217,9 @@ def forward( v: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, + *, + # Keyword-only so the TP sharding's local_map treats it as replicated + # pass-through metadata (like attention_masks), not a sharded input. cu_seqlens: torch.Tensor | None = None, ) -> torch.Tensor: # Expand Q/K heads to match V when n_value_heads > n_key_heads @@ -432,7 +435,7 @@ def forward( 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, cu_seqlens) + output = self.kernel(xq, xk, xv, g, beta, cu_seqlens=cu_seqlens) xz = xz.view(bs, seqlen, -1, self.value_head_dim) output = self.norm(output, xz)