Quantization-Aware Training (QAT) in NeMo Automodel

NeMo Automodel supports Quantization-Aware Training (QAT) for Supervised Fine-Tuning (SFT) using TorchAO. 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

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:

$
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

1
# Enable QAT with Int8 Dynamic Activation + Int4 Weight quantization
2
qat:
3
  enabled: true
4
  qat_config:
5
    _target_: nemo_automodel.components.quantization.qat.QATConfig
6
    quantizer_type: int8_dynact_int4weight
7
    groupsize: 256

Int4 Weight-Only Configuration

1
# Enable QAT with Int4 Weight-Only quantization
2
qat:
3
  enabled: true
4
  qat_config:
5
    _target_: nemo_automodel.components.quantization.qat.QATConfig
6
    quantizer_type: int4_weight_only
7
    groupsize: 256

Configuration Parameters

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.

1
qat:
2
  enabled: true
3
  fake_quant_after_n_steps: 0
4
  qat_config:
5
    _target_: nemo_automodel.components.quantization.qat.QATConfig
6
    quantizer_type: int8_dynact_int4weight
7
    groupsize: 256

Training Workflow

1. Prepare Your Configuration

Create a YAML configuration file with QAT enabled:

1
recipe: TrainFinetuneRecipeForNextTokenPrediction
2
3
step_scheduler:
4
  global_batch_size: 64
5
  local_batch_size: 8
6
  max_steps: 10000
7
  ckpt_every_steps: 1000
8
9
model:
10
  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
11
  pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
12
13
distributed:
14
  strategy: fsdp2
15
  dp_size: none
16
  tp_size: 1
17
  cp_size: 1
18
19
qat:
20
  enabled: true
21
  fake_quant_after_n_steps: 0
22
  qat_config:
23
    _target_: nemo_automodel.components.quantization.qat.QATConfig
24
    quantizer_type: int8_dynact_int4weight
25
    groupsize: 256
26
27
loss_fn:
28
  _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy
29
30
dataset:
31
  _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset
32
  dataset_name: rajpurkar/squad
33
  split: train
34
35
dataloader:
36
  _target_: torchdata.stateful_dataloader.StatefulDataLoader
37
  collate_fn: nemo_automodel.components.datasets.utils.default_collater
38
  shuffle: false
39
40
optimizer:
41
  _target_: torch.optim.Adam
42
  betas: [0.9, 0.999]
43
  eps: 1e-8
44
  lr: 2.0e-5
45
  weight_decay: 0
46
47
checkpoint:
48
  checkpoint_dir: checkpoints/
49
  model_save_format: safetensors
50
  save_consolidated: final

2. Run Training

Launch training with your QAT-enabled configuration:

$
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:

1
from pathlib import Path
2
3
import torch
4
5
from torchao.quantization import Int8DynamicActivationInt4WeightConfig, quantize_
6
from transformers import AutoModelForCausalLM
7
8
# Resolve the final checkpoint selected by the recipe. On filesystems without
9
# symlink support, AutoModel writes the relative target to LATEST.txt instead.
10
checkpoint_root = Path("checkpoints")
11
latest_link = checkpoint_root / "LATEST"
12
if latest_link.is_symlink():
13
    checkpoint_dir = latest_link.resolve(strict=True)
14
else:
15
    latest_pointer = checkpoint_root / "LATEST.txt"
16
    checkpoint_dir = (checkpoint_root / latest_pointer.read_text().strip()).resolve(strict=True)
17
18
checkpoint_path = checkpoint_dir / "model" / "consolidated"
19
if not checkpoint_path.is_dir():
20
    raise FileNotFoundError(f"Consolidated checkpoint not found: {checkpoint_path}")
21
22
model = AutoModelForCausalLM.from_pretrained(checkpoint_path, dtype=torch.bfloat16)
23
model.eval()
24
25
# Apply actual quantization (not fake quantization)
26
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:

1
import torch
2
3
from torchao.quantization import Int8DynamicActivationInt4WeightConfig, quantize_
4
from transformers import AutoModelForCausalLM
5
6
model = AutoModelForCausalLM.from_pretrained(
7
    "meta-llama/Llama-3.2-1B",
8
    dtype=torch.bfloat16,
9
)
10
model.eval()
11
12
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:

1
qat:
2
  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:

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
• TorchAO QAT Guide
• Quantization Fundamentals
• INT8 Quantization for Deep Learning