[Usage]: Unable to build and run Llama 7B model with AWQ int4 quantization on 4070 and 4080 (OOM)

System Info

System Information:

• OS: Ubuntu 24.04 LTS
• CPU architecture: x86_64
• CPU/Host memory size: 32GB
• GPU name: tested on both [GeForce RTX 4070 Ti SUPER] and [GeForce RTX 4080]
• GPU memory size: 16GB
• TensorRT-LLM version: 0.21.0
• Python versoin: 3.12.7
• NVIDIA Driver Version: 550.54.14
• CUDA Version: 12.4

Detailed output:

(base) fuji@fuji-2404:~/trtllm-benchmark$ nvidia-smi
Fri Aug 15 04:12:37 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.54.14              Driver Version: 550.54.14      CUDA Version: 12.4     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 4070 ...    Off |   00000000:01:00.0 Off |                  N/A |
|  0%   32C    P8              6W /  285W |     295MiB /  16376MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
+-----------------------------------------------------------------------------------------+

(base) fuji@fuji-2404:~/trtllm-benchmark$ python --version
Python 3.12.7
(base) fuji@fuji-2404:~/trtllm-benchmark$ pip show tensorrt_llm
Name: tensorrt_llm
Version: 0.21.0
Summary: TensorRT-LLM: A TensorRT Toolbox for Large Language Models
Home-page: https://github.com/NVIDIA/TensorRT-LLM
Author: NVIDIA Corporation
Author-email:
License: Apache License 2.0
Location: /home/fuji/anaconda3/lib/python3.12/site-packages
Requires: accelerate, aenum, backoff, blake3, build, click, click_option_group, colored, cuda-python, datasets, diffusers, einops, etcd3, evaluate, fastapi, flashinfer-python, h11, h5py, lark, llguidance, matplotlib, meson, mpi4py, mpmath, ninja, numpy, nvidia-cuda-nvrtc-cu12, nvidia-ml-py, nvidia-modelopt, nvidia-nccl-cu12, nvtx, onnx, onnx_graphsurgeon, openai, opencv-python-headless, optimum, ordered-set, pandas, peft, pillow, polygraphy, psutil, pulp, pydantic, pynvml, pyzmq, sentencepiece, setuptools, StrEnum, tensorrt, torch, torchvision, tornado, transformers, triton, uvicorn, wheel, xgrammar
Required-by:

How would you like to use TensorRT-LLM

Relative issue: https://github.com/NVIDIA/TensorRT-LLM/issues/1221 (closed but actually not solved?)

Note: This issue does not occur when using [GeForce RTX 5090] or [GeForce RTX 4090].

Problem description

Got OOM error when running the build engine script described below.

Tested model: meta-llama/Llama-2-7b-chat-hf and meta-llama/Llama-3.1-8B-Instruct

(base) fuji@fuji-2404:~/trtllm-benchmark$ python create_engine.py --model_name meta-llama/Llama-3.1-8B-Instruct --engine_dir /app/tensorrt_llm/llama3_1_8B_awq_int4
...
Traceback (most recent call last):
  File "/home/fuji/trtllm-benchmark/create_engine.py", line 78, in <module>
    main()
  File "/home/fuji/trtllm-benchmark/create_engine.py", line 70, in main
    llm = LLM(model=model_name,
          ^^^^^^^^^^^^^^^^^^^^^
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/tensorrt_llm/llmapi/llm.py", line 911, in __init__
    super().__init__(model, tokenizer, tokenizer_mode, skip_tokenizer_init,
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/tensorrt_llm/llmapi/llm.py", line 827, in __init__
    super().__init__(model, tokenizer, tokenizer_mode, skip_tokenizer_init,
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/tensorrt_llm/llmapi/llm.py", line 202, in __init__
    self._build_model()
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/tensorrt_llm/llmapi/llm.py", line 591, in _build_model
    self._engine_dir, self._hf_model_dir = model_loader()
                                           ^^^^^^^^^^^^^^
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/tensorrt_llm/llmapi/llm_utils.py", line 680, in __call__
    return self._build_model(), self._hf_model_dir
           ^^^^^^^^^^^^^^^^^^^
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/tensorrt_llm/llmapi/llm_utils.py", line 814, in _build_model
    build_task(self.get_engine_dir())
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/tensorrt_llm/llmapi/llm_utils.py", line 768, in build_task
    model_loader(engine_dir=engine_dir)
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/tensorrt_llm/llmapi/llm_utils.py", line 269, in __call__
    pipeline()
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/tensorrt_llm/llmapi/llm_utils.py", line 209, in __call__
    self.step_forward()
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/tensorrt_llm/llmapi/llm_utils.py", line 238, in step_forward
    self.step_handlers[self.counter]()
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/tensorrt_llm/llmapi/llm_utils.py", line 464, in _load_model_from_hf
    model_cls.quantize(
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/tensorrt_llm/models/llama/model.py", line 577, in quantize
    super().quantize(hf_model_dir,
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/tensorrt_llm/models/modeling_utils.py", line 995, in quantize
    quantize_and_export(
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/tensorrt_llm/quantization/quantize_by_modelopt.py", line 817, in quantize_and_export
    export_tensorrt_llm_checkpoint(
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/modelopt/torch/export/model_config_export.py", line 555, in export_tensorrt_llm_checkpoint
    raise e
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/modelopt/torch/export/model_config_export.py", line 490, in export_tensorrt_llm_checkpoint
    ) in torch_to_tensorrt_llm_checkpoint(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/modelopt/torch/export/model_config_export.py", line 411, in torch_to_tensorrt_llm_checkpoint
    merge_qkv(model_config)
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/modelopt/torch/export/model_config_utils.py", line 254, in merge_qkv
    attention.qkv.weight = splitted_qkv.weight
                           ^^^^^^^^^^^^^^^^^^^
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/modelopt/torch/export/model_config.py", line 146, in weight
    return torch.cat((self.q.weight, self.k.weight, self.v.weight))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch.OutOfMemoryError: CUDA out of memory.

create_engine.py

### Generation with Quantization
import logging
import torch
from tensorrt_llm import LLM
from tensorrt_llm.llmapi import QuantAlgo, QuantConfig
import argparse

def main():
    # Parse command-line arguments
    parser = argparse.ArgumentParser(description="Run benchmark script with a specified model name.")
    parser.add_argument("--model_name", required=True, help="Name of the model to use for benchmarking.")
    parser.add_argument("--engine_dir", required=True, help="Path to engine directory.")
    args = parser.parse_args()
    model_name = args.model_name
    engine_dir = args.engine_dir
    quant_and_calib_configs.append(
        (QuantConfig(quant_algo=QuantAlgo.W4A16_AWQ), None))
    quant_config, calib_config  = quant_and_calib_configs[0]
    llm = LLM(model=model_name,
                quant_config=quant_config,
                calib_config=calib_config)
    llm.save(engine_dir)

if __name__ == "__main__":
    main()

I have tried the quantization script, but it also fails with an OOM error.
https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/quantization

Attempt 1:

python3 /home/fuji/TensorRT-LLM/examples/quantization/quantize.py --model_dir meta-llama/Llama-3.1-8B-Instruct --awq_block_size 128 --tp_size 1 --qformat int4_awq --output_dir llama3_1_8B_awq_int4 --calib_size 1

Error:

Same error as above, `torch.OutOfMemoryError: CUDA out of memory.`

Attempt 2 (added --device cpu):

python3 /home/fuji/TensorRT-LLM/examples/quantization/quantize.py --model_dir meta-llama/Llama-3.1-8B-Instruct --awq_block_size 128 --tp_size 1 --qformat int4_awq --output_dir llama3_1_8B_awq_int4 --calib_size 1 --device cpu

Engine built successfully, but got the following error when launching the server

(base) fuji@fuji-2404:~/trtllm-benchmark$ trtllm-serve serve llama3_1_8B_awq_int4
...
Traceback (most recent call last):
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/tensorrt_llm/llmapi/utils.py", line 35, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/tensorrt_llm/executor/worker.py", line 614, in worker_main
    clear_sched_affinity(pid)
  File "/home/fuji/anaconda3/lib/python3.12/site-packages/tensorrt_llm/llmapi/utils.py", line 456, in clear_sched_affinity
    os.sched_setaffinity(pid, set(range(psutil.cpu_count())))
OSError: [Errno 22] Invalid argument
...
    raise RuntimeError("Executor worker died during initialization")
RuntimeError: Executor worker died during initialization

The issue also presents on meta-llama/Llama-2-7b-chat-hf model.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

Start with examples/quantization/quantize.py and the stack-trace paths in tensorrt_llm/llmapi/llm_utils.py and tensorrt_llm/quantization/quantize_by_modelopt.py; reproduce the AWQ int4 build on the listed GPUs, including --device cpu. Also inspect tensorrt_llm/llmapi/utils.py for the serving failure. Done means the documented quantization flow builds and the resulting model starts with trtllm-serve without either reported error.