[megatron, model] fix: support GDN with packed THD sequence parallel

[dengoswei]

What does this PR do?

Enable Megatron GatedDeltaNet to run with verl packed THD inputs under sequence parallelism.

Megatron's current GDN forward path does not support packed_seq_params. In verl's packed THD path, sequence boundaries are carried by cu_seqlens; this PR passes those boundaries to the varlen-aware GDN kernels instead of rejecting the packed input.

The patch:

• Adds a focused apply_patch_megatron_gated_delta_net() hook.
• Uses full sequence length after sequence-parallel gather for downstream GDN reshapes.
• Passes cu_seqlens to FLA causal_conv1d and chunk_gated_delta_rule.
• Builds seq_idx from cu_seqlens for the causal_conv1d_fn fallback.
• Rejects packed deterministic mode explicitly because the torch fallback is not varlen-aware.

Scope is limited to the MCore GDN packed-THD compute path. The runtime-boundary
changes (MTP, LoRA, Megatron-Bridge, vLLM weight sync, optimizer offload) are
intentionally out of scope and are tracked separately in #5599.

Duplicate-work check

• GDN / GatedDeltaNet open PRs reviewed: [fsdp] fix: make GDN work correctly when using ulysses and varlen sequence #5346, [megatron] fix: Qwen3.5 LoRA & MTP support (with Megatron-Bridge) #5599, [perf] feat: Add Qwen3.5 MFU flops estimation #6389.

Conclusion: no duplicate. None of the open PRs touch the Megatron/MCore GDN
packed_seq_params compute path:

• [fsdp] fix: make GDN work correctly when using ulysses and varlen sequence #5346 handles FSDP/Ulysses/varlen GDN — different backend (FSDP, not Megatron).
• [megatron] fix: Qwen3.5 LoRA & MTP support (with Megatron-Bridge) #5599 handles Qwen3.5 Megatron-Bridge, MTP, LoRA, and vLLM weight-sync — bundles the runtime-boundary changes this PR deliberately avoids.
• [perf] feat: Add Qwen3.5 MFU flops estimation #6389 adds Qwen3.5 MFU flops estimation — unrelated.

Test

CPU unit tests (tests/utils/test_megatron_gdn_patch_on_cpu.py):

python -m pytest tests/utils/test_megatron_gdn_patch_on_cpu.py
pre-commit run --files verl/models/mcore/patch.py \
  verl/workers/engine/megatron/transformer_impl.py \
  tests/utils/test_megatron_gdn_patch_on_cpu.py --show-diff-on-failure

Results (torch 2.9.1, mcore 0.17.0):

• 3 tests pass: cu_seqlens is threaded into the FLA varlen conv + gated-delta-rule kernels; the causal_conv1d_fn fallback builds the correct seq_idx ([[0,0,1,1]] for two length-2 sequences); packed deterministic mode is rejected.
• Pre-commit on touched files passes.

Note: an earlier revision of the test stubbed sys.modules["megatron.core"]
before importing verl.models.mcore.patch, which broke the import on any
machine with torch installed (the import was previously skipped on a
torch-less machine). Fixed by importing the patch entrypoint before
installing the fakes; both/all tests now run (not skip) and pass.

GPU validation (8×H100, mcore 0.17.0, torch 2.9.1, FLA):

A real patched GatedDeltaNet.forward was driven with real nn.Linear /
nn.Conv1d / FLA kernels and shared weights, comparing packed-THD output
(cu_seqlens=[0,16,40,80,88]) against the unpacked per-sequence reference:

• fla.modules.convolution.causal_conv1d + chunk_gated_delta_rule (production path): max_abs_diff = 0.0 across all sequences — bit-identical, no cross-sequence leakage. Backward grads finite.
• causal_conv1d_fn + seq_idx fallback path: max_abs_diff = 0.0, backward grads finite.

Scope of this validation: it confirms that cu_seqlens / seq_idx are threaded
correctly and that the varlen conv + gated-delta-rule kernels keep packed
sequences independent on synthetic, shared-weight inputs. It does not by
itself establish end-to-end training numerical stability with real model weights
and data. Before broad rollout, an in-model A/B comparison (packed/THD vs
padded) on the target model — checking first-step grad_norm and loss — and a
multi-rank TP>1 + sequence_parallel=True run (exercising the Megatron SP
gather) are recommended.

API and Usage Example

No user-facing API change. The patch is applied from the existing Megatron engine patch hook.

Design & Code Changes

• verl/models/mcore/patch.py: adds apply_patch_megatron_gated_delta_net.
• verl/workers/engine/megatron/transformer_impl.py: applies the GDN patch alongside existing Megatron CUDA patches.
• tests/utils/test_megatron_gdn_patch_on_cpu.py: monkeypatch tests covering the fla path, the causal_conv1d_fn seq_idx fallback, and packed-deterministic rejection.

AI assistance

This PR is paired with Claude. AI assistance was used to prepare and validate this PR. The human submitter has reviewed every changed line, verified the duplicate-work analysis, and run the GPU validation above.

Checklist Before Submitting

• Read the Contribute Guide.
• Search for similar PRs. Query links are included above.
• Format the PR title as [{modules}] {type}: {description}.
• Apply pre-commit checks on touched files.
• Add focused test coverage.
• Run GPU validation for packed THD GDN (forward/backward numerical equivalence on 8×H100).

[CLAassistant]

All committers have signed the CLA.

[gemini-code-assist]

Code Review

This pull request introduces a patch to enable Megatron's GatedDeltaNet for packed THD sequence-parallel inputs by forwarding sequence boundaries to the FLA gated-delta-rule and a varlen-aware causal convolution path. It also adds corresponding unit tests and registers the patch in the transformer implementation. Feedback was provided to address a performance bottleneck in _build_seq_idx_from_cu_seqlens where evaluating a CUDA tensor value in Python control flow (cu_seqlens[-1] < total_tokens) causes a device-to-host synchronization. A suggestion was made to unconditionally append total_tokens to avoid this check.

[gemini-code-assist]

The comparison cu_seqlens[-1] < total_tokens triggers a device-to-host synchronization because it evaluates a CUDA tensor value in Python control flow. This synchronization blocks the CPU until the GPU finishes all preceding operations, causing a significant performance bottleneck during training.

By unconditionally appending total_tokens to cu_seqlens, we can avoid the conditional check entirely. Since total_tokens >= cu_seqlens[-1], the difference will be non-negative, and any extra 0 in seq_lengths is naturally ignored by torch.repeat_interleave without changing the output.

Suggested change

	if cu_seqlens[-1] < total_tokens:
	cu_seqlens = torch.cat([cu_seqlens, cu_seqlens.new_tensor([total_tokens])])
	seq_lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).clamp(min=0)
	cu_seqlens = torch.cat([cu_seqlens, cu_seqlens.new_tensor([total_tokens])])
	seq_lengths = cu_seqlens[1:] - cu_seqlens[:-1]

[wuxibin89]

Megatron-LM@dev has supported GDN with THD+CP:

• [dev] feat(moe): Support packed sequence for gated delta net (GDN) NVIDIA/Megatron-LM#2644
• [dev] fix(ssm): handle alignment padding in GDN packed seq + CP NVIDIA/Megatron-LM#4230

[dengoswei]

YES, I am aware of that.

T'o use the sp & cp support in Megatron-LM, we need to upgrade the VeRL to using megatron mcore 0.18. I am not the timeline of VeRL + mcore 0.18, so..

[tnlin]

That's awesome. When will the PRs be merged? I can't wait to use them

[wuxibin89]

Please upgrade megatron>=0.18.0