feat(model): add masked lm support. by bbuschkaemper · Pull Request #4735 · NVIDIA-NeMo/Megatron-Bridge

bbuschkaemper · 2026-07-08T12:13:51Z

What does this PR do ?

Adds PreTrainedMaskedLM, a lightweight HF-pretrained wrapper for *ForMaskedLM models alongside the existing PreTrainedCausalLM.

Changelog

hf_pretrained/masked_lm.py: new PreTrainedMaskedLM wrapper class
conversion/auto_bridge.py: _resolve_pretrained_wrapper_cls dispatch logic to select PreTrainedMaskedLM vs. PreTrainedCausalLM by architecture
Unit tests for the wrapper and dispatch logic

GitHub Actions CI

See the CI section in the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation?
Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
- Reviewer: Does the PR have correct import guards for all optional libraries?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

Related to # (issue)

Signed-off-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com>

copy-pr-bot · 2026-07-08T12:13:54Z

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 extends the Hugging Face “hf_pretrained” wrapper layer and AutoBridge loading path to recognize encoder-only masked language model architectures (*ForMaskedLM) by introducing a dedicated PreTrainedMaskedLM wrapper and dispatching to it from AutoBridge.from_hf_pretrained.

Changes:

Add PreTrainedMaskedLM (lazy-loading wrapper using AutoModelForMaskedLM with AutoModel fallback).
Update AutoBridge to allowlist *ForMaskedLM and select the appropriate pretrained wrapper class during load.
Add unit tests covering wrapper behavior and AutoBridge dispatch/support semantics.

Reviewed changes

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

Show a summary per file

File	Description
tests/unit_tests/models/test_auto_bridge.py	Adds coverage for masked-LM architecture allowlisting and wrapper dispatch behavior.
tests/unit_tests/models/hf_pretrained/test_masked_lm.py	New unit tests for `PreTrainedMaskedLM` lazy-loading, device helpers, and basic API surface.
src/megatron/bridge/models/hf_pretrained/README.md	Documents the new `PreTrainedMaskedLM` wrapper usage.
src/megatron/bridge/models/hf_pretrained/masked_lm.py	Implements `PreTrainedMaskedLM` wrapper (config/tokenizer/model lazy loading, encode/decode, repr, save).
src/megatron/bridge/models/hf_pretrained/init.py	Exports `PreTrainedMaskedLM` from the package namespace.
src/megatron/bridge/models/conversion/auto_bridge.py	Adds `_resolve_pretrained_wrapper_cls` and routes HF loading/isinstance checks through a wrapper-type allowlist.

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

Signed-off-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com>

yaoyu-33

Overall direction looks good. I found two correctness/contract edges worth addressing inline.

yaoyu-33 · 2026-07-08T22:42:13Z

+
+        try:
+            model = AutoModelForMaskedLM.from_pretrained(self.model_name_or_path, **model_kwargs)
+        except ValueError:


Could we avoid treating every ValueError as proof that no masked-LM head exists? from_pretrained can also raise ValueError for invalid loader/config options or failures in a custom MLM implementation; retrying through AutoModel can then silently return a base encoder without the MLM head. Please decide the fallback from the config/model mapping (or narrow this to the unsupported-config case) and re-raise unrelated errors, with a regression test.

Done in bc62f43

yaoyu-33 · 2026-07-08T22:42:13Z

+        ) as mock_safe_load_config:
+            mock_safe_load_config.return_value = bert_masked_lm_config
+
+            assert AutoBridge.can_handle("bert-base-uncased") is True


This assertion conflicts with test_from_hf_pretrained_with_masked_lm_architecture_and_no_registered_bridge: BERT has no registered MegatronModelBridge, so can_handle() returns True even though the subsequent load is guaranteed to fail during validation. Since can_handle() is documented as a compatibility preflight, could we include the registry check here (for example by reusing _validate_config) or keep this False until a BERT bridge is registered?

Done in 9f54794. We can update the test assertion in #4737 once this PR lands.

Signed-off-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com>

yaoyu-33 · 2026-07-09T20:15:42Z

/ok to test bc62f43

huvunvidia · 2026-07-10T03:14:32Z

Hi @bbuschkaemper ,
All the CICD tests passed but it requires tests coverage (codecov/patch).
Can you please add more appropriate tests?

Signed-off-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com>

Signed-off-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com> # Conflicts: # src/megatron/bridge/models/conversion/auto_bridge.py

bbuschkaemper · 2026-07-10T09:47:06Z

Hi @bbuschkaemper , All the CICD tests passed but it requires tests coverage (codecov/patch). Can you please add more appropriate tests?

Done. Also merged main into branch since new commits on main created merge conflicts. Could you rerun CICD?

huvunvidia · 2026-07-10T14:01:59Z

/ok to test 9cdb1d3

huvunvidia · 2026-07-14T03:28:53Z

Hi @bbuschkaemper ,
It seems like there's a small conflict for merging. Can you have a look?
I can trigger the test right after.

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

yaoyu-33 · 2026-07-14T04:42:26Z

/ok to test e490712

Signed-off-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com> Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com> Co-authored-by: yaoyu-33 <yaoyu.094@gmail.com>

feat(model): add masked lm support.

df4a13e

Signed-off-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com>

Copilot AI review requested due to automatic review settings July 8, 2026 12:13

github-actions Bot added the community-request label Jul 8, 2026

Copilot started reviewing on behalf of bbuschkaemper July 8, 2026 12:14 View session

Copilot AI reviewed Jul 8, 2026

View reviewed changes

Comment thread src/megatron/bridge/models/hf_pretrained/README.md

Comment thread src/megatron/bridge/models/hf_pretrained/masked_lm.py

bbuschkaemper mentioned this pull request Jul 8, 2026

feat(model): add bert support. #4737

Open

5 tasks

fix(model): fix invalid config check and docs.

f637ffe

Signed-off-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com>

yaoyu-33 added area:model Model implementations and HF bridge logic feature New capabilities, enhancements, or enablement work needs-review PR is ready for code review and waiting on a reviewer labels Jul 8, 2026

yaoyu-33 reviewed Jul 8, 2026

View reviewed changes

bbuschkaemper added 2 commits July 9, 2026 09:45

fix(model): correctly check for valid config in auto bridge can_handle.

9f54794

Signed-off-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com>

fix(model): correctly handle ValueError in masked lm.

bc62f43

Signed-off-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com>

yaoyu-33 added ready-to-merge PR is approved, current, and only waiting for CI to pass before merge and removed needs-review PR is ready for code review and waiting on a reviewer labels Jul 9, 2026

copy-pr-bot Bot temporarily deployed to public July 9, 2026 20:29 Inactive

copy-pr-bot Bot temporarily deployed to test July 9, 2026 20:35 Inactive

copy-pr-bot Bot temporarily deployed to public July 9, 2026 21:03 Inactive

copy-pr-bot Bot temporarily deployed to public July 9, 2026 21:07 Inactive

copy-pr-bot Bot temporarily deployed to public July 9, 2026 21:09 Inactive

copy-pr-bot Bot temporarily deployed to public July 9, 2026 21:38 Inactive

bbuschkaemper added 2 commits July 10, 2026 09:39

test(model): increase test coverage for masked lm.

b0351dc

Signed-off-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com>

Merge branch 'main' into feat-add-masked-lm-wrapper

9cdb1d3

Signed-off-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com> # Conflicts: # src/megatron/bridge/models/conversion/auto_bridge.py

copy-pr-bot Bot temporarily deployed to public July 10, 2026 14:02 Inactive

copy-pr-bot Bot temporarily deployed to test July 10, 2026 14:03 Inactive

copy-pr-bot Bot temporarily deployed to public July 10, 2026 14:45 Inactive

copy-pr-bot Bot temporarily deployed to public July 10, 2026 15:13 Inactive

huvunvidia requested a review from yaoyu-33 July 13, 2026 13:41

yaoyu-33 previously approved these changes Jul 14, 2026

View reviewed changes

Merge main into feat-add-masked-lm-wrapper

e490712

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

yaoyu-33 dismissed their stale review via e490712 July 14, 2026 04:42

copy-pr-bot Bot temporarily deployed to public July 14, 2026 04:43 Inactive

copy-pr-bot Bot temporarily deployed to test July 14, 2026 04:43 Inactive

yaoyu-33 merged commit d9b63a9 into NVIDIA-NeMo:main Jul 14, 2026
15 checks passed

bbuschkaemper deleted the feat-add-masked-lm-wrapper branch July 14, 2026 07:29

copy-pr-bot Bot temporarily deployed to public July 14, 2026 09:23 Inactive

copy-pr-bot Bot temporarily deployed to public July 14, 2026 09:24 Inactive

copy-pr-bot Bot temporarily deployed to public July 14, 2026 09:52 Inactive

cuichenx mentioned this pull request Jul 28, 2026

[NeMo FW 26.08 Release] MBridge 0.6.0 Roadmap #5144

Open

Uh oh!

Conversation

bbuschkaemper commented Jul 8, 2026

What does this PR do ?

Changelog

GitHub Actions CI

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented Jul 8, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

yaoyu-33 left a comment

Choose a reason for hiding this comment

Uh oh!

yaoyu-33 Jul 8, 2026

Choose a reason for hiding this comment

Uh oh!

bbuschkaemper Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

yaoyu-33 Jul 8, 2026

Choose a reason for hiding this comment

Uh oh!

bbuschkaemper Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

yaoyu-33 commented Jul 9, 2026

Uh oh!

huvunvidia commented Jul 10, 2026

Uh oh!

bbuschkaemper commented Jul 10, 2026

Uh oh!

huvunvidia commented Jul 10, 2026

Uh oh!

huvunvidia commented Jul 14, 2026

Uh oh!

yaoyu-33 commented Jul 14, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants