Llama-3.2-3B-Instruct TensorRT-LLM checkpoint (FP8 weight + FP8 KV)

TensorRT-LLM checkpoint for Llama-3.2-3B-Instruct, with FP8 (W8A8) weight quantization and FP8 KV cache. Use with trtllm-build to produce an engine for inference.

Model details

Item Value
Base model Llama-3.2-3B-Instruct
Framework TensorRT-LLM 1.2.0 (checkpoint format)
Weight quantization FP8 E4M3 (W8A8), per-tensor scale
KV cache FP8
Calibration 512 samples from cnn_dailymail, seq_len=512
Producer TensorRT-Model-Optimizer llm_ptq (nvidia-modelopt 0.41.0) + TensorRT-LLM 1.2.0 convert_checkpoint (--use_fp8 --fp8_kv_cache)
TensorRT 10.14.1
Architecture LlamaForCausalLM (28 layers, hidden_size=3072, GQA 24/8 heads)
Tensor parallelism 1 (single GPU)

Build (how to produce this checkpoint)

FP8 requires a two-step pipeline: (1) run Model Optimizer llm_ptq to calibrate and quantize the HF model to FP8; (2) run TensorRT-LLM convert_checkpoint with the PTQ output.

1. Environment and dependencies

pip install tensorrt_llm==1.2.0 --extra-index-url https://pypi.nvidia.com
# Install TensorRT-Model-Optimizer (required for FP8 calibration)
# See https://github.com/NVIDIA/TensorRT-Model-Optimizer
pip install "nvidia-modelopt[hf]"

2. Quantize base model to FP8 (llm_ptq)

Run Model Optimizer's llm_ptq to produce an FP8-quantized HF-format directory, then convert:

# Step 1: PTQ calibration (512 samples, FP8 weight + FP8 KV)
# This produces saved_models_* with FP8 weights + activation/KV scales
bash TensorRT-Model-Optimizer/examples/llm_ptq/scripts/huggingface_example.sh \
  --model meta-llama/Llama-3.2-3B-Instruct \
  --quant fp8 --kv_cache_quant fp8 \
  --calib 512 --tasks quant

# Step 2: Convert PTQ output to TRT-LLM checkpoint format
# TRT-LLM 1.2.0 path: examples/models/core/llama/
python TensorRT-LLM/examples/models/core/llama/convert_checkpoint.py \
  --model_dir ./path/to/saved_models_* \
  --output_dir ./llama-3.2-3b-instruct-trtllm-ckpt-wq_fp8-kv_fp8 \
  --dtype float16 \
  --tp_size 1 \
  --use_fp8 \
  --fp8_kv_cache \
  --load_model_on_cpu

Important: --use_fp8 is required. Without it, convert_checkpoint sets quant_algo=null and leaves weights as float16. The resulting engine runs float16 weights with FP8 KV cache (scale=1.0), causing KV overflow, NaN logits, and garbage output.

3. Output

After conversion, --output_dir contains config.json and rank0.safetensors; that is the checkpoint in this repo.

How to use

1. Build engine

Requires TensorRT-LLM 1.2.0 and tensorrt_llm installed:

git clone https://huggingface.co/rungalileo/llama-3.2-3b-instruct-trtllm-ckpt-wq_fp8-kv_fp8
cd llama-3.2-3b-instruct-trtllm-ckpt-wq_fp8-kv_fp8

trtllm-build --checkpoint_dir . --output_dir ./engine \
  --max_batch_size 1 --max_input_len 4160 --max_seq_len 4224 \
  --gemm_plugin auto --gpt_attention_plugin auto

2. Run inference

Use a tokenizer from the base model:

trtllm-serve ./engine --tokenizer meta-llama/Llama-3.2-3B-Instruct --port 8000
# OpenAI-compatible API: http://localhost:8000/v1/completions

Checkpoint contents

File Description Size
config.json TRT-LLM model config (architecture + quantization metadata) 2.4 KB
rank0.safetensors Rank 0 weights: 140 FP8 E4M3 weight tensors + 336 FP32 scale tensors + 59 FP16 non-quantized tensors (embeddings, layernorms) 4.5 GB

Quantization details

  • Weight tensors (140): stored as FP8 E4M3 with per-tensor weights_scaling_factor (FP32)
  • Activation scales (140): per-tensor activation_scaling_factor (FP32), calibrated from 512 cnn_dailymail samples
  • KV cache scales (28 layers × 2): kv_cache_scaling_factor + kv_cache_rcp_scaling_factor (FP32)
  • Non-quantized (59): vocab embedding, all layernorms, lm_head remain FP16

References

Downloads last month
8
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for rungalileo/llama-3.2-3b-instruct-trtllm-ckpt-wq_fp8-kv_fp8

Finetuned
(1926)
this model