Refactor llm_qat: decouple dataset blend size from training runtime params

Dataset configs are now self-contained with blend_size and splits ratios.
train_samples/eval_samples in train configs become runtime caps that don't
invalidate the dataset cache. quantize.py only needs --dataset_config and
--recipe instead of a full train config. Convert int4_blockwise_weight_only
recipe to new list-of-dicts quant_cfg format.

Signed-off-by: realAsma <akuriparambi@nvidia.com>

Author: realAsma

.pre-commit-config.yaml

@@ -101,7 +101,7 @@ repos:
               examples/llm_eval/lm_eval_hf.py|
               examples/llm_eval/mmlu.py|
               examples/llm_eval/modeling.py|
-              examples/llm_qat/main.py|
+              examples/llm_qat/train.py|
               examples/llm_sparsity/weight_sparsity/finetune.py|
               examples/specdec_bench/specdec_bench/models/specbench_medusa.py|
               examples/speculative_decoding/main.py|
@@ -129,6 +129,21 @@ repos:
         args: ["-c", "pyproject.toml", "-q"]
         additional_dependencies: ["bandit[toml]"]
 
+  - repo: local
+    hooks:
+      - id: generate-arguments-md
+        name: Regenerate examples/llm_qat/ARGUMENTS.md
+        entry: bash -c 'python examples/llm_qat/train.py --generate_docs examples/llm_qat/ARGUMENTS.md'
+        language: system
+        files: >-
+          (?x)^(
+            examples/llm_qat/arguments\.py|
+            examples/llm_qat/train\.py|
+            modelopt/torch/opt/plugins/transformers\.py|
+            modelopt/torch/quantization/plugins/transformers_trainer\.py
+          )$
+        pass_filenames: false
+
   - repo: https://github.com/DavidAnson/markdownlint-cli2
     rev: v0.18.1
     hooks:

CHANGELOG.rst

@@ -14,6 +14,9 @@ Changelog
 - Enable PTQ workflow for the Step3.5-Flash MoE model with NVFP4 W4A4 + FP8 KV cache quantization. See `modelopt_recipes/models/Step3.5-Flash/nvfp4-mlp-only.yaml <https://github.com/NVIDIA/Model-Optimizer/blob/main/modelopt_recipes/models/Step3.5-Flash/nvfp4-mlp-only.yaml>`_ for more details.
 - Add support for vLLM fakequant reload using ModelOpt state for HF models. See `examples/vllm_serve/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/vllm_serve#load-qatptq-model-and-serve-in-vllm-wip>`_ for more details.
 - [Early Testing] Add Claude Code PTQ skill (``.claude/skills/ptq/``) for agent-assisted post-training quantization. The skill guides the agent through environment detection, model support checking, format selection, and execution via the launcher or manual SLURM/Docker/bare GPU paths. Includes handling for unlisted models with custom module patching. This feature is in early testing — use with caution.
+- Refactor ``llm_qat`` example with unified YAML-based configuration and flexible dataset blending.
+  ``ModelOptArgParser`` adds ``--config`` YAML support with CLI overrides and auto-generates ``ARGUMENTS.md`` from dataclass definitions.
+  Dataset blending (``configs/dataset/blend.yaml``) supports HuggingFace datasets, local JSON/JSONL/Parquet files, and weighted multi-source blends.
 
 **Backward Breaking Changes**
 

examples/llm_qad/README.md

@@ -2,6 +2,8 @@
 
 Quantization-Aware Distillation (QAD) training scripts for language models using Megatron-LM. These scripts enable training quantized (e.g., NVFP4) student models with knowledge distillation from full-precision teacher models.
 
+> **Note:** For Hugging Face LLM QAD, see the [LLM QAT QAD section](../llm_qat/README.md#end-to-end-qad-example).
+
 ## Overview
 
 | Script | Purpose |

examples/llm_qat/.gitignore

@@ -0,0 +1 @@
+.cache/

examples/llm_qat/ARGUMENTS.md

@@ -0,0 +1,50 @@
+# Argument Reference
+
+_Auto-generated — do not edit by hand._
+
+## DistillArguments
+
+| Argument | Type | Default | Description |
+|----------|------|---------|-------------|
+| `--distill` | `bool` | `False` | Enable training with knowledge distillation. |
+| `--teacher_model` | `str` | `None` | The name or path of the teacher model to use for distillation. |
+| `--criterion` | `str` | `"logits_loss"` | Distillation loss criterion. Currently only 'logits_loss' is supported. |
+
+## DataArguments
+
+| Argument | Type | Default | Description |
+|----------|------|---------|-------------|
+| `--dataset_config` | `str` | `"configs/dataset/blend.yaml"` | Path to a dataset blend YAML config file. |
+| `--train_samples` | `int` | `20000` | Number of training samples to use. |
+| `--eval_samples` | `int` | `2000` | Number of evaluation samples to use. |
+| `--dataset_seed` | `int` | `42` | Random seed for dataset shuffling. |
+| `--dataset_cache_dir` | `str` | `".dataset_cache/tokenized"` | Directory for caching tokenized datasets. |
+| `--shuffle` | `bool` | `True` | Whether to shuffle dataset sources (reservoir sampling). |
+| `--shuffle_buffer` | `int` | `10000` | Buffer size for streaming shuffle. |
+| `--num_proc` | `int` | `16` | Number of CPU workers for tokenization. |
+
+## ModelArguments
+
+| Argument | Type | Default | Description |
+|----------|------|---------|-------------|
+| `--model_name_or_path` | `str` | `"meta-llama/Llama-2-7b-hf"` |  |
+| `--model_max_length` | `int` | `4096` | Maximum sequence length. Sequences will be right padded (and possibly truncated). |
+
+## QuantizeArguments
+
+| Argument | Type | Default | Description |
+|----------|------|---------|-------------|
+| `--recipe` | `str` | `None` | Path to a quantization recipe YAML file (built-in or custom). Built-in recipes can be specified by relative path, e.g. 'general/ptq/nvfp4_default-fp8_kv'. |
+| `--calib_size` | `int` | `512` | Specify the calibration size for quantization. The calibration dataset is used to setup the quantization scale parameters. |
+| `--calib_batch_size` | `int` | `1` | Batch size for calibration data during quantization. |
+| `--compress` | `bool` | `False` | Whether to compress the model weights after quantization for QLoRA. This is useful for reducing the model size. |
+| `--quantize_output_dir` | `str` | `"quantized_model"` | Directory to save the quantized model checkpoint. |
+
+## TrainingArguments
+
+Extends [HuggingFace TrainingArguments](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments). Only additional/overridden arguments are shown below.
+
+| Argument | Type | Default | Description |
+|----------|------|---------|-------------|
+| `--cache_dir` | `str` | `None` |  |
+| `--lora` | `bool` | `False` | Whether to add LoRA (Low-Rank Adaptation) adapter before training. When using real quantization, the LoRA adapter must be set, as quantized weights will be frozen during training. |