mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
Shengliang XuandClaude Opus 5
d0142c9dca
### What does this PR do?
**Type of change:** New feature (recipe loading) + one bug fix
Two things, the second built on the first:
1. **A recipe can now reuse another recipe whole.** A top-level
`$import` brings in the imported recipe's entire body; keys given
alongside it override the imported ones. `metadata.recipe_type` becomes
optional and is deprecated along the way.
2. **The deprecated `recipe_type` is swept out of every shipped recipe,
and the checkpoint backfill starts with two published checkpoints
recorded as aliases** that reuse a portable recipe wholesale — the first
users of the alias mechanism — plus a fix to two existing Nemotron NVFP4
recipes.
#### Declaring what kind of recipe a file is
`load_recipe` read `metadata.recipe_type` out of the raw YAML *before*
resolving imports, because it needs the schema class to hand to
`load_config`. That made the field impossible to inherit, so a recipe
reusing another had to restate a line it could only have copied.
It is now optional, and the loader takes the first of these that
answers:
1. a `# modelopt-schema:` comment naming the recipe's schema class,
2. `metadata.recipe_type` — **deprecated**; still read and still
honoured, so a recipe outside this repo keeps working unchanged,
3. the recipe it delegates to via a top-level `$import`.
Whatever a recipe *does* state must be true, in both directions. A
schema comment contradicting a `recipe_type` is rejected, and so is a
recipe importing a different kind of recipe — that used to surface as
whatever pydantic made of, say, an `eagle` section spliced into a PTQ
schema. The concrete recipe classes carry a `RECIPE_TYPE` ClassVar as
the single source of truth.
Only a recipe that another file **imports** needs the schema comment —
that is what `$import` resolution requires to validate the payload. The
sweep here drops `metadata.recipe_type` from all 78 shipped recipes that
carried it and gives the imported ones a `# modelopt-schema:` comment
instead, so nothing in-tree depends on the deprecated field.
A directory recipe's `metadata.yml` resolves its kind the same way —
schema comment first, `recipe_type` as the fallback — it just has no
`$import` to delegate through, since a directory recipe has no body of
its own to hand off. (Follow-up commit, after this PR's initial review:
it originally still required `recipe_type` unconditionally, the one
place the deprecation didn't reach.)
#### Checkpoint aliases
Two checkpoints NVIDIA has published in quantized form use a scheme a
portable recipe already produces, with no checkpoint-specific deviation,
so each is recorded as a thin **alias** (top-level `$import`, overriding
only `metadata`) at its own model-hub path -- the *source* checkpoint's
path, not the published quantized one's:
-
**`models/moonshotai/Kimi-K2.6/ptq/nvfp4_experts_only_mse-kv_fp8_cast`**
delegates to `general/ptq/nvfp4_experts_only_mse-kv_fp8_cast` —
expert-only NVFP4 (MSE static weights, dynamic inputs) with an FP8 KV
cache in cast mode — published as `nvidia/Kimi-K2.6-NVFP4`.
-
**`models/Qwen/Qwen3.5-397B-A17B/ptq/nvfp4_experts_mse-fp8_rest-kv_fp8`**
delegates to the `qwen3_5_moe` architecture recipe
`model_type/qwen3_5_moe/ptq/nvfp4_experts_mse-fp8_rest-kv_fp8` — NVFP4
(MSE static weights) on the routed experts, ModelOpt-default FP8
elsewhere, FP8 KV cache — published as
`nvidia/Qwen3.5-397B-A17B-NVFP4-V2`.
(Follow-up commit, after this PR's initial review: the Qwen entry
originally lived at `models/nvidia/Qwen3.5-397B-A17B/` -- nvidia is the
*published* checkpoint's org, not Qwen3.5-397B-A17B's own. Moved to
match the source model's actual hub path, same as the Kimi-K2.6 entry
above.)
Editing the base recipe changes every alias that points at it; nothing
is duplicated.
#### One fix
- **The Nemotron-3 Super and Ultra NVFP4 recipes** quantized the MTP
block on the **Megatron-Core** path, where it is a live `model.mtp`
submodule their broad `*mixer.*` patterns matched into, contrary to
their own descriptions. They now disable `mtp.*` explicitly. Hugging
Face runs were unaffected — `NemotronHPreTrainedModel` sets
`_keys_to_ignore_on_load_unexpected = [r"mtp.*"]` and builds no MTP
module.
### Usage
A checkpoint alias resolves through `--recipe` to the recipe it
delegates to:
```bash
python examples/hf_ptq/hf_ptq.py \
--pyt_ckpt_path <checkpoint> \
--recipe models/moonshotai/Kimi-K2.6/ptq/nvfp4_experts_only_mse-kv_fp8_cast \
--export_path <output>
```
A recipe that reuses another whole — the shape the aliases use:
```yaml
imports:
base: general/ptq/nvfp4_experts_only_mse-kv_fp8_cast
$import: base
metadata:
description: What this checkpoint uses the base recipe for.
```
### Testing
- **`tests/unit/recipe/test_loader.py`** — 28 new cases covering
whole-recipe reuse with no `metadata` at all; kind resolution from each
of the three sources, from a delegation chain and from a `$import` list;
a delegation cycle failing with `ValueError` rather than recursing;
`peek_declared_schema` including a comment placed below the first YAML
line; `recipe_type` being optional, filled per class, and rejected when
it contradicts; a directory recipe resolving its kind from a schema
comment the same way, rejecting a comment/`recipe_type` disagreement,
and still requiring one or the other; and delegating across kinds being
an error.
- **`tests/unit/recipe/test_recipe_docs.py`** — the
model-specific-recipe check now also covers the two new alias folders,
which must be listed in `ptq.md` like every other
`models/<org>/<model_id>` entry.
- **Recipe validation** (`tools/precommit/check_modelopt_recipes.py`)
and **`pre-commit`** pass on the changed files. The full
`tests/unit/recipe/` suite is left to CI — a broken `transformer_engine`
in the local dev venv keeps the `mtq.quantize`-based cases from running
there.
Not covered: **numerics**. Nothing here asserts accuracy, or that
running one of these recipes reproduces a released checkpoint's weights.
### Before your PR is "*Ready for review*"
- Is this change backward compatible?: ✅ — `metadata.recipe_type` is
still read and honoured for recipes outside this repo, the schema
comments are inert for direct loads, and the loader change only relaxes
a check.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ — no new
dependencies.
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅ — two feature entries, one deprecation, and one bug fix under 0.48.0.
- Did you get Claude approval on this PR?: ❌ — not yet run.
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* Recipes can delegate configurations, support checkpoint aliases, and
apply local metadata overrides.
* Recipe types can be inferred from schema declarations or delegated
recipes, with stronger consistency validation.
* Added unquantized KV-cache options, layerwise export, broader operator
calibration, and new PTQ examples.
* Added checkpoint-specific recipes and MLflow experiment references.
* **Bug Fixes**
* Improved ONNX calibration, FSDP2 export, and fused-MoE quantization
handling.
* Nemotron-3 recipes keep MTP blocks in BF16.
* **Documentation**
* Expanded guidance for aliases, delegation, schema declarations, and
recipe selection.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
91 lines
2.9 KiB
YAML
91 lines
2.9 KiB
YAML
# Domino speculative-decoding training recipe.
|
|
#
|
|
# Domino reuses the DFlash mode/pipeline and adds a causal correction head
|
|
# (GRU + projection), selected via dflash_architecture_config.projector_type=domino.
|
|
# Online training is the default path (data.mode=online). Override fields via an
|
|
# OmegaConf dotlist on the CLI.
|
|
|
|
# modelopt-schema: modelopt.recipe.config.ModelOptDFlashRecipe
|
|
metadata:
|
|
description: Domino training recipe (DFlash backbone + causal correction head).
|
|
|
|
# maps to ModelArguments (main.py)
|
|
model:
|
|
model_name_or_path:
|
|
trust_remote_code: false
|
|
use_fake_base_for_offline: false
|
|
|
|
# maps to DataArguments (main.py)
|
|
data:
|
|
mode: online
|
|
data_path:
|
|
offline_data_path:
|
|
# Jinja chat template with {% generation %} tags for answer_only_loss.
|
|
chat_template:
|
|
|
|
# maps to TrainingArguments (main.py)
|
|
training:
|
|
# --- commonly modified ---
|
|
output_dir:
|
|
num_train_epochs: 6
|
|
per_device_train_batch_size: 1
|
|
learning_rate: 6.0e-4
|
|
warmup_ratio: 0.04
|
|
training_seq_len: 3072
|
|
logging_steps: 50
|
|
save_steps: 2000
|
|
cp_size: 1
|
|
dp_shard_size: 1
|
|
disable_tqdm: true
|
|
# Keep off: eval runs the DFlash backbone only (correction head not applied
|
|
# yet), so AR would reflect the backbone alone, not the trained model.
|
|
estimate_ar: false
|
|
ar_validate_steps: 0
|
|
answer_only_loss: true
|
|
|
|
# --- rarely modified ---
|
|
do_eval: false
|
|
lr_scheduler_type: linear
|
|
save_strategy: steps
|
|
weight_decay: 0.0
|
|
max_grad_norm: 1.0
|
|
dataloader_drop_last: true
|
|
bf16: true
|
|
tf32: true
|
|
remove_unused_columns: false
|
|
# Required: while lambda_base == 1 the head params are unused in the backward
|
|
# graph, so DDP needs unused params allowed.
|
|
ddp_find_unused_parameters: true
|
|
ddp_timeout: 1800
|
|
report_to: tensorboard
|
|
|
|
# maps to DFlashConfig (modelopt/torch/speculative/config.py).
|
|
dflash:
|
|
dflash_block_size: 16
|
|
dflash_num_anchors: 256
|
|
dflash_use_torch_compile: false
|
|
# Domino does not use target-logit KD; it trains its own base/final CE losses.
|
|
dflash_self_logit_distillation: false
|
|
# gamma for exponential loss decay (block_size=16 -> 7).
|
|
dflash_loss_decay_factor: 7.0
|
|
# Qwen3 has no native mask token; 151669 is an unused id used by the reference.
|
|
dflash_mask_token_id: 151669
|
|
# lambda_base curriculum: start fully on base loss, decay to final over all steps.
|
|
dflash_lambda_base_start: 1.0
|
|
dflash_lambda_base_decay_ratio: 1.0
|
|
dflash_architecture_config:
|
|
num_hidden_layers: 5
|
|
# Draft attention/MLP dims. DFlash's draft is an independent model and does
|
|
# NOT auto-inherit these from the base (a fresh Qwen3Config already carries
|
|
# defaults, so modify()'s inherit-if-missing guard is a no-op). Set them
|
|
# explicitly to match the Qwen3-8B reference drafter (GQA: 8 KV heads).
|
|
num_attention_heads: 32
|
|
num_key_value_heads: 8
|
|
head_dim: 128
|
|
intermediate_size: 12288
|
|
projector_type: domino
|
|
emb_dim: 256
|
|
gru_hidden_dim: 1024
|
|
pure_draft_prefix_len: 1
|
|
shift_label: true
|