(Memory leak) trtllm-build gets OOM without GPTAttentionPlugin

[idantene]

System Info

• CPU architecture: x86_64
• CPU/Host memory size: 1TB
• GPU name: NVIDIA A100-SXM4-40GB (up to 8 of these available)
• Libraries
  • tensorrt_llm==0.17.0.dev2024121700
  • tensorrt==10.7.0
  • nvidia-modelopt==0.19.0
  • CUDA: 12.6.3
• OS: Ubuntu 22.04.5 LTS

Who can help?

@kaiyux @byshiue

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Trying to build Llama 3.1 8B as per the docs with:

trtllm-build --checkpoint_dir ./dummy_llama_converted_ckpt --output_dir ./dummy_llama_engine --max_batch_size 1 --max_input_len 1024 --max_seq_len 2048 --kv_cache_type disabled --gpt_attention_plugin disable --context_fmha disable --remove_input_padding disable --log_level verbose --gemm_plugin auto

Expected behavior

The model should compile into a TensorRT engine file without the GPTAttentionPlugin.

actual behavior

Getting out-of-memory (Killed) in a machine with over 1TB of RAM.

additional notes

This seems like an obvious memory leak somewhere. We were not able to identify the cause.

[nv-guomingz]

hi @idantene
could u plz confirm the above cmd is correct?
I got this error msg
build.py: error: argument --kv_cache_type: invalid KVCacheType value: 'disable'

[idantene]

Sorry @nv-guomingz, that was an obvious typo, it should be disabled (I've edited the post above).

[idantene]

Any updates @nv-guomingz?