Add support for packed thd in BERT language module.

[bbuschkaemper]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do ?

Adds packed thd input support for BERT language module to support sequence packing.

Issue tracking

Linked issue: #5539

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

[copy-pr-bot]

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.

[Copilot]

Pull request overview

This PR adds packed thd (sequence packing) input support to the Megatron Core BertModel, enabling packed-sequence execution paths (no explicit attention mask, packed-aware position ids, and packed-aware RoPE length/params). It also adds a unit test to validate that packed inputs bypass attention-mask expansion and use packed position ids / RoPE invocation.

Changes:

• Add packed_seq_params plumbing to BertModel.forward() and route packed runs to skip extended attention mask creation.
• Generate packed position ids from cu_seqlens_q when packed_seq_params is provided.
• Add a unit test that asserts packed forward uses packed position ids and does not build an extended attention mask.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
megatron/core/models/bert/bert_model.py	Adds packed_seq_params support in forward path, packed position-id generation, and packed-aware RoPE setup.
tests/unit_tests/models/test_bert_model.py	Adds a unit test validating packed-THD forward behavior (position ids, no extended mask, packed params propagation).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Phlip79]

/ok to test ab4aefc

[Phlip79]

/claude review

[claude]

Nice change — the packed thd position-id logic and validation are correct, and the new tests cover the validation paths well.

One non-blocking edge case: in BertModel.forward, the return_embeddings post-process branch computes:

masks = torch.sum(attention_mask, dim=1)

When packed_seq_params is provided, attention_mask is None, so this would raise a TypeError if a caller ever combined packed input with return_embeddings=True and post_process=True. It's an unusual flag combination (and the added tests use post_process=False), so it may be out of scope — but worth a guard or an explicit 'unsupported' error to fail clearly rather than deep in the sum.

[claude]

LGTM

[bbuschkaemper]

@NVIDIA/mcore-oncall

[ericharper]

/claude fix

[svcnvidia-nemo-ci]

❌ Claude fix stopped because a workflow step failed. Inspect the run.

[bbuschkaemper]

@ericharper can you retry? Fixed the merge conflict.

[bbuschkaemper]

@NVIDIA/mcore-oncall

[guihong-nv]

The packed path still needs fixes for padded or context-parallel metadata and for packed-aware post-processing. The inline comments describe the affected cases and expected behavior.

[bbuschkaemper]

@guihong-nv Should be fixed now, can you check?

• position ids are now built from physical padded boundaries. forward() accepts precomputed position_ids and raises error when asking for them with cp>1, covered by test_packed_forward_uses_cu_seqlens_positions_and_no_attention_mask and test_packed_forward_accepts_precomputed_cp_position_ids
• pooler is now indexed with first token of every sequence, covered by test_packed_binary_head_pools_each_sequence
• replaced attention_mask sum with packed-aware aggregation, packed + sp/cp are rejected explicitly, covered by test_packed_return_embeddings_aggregates_each_sequence
• all .item() and torch.any() calls are gone

[ilml]

@guihong-nv ptal, thanks