[shengliangxu]

What does this PR do?

Type of change: Refactor

Reference only, not for merge. This is the full change in one diff, for readers who want the whole picture. It is split for review into the three stacked PRs below; this PR will be closed as they land.

examples/hf_ptq decided how to load, calibrate, preview and export a model by testing its name (model_type == "t5", startswith(("deepseek", "kimi_k2")), "nemotron" in architectures, "vila" in the checkpoint path, ...). That knowledge belongs with the model. This series moves it into modelopt/torch/models, the per-model-type registry from #1828 and #2576-#2580, so the example dispatches through generic per-model data (PTQSpec) and optional hooks and contains no branch on a named model. Adding a model's PTQ facts becomes a registry entry, not an edit to the example. The example still dispatches through the registry; what leaves it is every condition that names a model.

The series

In merge order; each part builds, carries its own tests and passes CI without the later ones.

PR	Title	Added source lines
1/3 #3223	Move the speculative-draft check and the TRT-LLM enc-dec list into library modules	14
2/3 #3224	Add PTQSpec: per-model facts for loading, preview and FSDP2 in the model registry	162
3/3 #3225	Move VLM-wrapper and Whisper calibration behind model hooks	183

By the review guideline's count (lines added to Python source, without blank lines, comments and docstrings; tests excluded; a move git detects as a rename not counted as new), the whole series is 358 lines: 269 in modelopt/ and 89 in examples/hf_ptq. That includes 8 new two-file model packages (the registry's convention is one package per model type); the 78 lines of nemotron_vl_calib.py that move verbatim into the library are not counted.

Base. This branch is on main c785ca9 and has nine commits: the form of the change that was reviewed and that the GPU runs below ran on. The stack is on main 219b34a, which includes #3139 (expert-parallel PTQ from hf_ptq.py). #3139 rewrote the example functions this change edits, so the stack carries the change merged with it (conflicts resolved by keeping #3139's structure and applying the registry and hooks on top), and this branch now conflicts with main in the example. It is not rebased, so the tables below describe the earlier form of the change; each PR of the stack has the results for its own tip.

What moves where

Was in the example	Now
model_type == "t5" / "bart" -> device_map=None	no_device_map on t5, bart
"DeciLM" in architecture -> torch_dtype keyword	dropped: no spec names the DeciLM model types, so those checkpoints load through the dtype keyword (see below)
startswith(("deepseek", "kimi_k2")), == "nemotron_h" -> skip the generation preview	preview_unsupported on deepseek, deepseek_v3, deepseek_v32, deepseek_v4, kimi_k2, kimi_k25 and nemotron_h
not in ("llama4", "llama4_text") -> skip the post-quantization generation	post_ptq_generation_unsupported on llama4, llama4_text
"Deepseek" in architecture -> bundled checkpoint code	the existing prefers_bundled_code hook, through a new deepseek.bundled_code plugin registered for the same DeepSeek types and kimi_k2
is_nemotron_vl ("nemotron" in the architectures) and examples/hf_ptq/nemotron_vl_calib.py	vlm_wrapper and no_auto_device_map on NemotronH_Nano_VL_V2 and nemotron_h_omni, and a forward_calibration_batch hook; the file moves to modelopt/torch/models/nemotron_h_nano_vl_v2/calibration.py (git mv; the calibration forward is unchanged, the hook wrapper around it is new)
model_type == "whisper" in five places (processor, speech calibration data, batch-probe sample shape, preview input key, preview decoding)	hooks load_calibration_processor, make_calibration_dataloader, calibration_sample_input in modelopt/torch/models/whisper/calibration.py; the preview takes input_ids when the batch has them, else input_features
is_diffusion_gemma (sequential multi-GPU map; generate() includes the prompt)	sequential_device_map, generate_includes_prompt on diffusion_gemma
"vila" in args.local_checkpoint_path (FSDP2 refusal)	fsdp2_unsupported on llava_llama (VILA, NVILA, NVILA-Lite; LLaVA-repo fine-tunes share the type and are covered too)
SPECULATIVE_MODEL_LIST (Eagle, Medusa in the architecture)	modelopt.torch.speculative.utils.is_speculative_config (a draft keeps its base model's model_type, so no model-type registry can tell it apart)
is_trtllm_enc_dec_export (t5, mt5, umt5, bart, mbart, whisper)	modelopt/torch/export/trtllm/enc_dec.py, which goes away with that exporter in 0.49.0

PTQSpec is a new section of ModelSpec, read through get_ptq_spec(model_type), which returns an all-default section for a model without one. Its fields are reasons or keywords, so each special case carries its why. hf.py gains four optional hooks (forward_calibration_batch, load_calibration_processor, make_calibration_dataloader, calibration_sample_input), and a model type may now list several plugin modules (two that define the same hook are refused).

Two device-map precedences are kept exactly, which is why there are two fields: no_auto_device_map (Nemotron VL) still yields to an explicit --use_seq_device_map, since users shard these wrappers; no_device_map (T5, BART) yields to nothing.

Model types the series adds to the registry

The registry is keyed by exact model_type, where the example matched by prefix and by architecture substring. It does not enumerate every type those patterns reached: a type gets a spec only when the old code or its tests name it, or, for VILA, when no other key can hold the old check. Each new spec is its own package, the registry's convention.

Model type	Added in	Checkpoints	Why it is added
t5	#3224	T5	The example loaded it without a device map (model_type == "t5"): device_map="auto" can shard its tied encoder/decoder embeddings and position-bias buffers across GPUs, which non-deterministically gives NaN activations in multi-GPU calibration (transformers#21093). Carried as no_device_map.
bart	#3224	BART	The example loaded it without a device map (model_type == "bart"): device_map="auto" and "cuda" raise a meta-tensor error from its safetensors. Carried as no_device_map.
deepseek_v32	#3224	DeepSeek-V3.2	The example skipped the generation preview for every deepseek* type ("DeepSeek generation may go OOM"), and its tests name deepseek_v32. Carried as preview_unsupported. Like the other DeepSeek types, it also maps to the deepseek.bundled_code plugin, which replaces the "Deepseek" in architecture test.
kimi_k2	#3224	Kimi-K2 (remote code; a DeepSeek-V3 architecture under its own type)	The example named it in the same preview skip (startswith(("deepseek", "kimi_k2"))), and its tests name it. Carried as preview_unsupported. It maps to the same bundled-code plugin, since its architecture, DeepseekV3ForCausalLM, matched the old test.
diffusion_gemma	#3224	DiffusionGemma	The example detected it (is_diffusion_gemma) to load it with device_map="sequential" on several GPUs, since the balanced map can split a tied encoder/decoder pair, and to know that its generate() returns the prompt with the output. Carried as sequential_device_map and generate_includes_prompt.
llava_llama	#3224	VILA, VILA1.5, NVILA, NVILA-Lite (remote code); LLaVA-repo fine-tunes share the type	The example refused --use_fsdp2 when the checkpoint path contained "vila" ("custom builder + non-standard layer layout"). VILA has no model type of its own: its checkpoints declare llava_llama, the only key a per-model-type registry can hold that refusal under. Carried as fsdp2_unsupported.
NemotronH_Nano_VL_V2 (snake_case package nemotron_h_nano_vl_v2)	#3225	Nemotron Nano VL v2 (remote code)	The example's Nemotron VL support was written for this wrapper: nemotron_vl_calib.py runs a calibration forward in place of the wrapper's own forward(), which calls torch.distributed.get_rank() unconditionally and assumes past_key_values, and the example's tests name the type. The example found it by "nemotron" in the architecture of any multimodal model; the registry needs the exact type. The spec carries vlm_wrapper and no_auto_device_map, and the package holds the moved forward behind the forward_calibration_batch hook.
whisper	#3225	Whisper	The example tested model_type == "whisper" in five places. The package holds the three hooks that replace them, in whisper/calibration.py. Its spec records only the transformers floor, since every package in the registry has a specs.py (models README).

Types already in the registry gain a ptq_spec: deepseek, deepseek_v3, deepseek_v4, kimi_k25 and nemotron_h (preview_unsupported), llama4 and llama4_text (post_ptq_generation_unsupported), and nemotron_h_omni (the Nemotron VL fields and calibration hook: the old architecture match covered it too, and the moved forward already handles its layout).

Behavior differences (deliberate)

• Models that the old checks reached only through a prefix or a substring, and that no spec names, now take the flow's generic path. The other deepseek* types (deepseek_v2, deepseek_vl*, ...) are not skipped by the generation preview and do not prefer the bundled code. DeciLM checkpoints (deci, nemotron-nas) load through the dtype keyword instead of torch_dtype. A text-only DiffusionGemma type (diffusion_gemma_text) gets the generic device map and preview. Llama-Nemotron-Nano-VL, Nemotron-Parse and plain InternVL wrappers get the generic load, calibration and preview, without the image-text calibration default or the wrapper preview.
• DiffusionGemma is no longer recognised from its architecture alone or from a nested text_config (the real root type is diffusion_gemma).
• The FSDP2 refusal for VILA keys on the llava_llama model type instead of the checkpoint path containing "vila". The type is shared with LLaVA-repo fine-tunes (15 of the 68 llava_llama repos that a hub search for llava or vila finds), which are now refused too, and a renamed VILA checkpoint no longer slips through. The other way round, two kinds of checkpoint that the path test refused are no longer refused: those that declare the model type vila (VILA's four -hf-preview ports and one third-party port, 11 to 21 downloads each), and unrelated models whose path contains "vila" (for example Qwen2 fine-tunes named vilaw-*). This is by design: the registry holds one spec per model type, not one per variant.
• New stdout lines when a spec acts: Loading without a device map: ..., Skipping the generation preview: ..., Using device_map='sequential': .... Log wording: "VLM wrapper" replaces "Nemotron VL"; the FSDP2 MTP refusal says "MTP weights (...)".
• The generation preview takes input_ids when the batch has them, else input_features (it was Whisper by name).

Left out on purpose: dataset_utils.model_type_is_enc_dec, export.model_utils.is_multimodal_model (a Nemotron-Parse name test) and the /no_think prompt in vlm_utils are library code with the same smell, and are not touched here.

Usage

# modelopt/torch/models/<model_type>/specs.py: a model's PTQ facts are a registry entry
from ..specs import ModelSpec, PTQSpec, register

register(
    ModelSpec(
        model_type="my_model",
        min_transformers_version="5.5",
        ptq_spec=PTQSpec(preview_unsupported="its generation runs out of memory"),
    )
)

# modelopt/torch/models/<model_type>/calibration.py: an optional hook, listed in hf._PLUGINS["my_model"]
__all__ = ["forward_calibration_batch"]


def forward_calibration_batch(model, batch) -> bool:
    ...  # run one calibration batch and return True, or return False to leave it to the default

Testing

The CPU parity, focused and torch-only rows are on this branch's tip, rebased onto main c785ca9. The other rows ran on earlier forms of the series, before two rebases onto main and before the last commit dropped six model-type specs. That commit only removes registry entries, a field and the plugin routes of the dropped types, and none of the model-loading runs below used those types (T5, Llama, Mixtral, Whisper, Qwen3-VL).

Check	Result
CPU, wider parity against main: tests/examples/hf_ptq and tests/unit/torch/{models,export,speculative}, transformers 5.5.4 (the floor)	main: 66 failed / 1681 passed / 66 skipped; this branch: 66 failed / 1716 passed / 66 skipped. The 66 are the same ids (no GPU or diffusers in the local environment); the 35 extra passes are the series' new tests
Same, transformers 5.18	main: 66 failed / 1852 passed / 17 skipped; this branch: 66 failed / 1887 passed / 17 skipped; the same 66 ids
CPU, focused set (the example tests, tests/unit/torch/models, the enc-dec export and speculative-config tests)	558 passed, 20 skipped (5.5.4); 700 passed (5.18)
CPU, tests/unit/torch with transformers and the ONNX packages hidden (the selection of the torch-only CI job)	2953 passed, 236 skipped, 0 failed
pre-commit, every hook including mypy and bandit, over main..HEAD	clean
Real checkpoints on CPU through the example's own get_model, processor, dataloader, batch-probe sample, pre_quantize and post_quantize, main vs an earlier form of this branch	t5-small, whisper-tiny (synthetic audio in place of the audio dataset, which needs torchcodec and ffmpeg), a tiny Llama and a tiny Qwen3-VL: identical loaded class, devices, weights digest, preview input and generated text. The only difference is the deliberate new Loading without a device map: ... line for T5
Mutation checks (weak vs strong device-map precedence, the kimi_k2 spec, the encoder-decoder decline)	each caught by at least one test
GPU, tests/examples/hf_ptq/test_llm_ptq.py, main vs this branch, 1 and 2 GPUs, transformers 5.18 (each arm shown to import its own modelopt)	per-id identical at both widths (13 passed at each; the failures are one missing mpi4py in the TRT-LLM deploy step of the test image, identical at main); test_ptq_t5 passes at both
GPU, T5 end to end (test_ptq_t5 alone, stdout shown)	the new Loading without a device map: ... line appears once on this branch and not on main; both take the TRT-LLM enc-dec export; all 10 exported files are byte-identical, including both safetensors
GPU, tests/gpu/torch/{models,export} (registration, unified HF export), 1 GPU, transformers 5.9	per-id identical (9 passed; 3 hit the same fixed 120 s per-test timeout on both)
Not run: Whisper end to end on GPU	the GPU image lacks torchcodec and ffmpeg, which the audio dataset needs; covered on CPU above with synthetic audio
Open: Qwen3-VL, tests/examples/hf_ptq/test_vlm_ptq.py, GPU, transformers 5.9	main: 2 passed in 111 s. First run of this branch: the fp8 case failed (no traceback captured) and the nvfp4 case was still running after 16 minutes. A controlled rerun (export hashes, deploy, first traceback) is pending. Qwen3-VL has no spec in this series and the example's own load, calibration and preview path on a tiny Qwen3-VL is identical to main on CPU, so the PTQ path is unchanged by construction, but that is not yet shown on GPU

Before your PR is "Ready for review"

This PR is never marked ready; the checklist is filled in on each sub-PR. Commits are signed (git commit -s -S).

• Is this change backward compatible?: ✅ (example behavior is unchanged except for the differences listed above; registry additions only)
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: N/A (internal refactor)
• Did you get Claude approval on this PR?: N/A (reference only)

Additional Information

Tracked by OMNIML-6429, for ModelOpt v0.49.0. The earlier reference draft of the same kind was #2566 (closed; superseded by #2576-#2580).

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-3217/
Built to branch gh-pages at 2026-11-04 20:45 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

❌ Patch coverage is 93.19728% with 10 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.34%. Comparing base (c785ca9) to head (1996aa0).

Files with missing lines	Patch %	Lines
modelopt/torch/models/hf.py	75.00%	8 Missing ⚠️
modelopt/torch/models/deepseek/bundled_code.py	66.66%	1 Missing ⚠️
.../torch/models/nemotron_h_nano_vl_v2/calibration.py	87.50%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3217      +/-   ##
==========================================
+ Coverage   81.31%   81.34%   +0.02%     
==========================================
  Files         684      704      +20     
  Lines       79292    79482     +190     
==========================================
+ Hits        64479    64655     +176     
- Misses      14813    14827      +14     

Flag	Coverage Δ
unit	63.72% <93.19%> (+0.08%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.