> For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server. # Quantization-Aware Training (QAT) in NeMo Automodel NeMo Automodel supports Quantization-Aware Training (QAT) for Supervised Fine-Tuning (SFT) using [TorchAO](https://github.com/pytorch/ao). QAT simulates quantization effects during training so model weights can adapt to the selected lower-precision scheme. Whether QAT improves task quality over post-training quantization (PTQ), and whether the converted model improves deployment efficiency, depends on the model, task, hardware, and inference backend; measure both approaches on the target workload. ## What is Quantization-Aware Training? Quantization-Aware Training introduces fake-quantization operations in the forward pass while keeping model weights trainable. This gives the model an opportunity to adapt to quantization error before the weights are converted to an actual quantized representation. ### Benefits of QAT * **Adaptation during training**: Model weights are optimized while fake quantization is active * **A PTQ alternative**: QAT can be evaluated when PTQ does not meet the workload's quality target * **Deployment experimentation**: Converted models can be benchmarked with quantized kernels on the intended hardware and backend * **Explicit trade-off measurement**: Validation quality, latency, throughput, and memory can be compared against BF16 and PTQ baselines ### QAT vs. Post-Training Quantization | Aspect | QAT | Post-Training Quantization | | ------------------- | ---------------------------------------------------------------- | --------------------------------------------------------- | | **Accuracy** | Model weights can adapt to fake quantization; measure the result | No adaptation step; measure the result | | **Training time** | Longer - requires retraining | None - applied after training | | **Use case** | Consider when PTQ misses a measured quality target | Establish a quantized baseline before adding QAT training | | **Trainable state** | Fine-tunes model weights through fake quantizers | Quantizes fixed model weights after training | ## Requirements To use QAT in NeMo Automodel, you need: * **Software**: Use the repository's locked environment, which includes TorchAO * **Hardware**: An NVIDIA GPU and software stack supported by the selected TorchAO training and deployment kernels * **Training mode**: Full-parameter BF16 SFT; the current recipe rejects QAT with PEFT and rejects models containing non-BF16 parameters * **Model dimensions**: The model's target `nn.Linear` layers must satisfy the selected TorchAO quantizer's group and kernel shape constraints ## Install TorchAO TorchAO is a project dependency. Install the repository's locked environment: ```bash uv sync --locked ``` ## How QAT Works in NeMo Automodel NeMo Automodel integrates TorchAO's QAT quantizers into the training pipeline. During training: 1. **Model preparation**: The quantizer replaces eligible linear layers with TorchAO QAT linear modules 2. **Forward pass**: Eligible weights, and activations for 8da4w, pass through fake quantization 3. **Backward pass**: Gradients flow through the fake quantization operations 4. **Weight updates**: The optimizer updates the model weights; the current prepared modules do not add a separate set of trainable quantization parameters ### Supported Quantization Schemes NeMo Automodel supports two TorchAO QAT quantizers: #### Int8 Dynamic Activation + Int4 Weight (8da4w-qat) * **Quantizer**: `Int8DynActInt4WeightQATQuantizer` * **Activations**: INT8 with dynamic quantization * **Weights**: INT4 quantization * **Layer constraint**: With the current default `padding_allowed=False`, eligible linear layers must have `in_features` divisible by `groupsize`; incompatible layers remain unmodified #### Int4 Weight-Only (4w-qat) * **Quantizer**: `Int4WeightOnlyQATQuantizer` * **Activations**: BF16 in the current integration * **Weights**: INT4 quantization * **Layer constraint**: TorchAO targets bias-free linear layers; preparation requires compatible input dimensions, and conversion requires `out_features` divisible by 8 The current `QATConfig` sets both model-weight precision and weight-scale precision to BF16. For the documented group sizes, a target layer's `in_features` must be divisible by `groupsize`; the 4w path also requires divisibility by `inner_k_tiles * 16` (128 with the default `inner_k_tiles=8`). Layers left unmodified by a quantizer remain BF16, so inspect the prepared model rather than assuming every linear layer was quantized. ## Configuration To enable QAT in your training configuration, you need to specify the quantizer in your YAML configuration file. ### Basic Configuration ```yaml # Enable QAT with Int8 Dynamic Activation + Int4 Weight quantization qat: enabled: true qat_config: _target_: nemo_automodel.components.quantization.qat.QATConfig quantizer_type: int8_dynact_int4weight groupsize: 256 ``` ### Int4 Weight-Only Configuration ```yaml # Enable QAT with Int4 Weight-Only quantization qat: enabled: true qat_config: _target_: nemo_automodel.components.quantization.qat.QATConfig quantizer_type: int4_weight_only groupsize: 256 ``` ### Configuration Parameters | Parameter | Type | Description | | --------------------------- | ---- | ---------------------------------------------------------------------------------------- | | `enabled` | bool | Enable or disable QAT | | `fake_quant_after_n_steps` | int | Requested fake-quant activation step read by the recipe; use `0` in the current revision | | `qat_config._target_` | str | Fully qualified class name of the NeMo AutoModel QAT configuration | | `qat_config.quantizer_type` | str | QAT scheme: `int8_dynact_int4weight` or `int4_weight_only` | | `qat_config.groupsize` | int | Group size for weight quantization (typically 128 or 256) | ### Fake Quantization Start Step The recipe accepts `fake_quant_after_n_steps` as the requested activation threshold. In the current revision, however, the enable and disable helpers are called on top-level model parts without traversing their nested TorchAO QAT linear modules. Positive values therefore do not reliably delay fake quantization in those nested layers. Use `0` and do not rely on delayed activation until that traversal is implemented. ```yaml qat: enabled: true fake_quant_after_n_steps: 0 qat_config: _target_: nemo_automodel.components.quantization.qat.QATConfig quantizer_type: int8_dynact_int4weight groupsize: 256 ``` ## Training Workflow ### 1. Prepare Your Configuration Create a YAML configuration file with QAT enabled: ```yaml recipe: TrainFinetuneRecipeForNextTokenPrediction step_scheduler: global_batch_size: 64 local_batch_size: 8 max_steps: 10000 ckpt_every_steps: 1000 model: _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained pretrained_model_name_or_path: meta-llama/Llama-3.2-1B distributed: strategy: fsdp2 dp_size: none tp_size: 1 cp_size: 1 qat: enabled: true fake_quant_after_n_steps: 0 qat_config: _target_: nemo_automodel.components.quantization.qat.QATConfig quantizer_type: int8_dynact_int4weight groupsize: 256 loss_fn: _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy dataset: _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset dataset_name: rajpurkar/squad split: train dataloader: _target_: torchdata.stateful_dataloader.StatefulDataLoader collate_fn: nemo_automodel.components.datasets.utils.default_collater shuffle: false optimizer: _target_: torch.optim.Adam betas: [0.9, 0.999] eps: 1e-8 lr: 2.0e-5 weight_decay: 0 checkpoint: checkpoint_dir: checkpoints/ model_save_format: safetensors save_consolidated: final ``` ### 2. Run Training Launch training with your QAT-enabled configuration: ```bash uv run automodel your_qat_config.yaml --nproc-per-node 8 ``` ### 3. Monitor Training During training, the model will: * Apply fake quantization to eligible weights and, for 8da4w, activations * Learn to minimize loss while accounting for quantization effects * Produce checkpoints that can be converted to actual quantized models ### 4. Deploy Quantized Model After training, convert the QAT checkpoint to a fully quantized model for deployment: ```python from pathlib import Path import torch from torchao.quantization import Int8DynamicActivationInt4WeightConfig, quantize_ from transformers import AutoModelForCausalLM # Resolve the final checkpoint selected by the recipe. On filesystems without # symlink support, AutoModel writes the relative target to LATEST.txt instead. checkpoint_root = Path("checkpoints") latest_link = checkpoint_root / "LATEST" if latest_link.is_symlink(): checkpoint_dir = latest_link.resolve(strict=True) else: latest_pointer = checkpoint_root / "LATEST.txt" checkpoint_dir = (checkpoint_root / latest_pointer.read_text().strip()).resolve(strict=True) checkpoint_path = checkpoint_dir / "model" / "consolidated" if not checkpoint_path.is_dir(): raise FileNotFoundError(f"Consolidated checkpoint not found: {checkpoint_path}") model = AutoModelForCausalLM.from_pretrained(checkpoint_path, dtype=torch.bfloat16) model.eval() # Apply actual quantization (not fake quantization) quantize_(model, Int8DynamicActivationInt4WeightConfig(group_size=256)) ``` ## Performance Considerations ### Training Performance * **Training time**: Fake-quantization operations add work to each forward pass; benchmark step time against the BF16 baseline * **Memory usage**: QAT modules change the training graph and state; measure peak allocated and reserved memory for the selected model and topology * **Convergence**: Tune training duration and learning rate from validation results rather than assuming BF16 hyperparameters transfer unchanged ### Inference Performance After converting to actual quantization, benchmark the exact inference stack you plan to deploy. The encoded weights use lower precision, but scales, metadata, unmodified layers, kernel availability, and runtime packing all affect the end-to-end result. * **Speed**: Report latency and throughput on the target hardware and backend; quantized storage alone does not guarantee faster kernels * **Memory**: Measure serialized size and runtime peak memory because they include more than packed weight bits * **Accuracy**: Evaluate the same task metrics for BF16, PTQ, and QAT checkpoints ### When to Use QAT Consider QAT when: * **PTQ misses a target**: A measured PTQ quality result is below the workload requirement * **Training is available**: You can run and validate an additional full-parameter SFT job * **The deployment stack supports the scheme**: The target hardware and backend have compatible quantized kernels ### When Not to Use QAT Prefer a simpler baseline when: * **Establishing feasibility**: PTQ can test the deployment scheme without another training run * **The backend lacks compatible kernels**: Conversion would not exercise the intended quantized implementation * **Training resources are limited**: QAT requires another full-parameter training run * **PTQ already meets the measured targets**: Additional QAT training may not be justified ## Best Practices ### 1. Start with Post-Training Quantization Before investing in QAT, try post-training quantization to establish a baseline: ```python import torch from torchao.quantization import Int8DynamicActivationInt4WeightConfig, quantize_ from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained( "meta-llama/Llama-3.2-1B", dtype=torch.bfloat16, ) model.eval() quantize_(model, Int8DynamicActivationInt4WeightConfig(group_size=256)) ``` If accuracy is acceptable, you may not need QAT. ### 2. Choose the Right Quantization Scheme * **8da4w-qat**: Fake-quantizes eligible activations to INT8 and weights to INT4 * **4w-qat**: Fake-quantizes eligible weights to INT4 while activations remain BF16 Train and benchmark both only when their layer-shape constraints and deployment kernels match your model; neither scheme is categorically more accurate or faster for every workload. ### 3. Tune Group Size The `groupsize` parameter affects the granularity of quantization: * **Smaller groups (128)**: Finer weight quantization granularity with more scale metadata * **Larger groups (256)**: Coarser granularity with fewer scales Both values must divide the target linear layer's `in_features`. Compare validation metrics and deployment measurements before choosing one. ### 4. Monitor Validation Metrics Track validation metrics closely during QAT training: * Compare against full-precision baseline * Watch for convergence issues * Compare learning-rate choices using the same validation protocol ### 5. Start Fake Quantization Immediately Use immediate fake quantization in the current revision: ```yaml qat: fake_quant_after_n_steps: 0 ``` Positive values are accepted by the recipe, but the current direct helper calls do not traverse nested TorchAO QAT linear modules, so they must not be relied upon to delay activation. ## Accuracy vs. Efficiency Trade-offs ### Measure the Trade-off Do not assume a fixed accuracy loss, compression ratio, or inference speedup. Record the following for BF16, PTQ, and each QAT scheme you evaluate: | Measurement | What to keep constant | | ---------------------- | ----------------------------------------------------------- | | Task quality | Dataset, preprocessing, decoding, and metric implementation | | Serialized size | Checkpoint contents and serialization format | | Runtime memory | Batch size, sequence length, cache settings, and backend | | Latency and throughput | Hardware, quantized kernels, warm-up, and concurrency | ### Optimization Strategies If accuracy is below expectations: 1. **Tune training duration**: Use validation curves to decide whether more steps help 2. **Tune learning rate**: Compare a small set of learning rates against the BF16 training baseline 3. **Compare schemes**: Evaluate 8da4w and 4w when both satisfy the model and backend constraints 4. **Compare group sizes**: Smaller groups provide finer-grained quantization but add scale metadata ## Limitations and Known Issues ### Current Limitations * **SFT only**: QAT is currently supported for Supervised Fine-Tuning tasks only * **BF16 only**: The infrastructure rejects QAT when any model parameter is not BF16 * **No PEFT**: The fine-tuning recipe rejects configurations that combine QAT and PEFT * **Model compatibility**: TorchAO only prepares eligible `nn.Linear` layers, and the 4w path has stricter bias and shape constraints * **Training overhead**: QAT adds computational overhead during training ### Troubleshooting #### Issue: Training diverges or doesn't converge **Solution**: Try these approaches: * Compare lower learning rates against the BF16 baseline * Compare supported group sizes whose values divide the target layer dimensions * Verify your baseline model trains successfully without QAT #### Issue: Accuracy is significantly worse than expected **Solution**: * Ensure you're comparing against the same baseline (same training steps, data, etc.) * Compare 8da4w quantization with 4w if both are supported by the deployment backend * Compare supported group sizes whose values divide the target layer dimensions * Use the validation curve to decide whether additional training steps help #### Issue: Out of memory during training **Solution**: * Measure peak memory against the BF16 run to determine the added cost for this model * Reduce batch size if needed * Use gradient accumulation to maintain effective batch size ## References * [TorchAO Documentation](https://github.com/pytorch/ao) * [TorchAO QAT Guide](https://github.com/pytorch/ao/tree/main/torchao/quantization/qat) * [Quantization Fundamentals](https://pytorch.org/docs/stable/quantization.html) * [INT8 Quantization for Deep Learning](https://arxiv.org/abs/1806.08342)