Model-Optimizer/examples/llm_qat/llama_factory

Quantization Aware Training and Distillation with LLaMA-Factory

This directory provides integration between LLaMA-Factory and ModelOpt for Quantization-Aware Training (QAT) and Quantization-Aware Distillation (QAD). This enables efficient training of large language models with quantization techniques while maintaining model quality. This README only covers QAT/QAD training with LLaMA-Factory. For more information on setting up other datasets, training different models, or customizing training configurations, please refer to LLaMA-Factory.

Quick Start

Basic QAT/QAD Training

./launch_llamafactory.sh llama_config.yaml

NOTE: The launch_llamafactory.sh script automatically installs LLaMA Factory if it's not already present in your environment.

By default, the script uses fsdp2.yaml for distributed training.

Use Custom FSDP Arguments:

Pass additional FSDP parameters using the FSDP_ARGS environment variable:

FSDP_ARGS="--fsdp_transformer_layer_cls_to_wrap Qwen3DecoderLayer" \
./launch_llamafactory.sh llama_config.yaml

NOTE: The default fsdp2.yaml uses Qwen3DecoderLayer as the transformer layer class. If your model uses a different layer class (e.g., LlamaDecoderLayer for Llama models), pass --fsdp_transformer_layer_cls_to_wrap <your_layer_class> to the launch_llamafactory.sh script or provide a custom FSDP configuration file. See Supported Backends for details.

Custom Config File:

Specify your own FSDP configuration:

./launch_llamafactory.sh llama_config.yaml \
  --accelerate_config /path/to/custom_fsdp.yaml

Training using CLI

For QAT/QAD training using llamafactory_cli, run

./launch_llamafactory.sh train llama_config.yaml

Configuration Guide

This section explains how to configure quantization parameters in your training setup.

YAML Configuration Structure

Your configuration file should follow same structure as sft example from llama3_full_sft.yaml with addition of modelopt configuration:

# Model Configuration
model:
  model_name_or_path: /path/to/your/model
  trust_remote_code: true

### method
stage: sft                    # Supervised Fine-Tuning
do_train: true
finetuning_type: full         # Full parameter fine-tuning

### dataset
dataset: your_dataset
cutoff_len: 4096             # Maximum sequence length
max_samples: 100000          # Maximum training samples

### train
per_device_train_batch_size: 1
gradient_accumulation_steps: 1
learning_rate: 1.0e-5
num_train_epochs: 1
bf16: true                   # Use bfloat16 precision

val_size: 0.1               # Validation split ratio

### ModelOpt Configuration
modelopt:
  recipe: general/ptq/nvfp4_default-kv_fp8  # Quantization recipe (built-in or custom path)
  calib_size: 1024                          # Calibration dataset size
  compress: false                           # Enable weight compression
  distill: false                            # Modify distill to true for QAD
  teacher_model: /path/to/teacher/model     # For QAD (optional)

NOTE: compress: true enables weight compression and will by default use ddp.yaml. NOTE: When training without cli, avoid using deepspeed option in the YAML configuration file.

Deployment

The final QAT/QAD model after training is similar in architecture to that of PTQ model. It simply has updated weights as compared to the PTQ model. It can be deployed to TensorRT-LLM (TRTLLM) or to TensorRT just like a regular ModelOpt PTQ model if the quantization format is supported for deployment.

To run QAT/QAD model with TRTLLM, run:

cd ../../hf_ptq

./scripts/huggingface_example.sh --model <path-to-QAT/QAD-model> --quant nvfp4

See more details on deployment of quantized model here.