Skip to content

[codex] Propagate SBHD padding masks across pipeline stages - #5544

Merged
xiaoyao0115 merged 1 commit into
NVIDIA:devfrom
xiaoyao0115:codex/propagate-sbhd-padding-mask
Aug 2, 2026
Merged

[codex] Propagate SBHD padding masks across pipeline stages#5544
xiaoyao0115 merged 1 commit into
NVIDIA:devfrom
xiaoyao0115:codex/propagate-sbhd-padding-mask

Conversation

@xiaoyao0115

Copy link
Copy Markdown
Contributor

What does this PR do?

This splits the SBHD validation padding-mask propagation fix out of #5535.

  • Emits a boolean physical padding_mask from SBHD validation samples, separate from the LM loss mask.
  • Partitions the mask with sequence tensors under context parallelism.
  • Broadcasts the mask through tensor and pipeline parallel stages, including intermediate PP stages.
  • Builds the SBHD validation dataset on every pipeline stage so local MoE layers receive matching padding metadata.

Why?

Prompt tokens can be loss-masked while still being valid model inputs, so the loss mask cannot identify physical padding. Intermediate pipeline stages also need physical padding metadata for their local layers.

Impact

Every pipeline stage receives the correct SBHD physical-padding mask without changing prompt-token participation or LM loss masking.

Validation

  • BASE_REF=dev CHECK_ONLY=true SKIP_DOCS=true bash tools/autoformat.sh
  • tests/unit_tests/data/test_varlen_dataset.py: 69 passed

@copy-pr-bot

copy-pr-bot Bot commented Jun 29, 2026

Copy link
Copy Markdown

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.

@xiaoyao0115
xiaoyao0115 force-pushed the codex/propagate-sbhd-padding-mask branch from 2b06396 to a423e3c Compare June 30, 2026 08:03
@xiaoyao0115

Copy link
Copy Markdown
Contributor Author

/ok to test a423e3c

@xiaoyao0115
xiaoyao0115 force-pushed the codex/propagate-sbhd-padding-mask branch from a423e3c to 0668921 Compare July 27, 2026 07:32
@xiaoyao0115

Copy link
Copy Markdown
Contributor Author

/ok to test 0668921

@xiaoyao0115

Copy link
Copy Markdown
Contributor Author

/ok to test 5332138

@svcnvidia-nemo-ci

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

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

@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Jul 30, 2026
@xiaoyao0115
xiaoyao0115 force-pushed the codex/propagate-sbhd-padding-mask branch from 37681aa to 226bbf8 Compare July 30, 2026 06:21
@xiaoyao0115

Copy link
Copy Markdown
Contributor Author

/ok to test 226bbf8

@xiaoyao0115
xiaoyao0115 enabled auto-merge July 30, 2026 06:25
@xiaoyao0115
xiaoyao0115 added this pull request to the merge queue Jul 30, 2026
@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to a conflict with the base branch Jul 30, 2026
@xiaoyao0115

Copy link
Copy Markdown
Contributor Author

/ok to test f7f921f

@svcnvidia-nemo-ci

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

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

@svcnvidia-nemo-ci

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

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

Signed-off-by: Tailai Ma <58548582+xiaoyao0115@users.noreply.github.com>
@xiaoyao0115

Copy link
Copy Markdown
Contributor Author

/ok to test 78ffb15

@svcnvidia-nemo-ci

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants