Support mbridge distillation for any_model#904

Closed

danielkorzekwa wants to merge 136 commits intodkorzekwa/any_modelfrom

dkorzekwa/any_model_mbridge_distillation

Commits on Mar 4, 2026

Add MBridge distillation support for AnyModel checkpoints
danielkorzekwa
committed
Add missing files from modelopt/torch/puzzletron/export/mbridge
danielkorzekwa
committed
A tutorial on mbridge distillation for puzzletron/any_model
danielkorzekwa
committed
Update distillation readme
danielkorzekwa
committed
Improve mbridge tutorial for anymodel
danielkorzekwa
committed
Fixing distillation for heterogenous models (call self.teacher.finalize() on DistillationProvider.provide()
danielkorzekwa
committed
Add original keval distill script.
danielkorzekwa
committed
Replace null tokenizer with a teacher tokenizer
danielkorzekwa
committed
Ensure exception is printed (and not lost during distributed run)
danielkorzekwa
committed
Add anymodel support to mbridge distillation.
danielkorzekwa
committed
Improve error handling
danielkorzekwa
committed
Improve mbridge distillation readme
danielkorzekwa
committed
Update mbridge distillation readme.
danielkorzekwa
committed
delete old code
danielkorzekwa
committed
Make mbrdige distillation to work with nvidian+nemo+26.02.rc5
danielkorzekwa
committed
rename distill_hf_keval.py to distill_hf.py
danielkorzekwa
committed
Restore deleted import_anymodel_to_mbridge script (could be useful yet)
danielkorzekwa
committed
Top-K KL Divergence loss (#747)

AAnoosheh
authored and
danielkorzekwa
committed
[5796745][ONNX][Autocast] Fix opset check for model with custom ops (#767)

gcunhase
authored and
danielkorzekwa
committed
remove duplicated RMSNorm and use LlamaRMSNorm from transformers (#774)

yeyu-nvidia
authored and
danielkorzekwa
committed
[1/2] Address security concerns in code (#626)

kevalmorabia97
authored and
danielkorzekwa
committed
[NVBUG 5801937] Disable dq_only by default (#777)

ajrasane
authored and
danielkorzekwa
committed
Add static per block MSE for NVFP4 weight (#613)

Fridah-nv
authored and
danielkorzekwa
committed
[5763448][ONNX][Autocast] Fix Resize input type mismatch error (#757)

gcunhase
authored and
danielkorzekwa
committed
Fix AWQ export when quantization of some layers are disabled (#721)

meenchen
authored and
danielkorzekwa
committed
Fix Qwen3 recipe and update autoquant example cmd (#749)

meenchen
authored and
danielkorzekwa
committed
chg: passing through trust_remote_code (#778)

ChenhanYu
authored and
danielkorzekwa
committed
Remove quantization_config in config.json from original deepseek models (#753)

Edwardf0t1
authored and
danielkorzekwa
committed
Change trust_remote_code default to False for security reason (#787)

ChenhanYu
authored and
danielkorzekwa
committed
[0.5/3] Diffusion ckpt export for NVFP4 & FP8 (#783)

jingyu-ml
authored and
danielkorzekwa
committed
[5763424][ONNX][Autocast] Fix ConstantOfShape layer output precision (#789)

gcunhase
authored and
danielkorzekwa
committed
[NVBug 5702186] Fix awq model export for Gemma3 (#793)

meenchen
authored and
danielkorzekwa
committed
[5750013][5591945][5360813]: AutoCast standalone implementation for type inference (#719)

galagam
authored and
danielkorzekwa
committed
[5676209] Fix duplicated calib data (#794)

gcunhase
authored and
danielkorzekwa
committed
Define kv cache scaling factor as amax / 448 (#790)

cjluo-nv
authored and
danielkorzekwa
committed
Add Quantizers for Qwen3VLMoeTextDecoderLayer (#666)

authored and
danielkorzekwa
committed
Support KIMI K2 Thinking int4 checkpoint PTQ (#669)

cjluo-nv
authored and
danielkorzekwa
committed
Revert onnxruntime-gpu version to 1.22.0 for Windows (#801)

hthadicherla
authored and
danielkorzekwa
committed
Add Security considerations in docs (#803)

kevalmorabia97
authored and
danielkorzekwa
committed
Add NAS to Minitron pruning for parameter based auto-pruning (#720)

kevalmorabia97
authored and
danielkorzekwa
committed
[1/3] Diffusion ckpt export for NVFP4 & FP8 (#781)

jingyu-ml
authored and
danielkorzekwa
committed
[1/3] Add the fastvideo support (#804)

jingyu-ml
authored and
danielkorzekwa
committed
Svdquant huggingface checkpoint export support (#754)

sychen52
authored and
danielkorzekwa
committed
Fix moe amax remedy for dsr1 and remove global barrier in quantization megatron plugins (#808)

ChenhanYu
authored and
danielkorzekwa
committed
Fix a nvfp4 weight amax attribute issue during export (#785)

Edwardf0t1
authored and
danielkorzekwa
committed
Support megatron generate for vlm (#773)

yueshen2016
authored and
danielkorzekwa
committed
[Minor] Force 'fuse_wgrad_accumulation' to false for TE GroupedLinear (#814)

realAsma
authored and
danielkorzekwa
committed
Feat: Context Parallel for Eagle3 Training (#745)

h-guo18
authored and
danielkorzekwa
committed
Ynankani/update windows benchmark md (#762)

ynankani
authored and
danielkorzekwa
committed
[5676209][ONNX][Autocast] Add support for single `npz` file with multiple samples (#815)

gcunhase
authored and
danielkorzekwa
committed
Support VLM calibration with image-text data (#755)

Edwardf0t1
authored and
danielkorzekwa
committed
[5725362] AutoCast Fixes for models with external data (#731)

galagam
authored and
danielkorzekwa
committed
add FP8 sweep option for static NVFP4 MSE (#758)

Fridah-nv
authored and
danielkorzekwa
committed
Change cnn_dailymail to abisee/cnn_dailymail (#819)

kevalmorabia97
authored and
danielkorzekwa
committed
[5525939] Allow user to select target opset in MOQ (#809)

authored and
danielkorzekwa
committed
Modelopt-windows documentation update (#812)

vishalpandya1990
authored and
danielkorzekwa
committed
Add support for MXFP8 PTQ (#736)

danisereb
authored and
danielkorzekwa
committed
Nenotrom Nano PTQ fix where MoELayer forward has additional named arguments (#823)

ChenhanYu
authored and
danielkorzekwa
committed
Support MLA nvfp4 quant for Deepseek for max perf (#582)

binghanc
authored and
danielkorzekwa
committed
Rename MLM teacher arg (#829)

AAnoosheh
authored and
danielkorzekwa
committed
Context parallelism for Megatron core models (#818)

yeyu-nvidia
authored and
danielkorzekwa
committed
GPTQ Lite implementation (#555)

sugunav14
authored and
danielkorzekwa
committed
Added column-major storage of weights and scales in INT4 quantization for model load time improvement in TRT-RTX (#811)

hthadicherla
authored and
danielkorzekwa
committed
Layerwise KD mode (#802)

AAnoosheh
authored and
danielkorzekwa
committed
Integrate Automated QDQ placement tool - Part 1 (#701)

willg-nv
authored and
danielkorzekwa
committed
Add Megatron-Bridge pruning example scripts (#800)

kevalmorabia97
authored and
danielkorzekwa
committed
Increase nighytly gpu test timeout to 150mins

kevalmorabia97
authored and
danielkorzekwa
committed
Fixes for Megatron Expert Parallel, GroupedMLP and SequentialMLP (#831)

realAsma
authored and
danielkorzekwa
committed
Noeyy/add test cases for the newly added checkpoints on HF (#827)

noeyy-mino
authored and
danielkorzekwa
committed
Update on the QuantModule & DynamicModule to accept external forward (#824)

jingyu-ml
authored and
danielkorzekwa
committed
[2/4] Diffusion Quantized ckpt export (#810)

jingyu-ml
authored and
danielkorzekwa
committed
Latent MOE & Repeated MTP support for NemotronH; fix KV cache quant export (#830)

jenchen13
authored and
danielkorzekwa
committed
Move parallel_state init and warnings to Quant DynamicModule + MBridge pruning doc update (#854)

kevalmorabia97
authored and
danielkorzekwa
committed
GLM-4.7 MTP support (#792)

Edwardf0t1
authored and
danielkorzekwa
committed
Integrate Automated QDQ placement tool - part 2.1 (#844)

willg-nv
authored and
danielkorzekwa
committed
Fix TEGroupedLinear quantization for expert parallelism (EP > 1) (#833)

yueshen2016
authored and
danielkorzekwa
committed
Track global_amax for weight FP4 MSE sweep; Refactor to NVFP4StaticQantizer, NVFP4MSECalibrator (#849)

realAsma
authored and
danielkorzekwa
committed
Fix Sequential MLP amax sync deadlock (#862)

ChenhanYu
authored and
danielkorzekwa
committed
Add contribution guidelines for experimental features (#867)

kaix-nv
authored and
danielkorzekwa
committed
[5868890][ONNX][Autocast] Fix: failure when checking input shape with unknown dimension (#859)

gcunhase
authored and
danielkorzekwa
committed
fix the path change in torch v2.10 for spec dec (#863)

yeyu-nvidia
authored and
danielkorzekwa
committed
Support Qwen3 Next MTP load and export (#860)

authored and
danielkorzekwa
committed
Add Dynamic Memory Sparsification (DMS) training and inference implementation (#877)

authored and
danielkorzekwa
committed
[3.1/4] Diffusion Quantized ckpt export - WAN 2.2 14B (#855)

jingyu-ml
authored and
danielkorzekwa
committed
[fix][5875912] Fix autoquant-autodeploy example (#878)

Fridah-nv
authored and
danielkorzekwa
committed
Add Megatron-Bridge recipe-free distillation example script (#861)

kevalmorabia97
authored and
danielkorzekwa
committed
Integrate Automated QDQ placement tool - part 2.2 (#845)

willg-nv
authored and
danielkorzekwa
committed
Integrate Automated QDQ placement tool - part 2.3 (#846)

authored and
danielkorzekwa
committed
Chenhany/megatron export per layer (#881)

ChenhanYu
authored and
danielkorzekwa
committed
Add Nemotron parse PTQ support (#786)

Edwardf0t1
authored and
danielkorzekwa
committed
MBridge pruning minor fix for saving pruned NemotronH (#887)

kevalmorabia97
authored and
danielkorzekwa
committed
Separate CI job for Megatron GPU tests (#888)

kevalmorabia97
authored and
danielkorzekwa
committed
[OMNIML-3232] Support full TE spec for NemotronH HF-to-Megatron import (#884)

yueshen2016
authored and
danielkorzekwa
committed
[OMNIML-3505] LTX-2 Distillation Trainer (#892)

mxinO
authored and
danielkorzekwa
committed
[fix][5889686] AutoCast: Fix logger (#890)

galagam
authored and
danielkorzekwa
committed
Mamba MOE Quant Configs + Fix Export Bug (#882)

jenchen13
authored and
danielkorzekwa
committed
Support MiniMax M2.1 (FP8 checkpoint) (#817)

cjluo-nv
authored and
danielkorzekwa
committed
[OMNIML-2850] [3/n] Adds sparse attention calibration (#538)

kaix-nv
authored and
danielkorzekwa
committed
Refactor: Eagle data loading (#668)

h-guo18
authored and
danielkorzekwa
committed
Diffusion export bug fixed for model_index.json (#901)

jingyu-ml
authored and
danielkorzekwa
committed
Fix: restore requires_grad in transformers5 reloading (#907)

h-guo18
authored and
danielkorzekwa
committed
[Bug fix] Fake quantized model save after HF accelerate hooks are added (#906)

realAsma
authored and
danielkorzekwa
committed
[NVBUG: 5804406] Auto detect MOE layers (#900)

cjluo-nv
authored and
danielkorzekwa
committed
add local hessian calibration (#788)

Fridah-nv
authored and
danielkorzekwa
committed
Support multiple-batch input for autocast calibration. (#760)

byte-deve
authored and
danielkorzekwa
committed
Fix DeepSeek PTQ script (#912)

cjluo-nv
authored and
danielkorzekwa
committed
Upgrade Dev containers for CICD to latest (#891)

kevalmorabia97
authored and
danielkorzekwa
committed
Remove test_llama_eval_sparse_attention (#914)

kaix-nv
authored and
danielkorzekwa
committed
Sync MOE layer input quantizer only (#903)

jenchen13
authored and
danielkorzekwa
committed
Added support to rotate in fp32 (optional) (#885)

kinjalpatel27
authored and
danielkorzekwa
committed
Improve megatron dataset preprocessing script and update docs (#918)

kevalmorabia97
authored and
danielkorzekwa
committed
Fix test_transformers_tp for torch 2.10 env (#915)

kevalmorabia97
authored and
danielkorzekwa
committed
flush print megatron tokenization stats and update readme (#927)

kevalmorabia97
authored and
danielkorzekwa
committed
SpecDec Bench: February Update (#875)

IzzyPutterman
authored and
danielkorzekwa
committed
Fix: quant config error on quantized offline eagle (#925)

h-guo18
authored and
danielkorzekwa
committed
Fix serializing a distillation checkpoint to also serialize fields from a student provider class, e.g., GenericHeterogeneousProvider (to serialize heterogeneous_layers_config_encoded_json
danielkorzekwa
committed
Create a script for export_mbridge_to_hf.py based on examples/conversion/convert_checkpoints.py
danielkorzekwa
committed
remove import functionality
danielkorzekwa
committed
Destroy process group before creating a new one + README for distillation
danielkorzekwa
committed
Add automodel_distillation example: KD with NeMo AutoModel for AnyModel checkpoints

Separius
authored and
danielkorzekwa
committed
Add qwen distillation results to distillation tutorial
danielkorzekwa
committed
code refactoring
danielkorzekwa
committed
Remove not needed method
danielkorzekwa
committed
clean up docs
danielkorzekwa
committed
Delete not needed file
danielkorzekwa
committed
Integration test for distill_hf
danielkorzekwa
committed
Use strict=False for test_distill_hf.py which uses small models (2 layers) with fewer layers
danielkorzekwa
committed
change mbs in the test from 2 to 1 (to match mbrdige distillation tutorial)
danielkorzekwa
committed
Improve comments
danielkorzekwa
committed
Verify that the distilled model can be loaded in HuggingFace format
danielkorzekwa
committed
Delete not needed import script
danielkorzekwa
committed
Update Dataset Preparation setep
danielkorzekwa
committed
Improve ## Setup section in mbridge distillation readme. Code refactoring to conftest.py
danielkorzekwa
committed
Set the current working dir in a docker container: -w /opt/Model-Optimizer
danielkorzekwa
committed
replace submit_job with srun

Separius
authored and
danielkorzekwa
committed
gpt-oss 20b support (#889)

authored and
danielkorzekwa
committed