[RL][qwen3_5] GDN: reset conv/recurrent state at packed sample boundaries by yichuan-w · Pull Request #3892 · pytorch/torchtitan

yichuan-w · 2026-07-09T06:05:09Z

Problem

The RL trainer packs (next-fits) several training samples into one row and
restarts positions at 0 at each sample boundary. Full-attention layers keep
packed samples independent with a block-diagonal attention mask, but
GatedDeltaNet was invoked as self.attn(h) with no boundary information.

GatedDeltaNet is linear attention: it carries a recurrent state across the
sequence and a causal conv1d whose window reaches back conv_kernel_size - 1
tokens. With no boundaries, both the recurrent state and the conv window bled
across packed sample boundaries, so every non-first packed sample in a row
saw leakage from the previous sample. The trainer (which packs) then computed
wrong logprobs for those tokens while the vLLM generator (which does not pack)
did not, yielding garbage importance ratios and flat RL reward. At long
sequence lengths with ~3 samples per row, this corrupts the majority of trained
tokens.

Fix

Derive FLA cu_seqlens from the packed positions and thread it through
Qwen35TransformerBlock -> GatedDeltaNet -> GatedDeltaKernel:

A positions == 0 index marks a sample start. This reuses the exact
convention already used for the softmax document mask in
models/common/attention.py (doc_ids = cumsum(positions == 0) - 1), so the
two attention types now agree on where sample boundaries are.
cu_seqlens is passed to the FLA chunk / fused-recurrent kernels and to
fla.causal_conv1d, both of which reset per segment. The recurrent state and
the conv window no longer cross sample boundaries.
A single, unpacked sequence has exactly one positions == 0 index, so
_cu_seqlens_from_positions returns None and the generator / eval / single
sequence path is bitwise unchanged.

The new arguments are optional (None default), so all existing callers are
unaffected.

Numerical validation

torchtitan/experiments/rl/tests/test_qwen3_5_gdn_packing.py (passes on H100):

test_cu_seqlens_from_positions (CPU): single sample -> None; packed
[0..29, 0..39] -> boundaries [0, 30, 70]; also covers 3D MRoPE positions.
test_gated_delta_net_packing_matches_separate (GPU): a packed two-sample
row processed with positions matches the two samples processed separately to
bf16 (max abs diff ~1e-4), while packing without positions is clearly
contaminated at the second sample (>10x larger diff). The first packed sample
is bitwise identical either way (nothing precedes it).

Notes / limitations

The packed conv path uses fla.causal_conv1d on flattened [1, B*L, C]
activations and asserts non-TP (FSDP) activations; channel-sharded TP conv
still uses the existing local_map path (which does not carry cu_seqlens).
This matches the current RL-training setup for GDN (FSDP, no TP). GDN training
runs eager, so the in-forward cu_seqlens computation is fine.
cu_seqlens is derived per GDN layer from the same positions; this is a
small per-layer nonzero and keeps the change local to the GDN model.

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).

tianyu-l · 2026-07-09T07:33:05Z

    return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)


+def _cu_seqlens_from_positions(positions: torch.Tensor) -> torch.Tensor | None:


Haven't looked into details, but I wonder if it can be refactored so that it follows the structure of how varlen metadata is sent to the model
https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/decoder.py#L327

good point, reuse the meta data helper can make things cleaner

Done in the follow-up commit. It now reuses create_varlen_metadata_for_document (same positions == 0 document convention as the full-attention masks) and computes the cu_seqlens once in Qwen35Model.forward, threading it to the layers instead of re-deriving it per GDN layer.

One subtlety worth calling out: I derive it from positions rather than reading it off attention_masks. GatedDeltaNet uses the FLA kernels and always needs the sample boundaries, independent of the full-attention backend — but attention_masks only carries cu_seqlens when full attention uses the varlen backend. With flex it's a BlockMask, and for a pipeline stage holding only GatedDeltaNet layers get_attention_masks returns None (decoder.py). Deriving from positions keeps GDN correct in all of those cases while still computing the metadata a single time per forward.

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.

…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.

shuhuayu · 2026-07-10T06:40:18Z

+            assert not isinstance(
+                x, DTensor
+            ), "packed GDN conv (cu_seqlens) assumes non-TP (FSDP) activations"


this prevents any training with tp enabled? why is this?

fixed, now we are good to use TP

shuhuayu · 2026-07-10T06:52:34Z

+        # 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:


i am not sure we should move to _fla_causal_conv1d. 2) if we use both _fla_causal_conv1d and torch's F.conv1d, we should make them both work with cu_seqlens, and cu_seqlens can be None and no checking needed, and we need to make it configurable if both are supported.

cc: @tianyu-l

F.conv1d can't reset the conv window at sample boundaries, so the packed path needs fla (it resets at cu_seqlens). Unpacked stays on F.conv1d; dispatch is just cu_seqlens is None.

I dropped the non-TP assert — the packed conv now runs through local_map on channel shards (depthwise conv is per-channel), so TP works too. Verified on 2 GPUs: packed GatedDeltaNet forward is bitwise-equal to non-TP (max_abs_diff 0).

Didn't add a backend switch since F.conv1d can't serve the packed case anyway — happy to add one if you'd prefer.

shuhuayu · 2026-07-10T06:54:20Z

+            bs, seqlen, channels = x.shape
+            weight = conv.weight
+            bias = conv.bias
+            if isinstance(weight, DTensor):


we should use local_map to handle non-DTensor part, see the following case as an example.

shuhuayu · 2026-07-10T06:57:36Z

+        # 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.
+


nit: remove

shuhuayu · 2026-07-10T06:59:02Z

+        # 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.


nit:

Suggested change

# 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 holds the packed sample boundaries

shuhuayu · 2026-07-10T07:05:30Z

+    """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].


nit: too long, please be succinct.

shuhuayu · 2026-07-10T07:06:52Z

    return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)


+def _gdn_cu_seqlens(positions: torch.Tensor) -> torch.Tensor | None:


this looks like a very thin wrapper, why not inline it in the forward method.

shuhuayu · 2026-07-10T07:11:42Z

+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:
+
+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.


plz frame this as a standard test, not a fix other people have no context.

shuhuayu · 2026-07-10T07:12:25Z

+    return (x - y).abs().max().item()
+
+
+def test_gdn_cu_seqlens():


given that gdn_cu_seqlens is a thin wrapper, this test should belong to create_varlen_metadata_for_document? can you check whether it is already there? if not, let's add one for create_varlen_metadata_for_document.

shuhuayu · 2026-07-10T07:16:18Z

+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="FLA GDN kernels need CUDA")
+def test_gated_delta_net_packing_matches_separate():


there is a unnecessary bug test within this test. we only need to test the standard gdn works. again, people do not have this bug context, it should be avoided, and frame this as a clean standard test.

tianyu-l · 2026-07-10T07:43:15Z

+        # (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(


It feels to me that this will create another difficulty for CP, so worth adding a NOTE, although CP is not supported for GDN anyway. cc @fegin

…onv, 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).

…nder 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).

yichuan-w requested review from fegin, tianyu-l, wconstab and wwwjn as code owners July 9, 2026 06:05

pytorch-bot Bot added ciflow/8gpu ciflow/rl labels Jul 9, 2026

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jul 9, 2026

yichuan-w changed the title ~~[qwen3_5] GDN: reset conv/recurrent state at packed sample boundaries~~ [RL][qwen3_5] GDN: reset conv/recurrent state at packed sample boundaries Jul 9, 2026

tianyu-l reviewed Jul 9, 2026

View reviewed changes

tianyu-l requested a review from shuhuayu July 9, 2026 07:33

yichuan-w requested a review from tianyu-l July 9, 2026 23:55

shuhuayu requested changes Jul 10, 2026

View reviewed changes

tianyu-l reviewed Jul 10, 2026

View reviewed changes

yichuan-w added 2 commits July 12, 2026 17:26

yichuan-w requested review from shuhuayu and tianyu-l July 13, 2026 00:42

		return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)


		def _cu_seqlens_from_positions(positions: torch.Tensor) -> torch.Tensor \| None:

		return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)


		def _gdn_cu_seqlens(positions: torch.Tensor) -> torch.Tensor \| None:



		@pytest.mark.skipif(not torch.cuda.is_available(), reason="FLA GDN kernels need CUDA")
		def test_gated_delta_net_packing_matches_separate():

Uh oh!

Conversation

yichuan-w commented Jul 9, 2026

Problem

Fix

Numerical validation

Notes / limitations

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants