Support mbridge distillation for any_model by danielkorzekwa · Pull Request #904 · NVIDIA/Model-Optimizer

danielkorzekwa · 2026-02-18T18:26:06Z

What does this PR do?

hf_to_mcore mbridge converter (examples for llama and qwen models)
distillation script

- Add distill_anymodel.py: Knowledge distillation script for AnyModel checkpoints - Add import_anymodel_to_mbridge.py: Import script to convert HF AnyModel to MBridge format - Update base.py: Simplify HeterogeneousBridgeMixin for AnyModel support

- Add __init__.py: Module initialization - Add llama.py: Llama bridge implementation - Add qwen3.py: Qwen3 bridge implementation

coderabbitai · 2026-02-18T18:26:18Z

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review (3)

main
release/.*
feature/.*

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment
Commit unit tests in branch dkorzekwa/any_model_mbridge_distillation

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

kevalmorabia97 · 2026-02-20T12:35:56Z

modelopt/torch/puzzletron/export/mbridge/base.py

+        OmegaConf conversion tries to access per_block_parameters which may not
+        be initialized when loading from YAML. Return empty list as fallback.
+        """
+        if name == "per_block_parameters":


Should this also check if that attr is not set then only return [] else return whatever is set?

kevalmorabia97 · 2026-02-20T12:47:10Z

modelopt/torch/puzzletron/export/mbridge/base.py

Do you think at some point we should just upstream this in Megatron-Bridge repo since its a standard Megatron feature without anything to do with model optimization?

pros and cons for both. cons = base.py is NAS related, e.g. no-op in block_configs. However, I think having more generic heterogenous support in mbridge would be useful. I captured it in TODOs.

kevalmorabia97 · 2026-02-20T12:49:02Z

examples/puzzletron/mbridge_distillation/README.md

+export WORKSPACE=/path/to/your/project
+```
+
+1. **Clone Megatron-Bridge:**


Megatron-Bridge is already cloned in the container at /opt/Megatron-Bridge. Why dont we just do following inside the container: cd /opt/Megatron-Bridge && git checkout 960a718cb8989676b258e107d538642717e22e39?

examples/puzzletron/mbridge_distillation/README.md

examples/puzzletron/mbridge_distillation/distill_anymodel.py

…ze() on DistillationProvider.provide() Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

AAnoosheh · 2026-02-24T14:07:34Z

modelopt/torch/puzzletron/export/mbridge/base.py

+
+
+@dataclass
+class GenericHeterogeneousProvider(GPTModelProvider, HeterogeneousTransformerConfig):


UPDATE: the ModelProvider design is now being replaced by ModelConfig which contains a Builder and a TransformerConfig, rather than trying to be both classes at the same time. (So far Mamba changed, but GPT incoming soon)
https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/7cbcf4a4f3f76b1d37cb5395bf2220f8abc44877/src/megatron/bridge/models/mamba/mamba_builder.py#L86

Just heads up that this will need to be refactored in near future.

Good to know. Will need fixes in pruning and distillation example as well. TODOs when 26.04 is out

AAnoosheh · 2026-02-24T14:15:14Z

examples/puzzletron/mbridge_distillation/distill_hf.py

+
+# Patch upstream module BEFORE importing distill() so isinstance checks work with our local DistillationProvider
+# This must happen before distill() is imported because distill.py imports DistillationProvider at module load time
+megatron.bridge.models.distillation_provider.DistillationProvider = DistillationProvider


This is the only significant difference from the existing script, right?

…illation

copy-pr-bot · 2026-02-24T18:08:45Z

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.

…om a student provider class, e.g., GenericHeterogeneousProvider (to serialize heterogeneous_layers_config_encoded_json Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

…ion/convert_checkpoints.py Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

…tion Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

…istillation

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

danielkorzekwa added 2 commits February 18, 2026 10:20

Add missing files from modelopt/torch/puzzletron/export/mbridge

562f46b

- Add __init__.py: Module initialization - Add llama.py: Llama bridge implementation - Add qwen3.py: Qwen3 bridge implementation

danielkorzekwa requested review from a team as code owners February 18, 2026 18:26

danielkorzekwa requested review from realAsma and removed request for a team February 18, 2026 18:26

kevalmorabia97 requested review from AAnoosheh and kevalmorabia97 and removed request for realAsma February 18, 2026 18:29

danielkorzekwa added 3 commits February 19, 2026 01:08

A tutorial on mbridge distillation for puzzletron/any_model

d244ca7

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Update distillation readme

018b208

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Improve mbridge tutorial for anymodel

5f30fa9

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

kevalmorabia97 reviewed Feb 20, 2026

View reviewed changes

danielkorzekwa added 12 commits February 20, 2026 11:03

Fixing distillation for heterogenous models (call self.teacher.finali…

5f73765

…ze() on DistillationProvider.provide() Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Add original keval distill script.

99400f3

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Replace null tokenizer with a teacher tokenizer

f2b65ec

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Ensure exception is printed (and not lost during distributed run)

e385b72

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Add anymodel support to mbridge distillation.

fbcf2a6

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Improve error handling

315e1c0

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Improve mbridge distillation readme

9a98e88

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Update mbridge distillation readme.

d5997cd

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

delete old code

2a611c6

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Make mbrdige distillation to work with nvidian+nemo+26.02.rc5

66b9fdd

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

rename distill_hf_keval.py to distill_hf.py

22cfa6c

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Restore deleted import_anymodel_to_mbridge script (could be useful yet)

4bbdaee

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

AAnoosheh reviewed Feb 24, 2026

View reviewed changes

Merge branch 'feature/compress' into dkorzekwa/any_model_mbridge_dist…

443fb16

…illation

danielkorzekwa requested review from a team as code owners February 24, 2026 18:08

danielkorzekwa requested review from ajrasane, cjluo-nv, h-guo18, kevalmorabia97, mxinO, realAsma, vishalpandya1990 and yueshen2016 and removed request for a team February 24, 2026 18:08

danielkorzekwa added 7 commits February 25, 2026 03:23

Fix serializing a distillation checkpoint to also serialize fields fr…

0b48b69

…om a student provider class, e.g., GenericHeterogeneousProvider (to serialize heterogeneous_layers_config_encoded_json Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Create a script for export_mbridge_to_hf.py based on examples/convers…

291f1b8

…ion/convert_checkpoints.py Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

remove import functionality

d9e6770

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Destroy process group before creating a new one + README for distilla…

6b2afe5

…tion Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Merge branch 'dkorzekwa/any_model' into dkorzekwa/any_model_mbridge_d…

5be52e7

…istillation

Add qwen distillation results to distillation tutorial

b3a01b5

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

code refactoring

a71c507

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>



		@dataclass
		class GenericHeterogeneousProvider(GPTModelProvider, HeterogeneousTransformerConfig):

Conversation

danielkorzekwa commented Feb 18, 2026

What does this PR do?

Uh oh!

coderabbitai bot commented Feb 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Uh oh!

kevalmorabia97 Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

kevalmorabia97 Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

danielkorzekwa Feb 23, 2026

Choose a reason for hiding this comment

Uh oh!

kevalmorabia97 Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

AAnoosheh Feb 24, 2026

Choose a reason for hiding this comment

Uh oh!

kevalmorabia97 Feb 24, 2026

Choose a reason for hiding this comment

Uh oh!

AAnoosheh Feb 24, 2026

Choose a reason for hiding this comment

Uh oh!

copy-pr-bot bot commented Feb 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

coderabbitai bot commented Feb 18, 2026 •

edited

Loading