[dev] fix(ssm): handle alignment padding in GDN packed seq + CP by yxs · Pull Request #4230 · NVIDIA/Megatron-LM

yxs · 2026-04-09T11:28:44Z

What does this PR do ?

Fix shape mismatch in GDN packed sequence + CP > 1 when alignment padding is present.

Summary

When cu_seqlens_q_padded is not provided, GDN falls back to cu_seqlens_q (actual token boundaries) for the per-sequence all-to-all unpack. Alignment padding tokens are lost, making the output shorter than hidden_states, crashing at the residual connection.

Fix: Extract cu_seqlens resolution into _resolve_cu_seqlens():

cu_seqlens_padded available → use it (unchanged behavior)
CP = 1 → use actual (no all-to-all, no issue)
Single sequence + padding → auto-extend to cover trailing padding
Multi sequence + padding → raise ValueError requesting cu_seqlens_padded

Test results

Unit tests: 5/5 passed

(2× H100, CP=2):

Single seq + padding + no cu_seqlens_padded → auto-extend, output shape correct
Multi seq + padding + no cu_seqlens_padded → ValueError raised correctly
With cu_seqlens_padded set → standard path works
No padding → passes through unchanged

Fixes #4194

PR for main: #2645

copy-pr-bot · 2026-04-09T11:28:50Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

yxs · 2026-04-09T11:32:00Z

cc @yuzhongw-nvidia

yuzhongw-nvidia · 2026-04-10T05:55:41Z

/claude review

claude

Looks good overall — clean fix with solid test coverage. Left one minor nit on the error message formatting.

yuzhongw-nvidia · 2026-04-10T06:05:03Z

/ok to test 8a5cbc3

yxs · 2026-04-11T08:11:30Z

@yuzhongw-nvidia is this PR ready to merge?

When cu_seqlens_q_padded is not provided and alignment padding exists, GDN's packed all-to-all loses padding tokens, causing shape mismatch at the residual connection. Fix: - Extract cu_seqlens resolution into _resolve_cu_seqlens() method - Single-sequence case: auto-extend cu_seqlens to cover trailing padding - Multi-sequence case: raise clear error requesting cu_seqlens_padded - Add unit tests for cu_seqlens validation logic Validated on 2x H100 (CP=2): - Single seq with padding + no cu_seqlens_padded: output shape correct - Multi seq with padding + no cu_seqlens_padded: ValueError raised - With cu_seqlens_padded set: standard path works - No padding case: passes through unchanged Fixes: NVIDIA#4194 Co-Authored-By: Claude <noreply@anthropic.com>

Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>

yuzhongw-nvidia · 2026-04-13T00:53:08Z

/ok to test a5b5110

svcnvidia-nemo-ci · 2026-04-13T04:56:13Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/24326368290

yuzhongw-nvidia · 2026-04-13T07:00:12Z

@yuzhongw-nvidia is this PR ready to merge?

Merged. Many thanks for your contribution!

Phlip79 · 2026-04-14T00:57:58Z

/claude copy

svcnvidia-nemo-ci · 2026-04-14T01:04:35Z

🤖 This PR has been cherry-picked to main in #4289.

…IA#4230) Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Yuzhong Wang <yuzhongw@nvidia.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>

Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Yuzhong Wang <yuzhongw@nvidia.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>

…IA#4230) Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Yuzhong Wang <yuzhongw@nvidia.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>

…DeltaNet Port the packed-sequence support from upstream NVIDIA/Megatron-LM PRs NVIDIA#2644 and NVIDIA#4230 into the GatedDeltaNet forward (adapted to this fork's GDN code): - Gate packed handling on `packed_seq_params.qkv_format == 'thd'` (not just `is not None`), so the sbhd PackedSeqParams Megatron passes uniformly is ignored and only genuine THD packing is processed. - Resolve cu_seqlens via `_resolve_cu_seqlens`, preferring *_padded and validating cu_seqlens[-1] == total_seq_len (alignment-padding handling). - Unpack the packed stream per-sequence around both CP all-to-all calls (cp2hp / hp2cp) via `_unpack_sequence`, and pass cu_seqlens into the causal_conv1d and gated_delta_rule kernels. - Guard torch_chunk_gated_delta_rule against cu_seqlens (unsupported). - Assert batch == 1 and non-deterministic mode for packed sequences. Also keep the independent fp32 numerics fix: cast A_log and dt_bias to fp32 before exp/softplus when computing g, matching the "In fp32" intent. Verified the cu_seqlens[-1] == seq_len invariant holds for the target training config (TP=4, sequence-parallel, CP=1, THD): seq_len = (L/tp)*sp*cp = L and cu_seqlens[-1] = L*cp = L. Note CP>1 + GDN is a known unsupported edge that would trip the validation. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

yxs requested review from a team as code owners April 9, 2026 11:28

github-actions Bot added the community-request label Apr 9, 2026

yxs force-pushed the fix/gdn-packed-cp-alignment branch from 45d1a49 to 7e32649 Compare April 9, 2026 11:31

yuzhongw-nvidia requested review from xiaoyao0115 and yuzhongw-nvidia April 10, 2026 00:44

yuzhongw-nvidia reviewed Apr 10, 2026

View reviewed changes

Comment thread megatron/core/ssm/gated_delta_net.py Outdated

yuzhongw-nvidia reviewed Apr 10, 2026

View reviewed changes

Comment thread megatron/core/ssm/gated_delta_net.py Outdated

yxs requested a review from yuzhongw-nvidia April 10, 2026 02:11

xiaoyao0115 reviewed Apr 10, 2026

View reviewed changes

Comment thread megatron/core/ssm/gated_delta_net.py Outdated

claude Bot reviewed Apr 10, 2026

View reviewed changes

Comment thread megatron/core/ssm/gated_delta_net.py Outdated

claude Bot reviewed Apr 10, 2026

View reviewed changes

svcnvidia-nemo-ci added this to the Core 0.16 milestone Apr 10, 2026

copy-pr-bot Bot temporarily deployed to test April 10, 2026 06:05 Inactive

yuzhongw-nvidia approved these changes Apr 10, 2026

View reviewed changes

yuzhongw-nvidia enabled auto-merge April 10, 2026 11:15

yxs and others added 5 commits April 12, 2026 17:52

Update the condition of GDN THD

9b155d1

Address review comments

43a115f

fix _resolve_cu_seqlens and add a UT

1256b85

Update megatron/core/ssm/gated_delta_net.py

a5b5110

Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>

yuzhongw-nvidia force-pushed the fix/gdn-packed-cp-alignment branch from 8a5cbc3 to a5b5110 Compare April 13, 2026 00:52

copy-pr-bot Bot temporarily deployed to test April 13, 2026 00:54 Inactive

Victarry approved these changes Apr 13, 2026

View reviewed changes

yuzhongw-nvidia added this pull request to the merge queue Apr 13, 2026

Merged via the queue into NVIDIA:dev with commit 0f6fcb0 Apr 13, 2026
61 checks passed

yuzhongw-nvidia mentioned this pull request Apr 13, 2026

GDN packed sequence (THD format) + CP + SP causes shape mismatch at residual connection. #4194

Closed

svcnvidia-nemo-ci added a commit that referenced this pull request Apr 14, 2026

Copy PR #4230 to main

3a1d51e

svcnvidia-nemo-ci mentioned this pull request Apr 14, 2026

[Copy to main] fix(ssm): handle alignment padding in GDN packed seq + CP #4289

Closed

yuzhongw-nvidia mentioned this pull request Apr 14, 2026

[main] feat(moe): Support packed sequence for gated delta net (GDN) #2645

Merged

6 tasks

yuzhongw-nvidia mentioned this pull request Apr 16, 2026

chore: nightly sync main into dev (14_04_2026) #4291

Merged

wuxibin89 mentioned this pull request Jun 2, 2026

[megatron, model] fix: support GDN with packed THD sequence parallel verl-project/verl#6559

Closed

6 tasks

dreamyang-liu mentioned this pull request Jun 26, 2026

GDN: THD packed-sequence support breaks under CP>1 dreamyang-liu/Megatron-LM#1

Open

Uh oh!

Conversation

yxs commented Apr 9, 2026 • edited by yuzhongw-nvidia Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Summary

Test results

Uh oh!

copy-pr-bot Bot commented Apr 9, 2026

Uh oh!

yxs commented Apr 9, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yuzhongw-nvidia commented Apr 10, 2026

Uh oh!

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Uh oh!

yuzhongw-nvidia commented Apr 10, 2026

Uh oh!

yxs commented Apr 11, 2026

Uh oh!

yuzhongw-nvidia commented Apr 13, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 13, 2026

Uh oh!

Uh oh!

yuzhongw-nvidia commented Apr 13, 2026

Uh oh!

Phlip79 commented Apr 14, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

yxs commented Apr 9, 2026 •

edited by yuzhongw-nvidia

Loading