CUDA OOM for Llama 2 with inferred max batch total token set automatically

[zoltan-fedor]

System Info

Using the latest official text-generation-inference docker image from ghcr.io/huggingface/text-generation-inference.

Trying to deploy the meta-llama/Llama-2-7b-hf model without quantization on a single A10G GPU with manually set token limits:

MAX_INPUT_LENGTH = "4095"
MAX_BATCH_PREFILL_TOKENS = "4096"
MAX_BATCH_TOTAL_TOKENS = "4096"
MAX_TOTAL_TOKENS = "4096"

When it starts up it overwrites the max batch total tokens to 16224 and takes up all of the 22Gb memory of the GPU:

│ 2023-07-19T21:00:08.123161Z  INFO text_generation_launcher: Args { model_id: "meta-llama/Llama-2-7b-hf", revision: None, validation_workers: 2, sharded: Some(false), num_shard: None, quantize: None, dtype: None, trust_remote_code: false, │
│ 2023-07-19T21:00:08.123274Z  INFO download: text_generation_launcher: Starting download process.                                                                                                                                              │
│ 2023-07-19T21:00:09.954549Z  INFO text_generation_launcher: Download file: model-00001-of-00002.safetensors                                                                                                                                   │
│                                                                                                                                                                                                                                               │
│ 2023-07-19T21:00:31.592309Z  INFO text_generation_launcher: Downloaded /data/models--meta-llama--Llama-2-7b-hf/snapshots/b3ab0f18ce41e58ac494e507dad91ef4fd6706c7/model-00001-of-00002.safetensors in 0:00:21.                                │
│                                                                                                                                                                                                                                               │
│ 2023-07-19T21:00:31.592409Z  INFO text_generation_launcher: Download: [1/2] -- ETA: 0:00:21                                                                                                                                                   │
│                                                                                                                                                                                                                                               │
│ 2023-07-19T21:00:31.592781Z  INFO text_generation_launcher: Download file: model-00002-of-00002.safetensors                                                                                                                                   │
│                                                                                                                                                                                                                                               │
│ 2023-07-19T21:00:44.004184Z  INFO text_generation_launcher: Downloaded /data/models--meta-llama--Llama-2-7b-hf/snapshots/b3ab0f18ce41e58ac494e507dad91ef4fd6706c7/model-00002-of-00002.safetensors in 0:00:12.                                │
│                                                                                                                                                                                                                                               │
│ 2023-07-19T21:00:44.004260Z  INFO text_generation_launcher: Download: [2/2] -- ETA: 0                                                                                                                                                         │
│                                                                                                                                                                                                                                               │
│ 2023-07-19T21:00:44.325179Z  INFO download: text_generation_launcher: Successfully downloaded weights.                                                                                                                                        │
│ 2023-07-19T21:00:44.325365Z  INFO shard-manager: text_generation_launcher: Starting shard rank=0                                                                                                                                              │
│ 2023-07-19T21:00:53.004083Z  INFO text_generation_launcher: Server started at unix:///tmp/text-generation-server-0                                                                                                                            │
│                                                                                                                                                                                                                                               │
│ 2023-07-19T21:00:53.034451Z  INFO shard-manager: text_generation_launcher: Shard ready in 8.707555769s rank=0                                                                                                                                 │
│ 2023-07-19T21:00:53.132287Z  INFO text_generation_launcher: Starting Webserver                                                                                                                                                                │
│ 2023-07-19T21:00:53.864803Z  WARN text_generation_router: router/src/main.rs:341: `--revision` is not set                                                                                                                                     │
│ 2023-07-19T21:00:53.864837Z  WARN text_generation_router: router/src/main.rs:342: We strongly advise to set it to a known supported commit.                                                                                                   │
│ 2023-07-19T21:00:54.111041Z  INFO text_generation_router: router/src/main.rs:363: Serving revision b3ab0f18ce41e58ac494e507dad91ef4fd6706c7 of model meta-llama/Llama-2-7b-hf                                                                 │
│ 2023-07-19T21:00:54.115027Z  INFO text_generation_router: router/src/main.rs:210: Warming up model                                                                                                                                            │
│ 2023-07-19T21:00:56.051182Z  WARN text_generation_router: router/src/main.rs:228: `--max-batch-total-tokens` is deprecated for Flash Attention models.                                                                                        │
│ 2023-07-19T21:00:56.051206Z  WARN text_generation_router: router/src/main.rs:232: Inferred max batch total tokens: 16224                                                                                                                      │
│ 2023-07-19T21:00:56.051210Z  INFO text_generation_router: router/src/main.rs:239: Setting max batch total tokens to 16224    

Unfortunately then it crashes with CUDA OOM errors:

:infer:send_error: text_generation_router::infer: router/src/infer.rs:554: Request failed during generation: Server error: CUDA out of memory. Tried to allocate 16.00 MiB (GPU 0; 22.20 GiB total capacity; 20.75 GiB already allocated; 5.12 MiB free; 22.20 GiB allowed; 20.86 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
2023-07-19T21:02:09.363221Z ERROR text_generation_launcher: Method Prefill encountered an error.
Traceback (most recent call last):
  File "/opt/conda/bin/text-generation-server", line 8, in <module>
    sys.exit(app())
  File "/opt/conda/lib/python3.9/site-packages/typer/main.py", line 311, in __call__
    return get_command(self)(*args, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 1130, in __call__
    return self.main(*args, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/typer/core.py", line 778, in main
    return _main(
  File "/opt/conda/lib/python3.9/site-packages/typer/core.py", line 216, in _main
    rv = self.invoke(ctx)
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 1657, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 1404, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 760, in invoke
    return __callback(*args, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/typer/main.py", line 683, in wrapper
    return callback(**use_params)  # type: ignore
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/cli.py", line 78, in serve
    server.serve(
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/server.py", line 169, in serve
    asyncio.run(
  File "/opt/conda/lib/python3.9/asyncio/runners.py", line 44, in run
    return loop.run_until_complete(main)
  File "/opt/conda/lib/python3.9/asyncio/base_events.py", line 634, in run_until_complete
    self.run_forever()
  File "/opt/conda/lib/python3.9/asyncio/base_events.py", line 601, in run_forever
    self._run_once()
  File "/opt/conda/lib/python3.9/asyncio/base_events.py", line 1905, in _run_once
    handle._run()
  File "/opt/conda/lib/python3.9/asyncio/events.py", line 80, in _run
    self._context.run(self._callback, *self._args)
  File "/opt/conda/lib/python3.9/site-packages/grpc_interceptor/server.py", line 159, in invoke_intercept_method
    return await self.intercept(
> File "/opt/conda/lib/python3.9/site-packages/text_generation_server/interceptor.py", line 21, in intercept
    return await response
  File "/opt/conda/lib/python3.9/site-packages/opentelemetry/instrumentation/grpc/_aio_server.py", line 82, in _unary_interceptor
    raise error
  File "/opt/conda/lib/python3.9/site-packages/opentelemetry/instrumentation/grpc/_aio_server.py", line 73, in _unary_interceptor
    return await behavior(request_or_iterator, context)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/server.py", line 71, in Prefill
    generations, next_batch = self.model.generate_token(batch)
  File "/opt/conda/lib/python3.9/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/flash_causal_lm.py", line 823, in generate_token
    raise e
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/flash_causal_lm.py", line 811, in generate_token
    out = self.forward(
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/flash_causal_lm.py", line 787, in forward
    return self.model.forward(
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/custom_modeling/flash_llama_modeling.py", line 471, in forward
    hidden_states = self.model(
  File "/opt/conda/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/custom_modeling/flash_llama_modeling.py", line 430, in forward
    hidden_states, residual = layer(
  File "/opt/conda/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/custom_modeling/flash_llama_modeling.py", line 373, in forward
    mlp_output = self.mlp(normed_attn_res_output)
  File "/opt/conda/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/custom_modeling/flash_llama_modeling.py", line 319, in forward
    return self.down_proj(self.act(gate_up_states[:, 0]) * gate_up_states[:, 1])
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 22.20 GiB total capacity; 20.75 GiB already allocated; 3.12 MiB free; 22.20 GiB allowed; 20.86 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

2023-07-19T21:02:09.363452Z ERROR batch{batch_size=1}:prefill:prefill{id=9 size=1}:prefill{id=9 size=1}: text_generation_client: router/client/src/lib.rs:33: Server error: CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 22.20 GiB total capacity; 20.75 GiB already allocated; 3.12 MiB free; 22.20 GiB allowed; 20.86 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

2023-07-19T21:02:09.364390Z ERROR HTTP request{otel.name=POST / http.client_ip= http.flavor=1.1 http.host=llama-2-7b-hf-service.app-tgi.svc.cluster.local http.method=POST http.route=/ http.scheme=HTTP http.target=/ http.user_agent=python-requests/2.31.0 otel.kind=server trace_id=ca4305abe673242ced42f1111bbf0057}:compat_generate{default_return_full_text=Extension(true)}:generate{parameters=GenerateParameters { best_of: None, temperature: None, repetition_penalty: None, top_k: None, top_p: None, typical_p: None, do_sample: false, max_new_tokens: 300, return_full_text: Some(false), stop: [], truncate: None, watermark: false, details: true, decoder_input_details: false, seed: None }}:generate{request=GenerateRequest { inputs: "Given the context please answer the question, but only if the context is relevant, otherwise return \"I don't know\". Context: Both Carla and Paul lives in Berlin; Question: Who lives in Berlin?; Answer:", parameters: GenerateParameters { best_of: None, temperature: None, repetition_penalty: None, top_k: None, top_p: None, typical_p: None, do_sample: false, max_new_tokens: 300, return_full_text: Some(false), stop: [], truncate: None, watermark: false, details: true, decoder_input_details: false, seed: None } }}:generate_stream{request=GenerateRequest { inputs: "Given the context please answer the question, but only if the context is relevant, otherwise return \"I don't know\". Context: Both Carla and Paul lives in Berlin; Question: Who lives in Berlin?; Answer:", parameters: GenerateParameters { best_of: None, temperature: None, repetition_penalty: None, top_k: None, top_p: None, typical_p: None, do_sample: false, max_new_tokens: 300, return_full_text: Some(false), stop: [], truncate: None, watermark: false, details: true, decoder_input_details: false, seed: None } }}:infer:send_error: text_generation_router::infer: router/src/infer.rs:554: Request failed during generation: Server error: CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 22.20 GiB total capacity; 20.75 GiB already allocated; 3.12 MiB free; 22.20 GiB allowed; 20.86 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

The issue is, that I don't see how I can get around the inferred max batch total token size, which overwrites the token limits I provide.

When using v0.8.0, then it all works (no inferred max batch total tokens being applied, so I assume it uses the numbers I have provided) and uses only 19.9Gb on the GPU.

Information

• Docker
• The CLI directly

Tasks

• An officially supported command
• My own modifications

Reproduction

Steps to reproduce:

1. Use the new llama-2-7b-hf model on an A10G (probably the issue is reproducible on other GPUs too)
2. Set the token limits manually
3. Observe the max batch total token getting automatically overwritten to max out the GPU
4. Observe that requests sent in now will cause CUDA OOM

Expected behavior

I guess either the inferred max batch total token should be set to lower, so there is no CUDA OOM or at least have an option to overwrite it.

[OlivierDehaene]

Do you have another process running on the same gpu by any chance?

[OlivierDehaene]

In #659, we add --cuda-memory-fraction to allow you to set the maximum fraction of VRAM taken by TGI.

[zoltan-fedor]

Thanks @OlivierDehaene ,

"Do you have another process running on the same gpu by any chance?" >> No, that is the only process running on the GPU which is dedicated to that pod (this is happening on EKS kubernetes cluster where each TGI pod gets its own A10G GPU node).

"In #659, we add --cuda-memory-fraction to allow you to set the maximum fraction of VRAM taken by TGI." >> Thanks, that is looking very promising! I will try it out as soon as it becomes available.
For now I just stick with v0.8.0

[shayan1897]

@OlivierDehaene
Still having OOM with latest image and changes on mem_get_info