FastAPI with LLM inference does not release accumulated VRAM
Author: variable
Pull request: no
Labels: name: bug

this is my fastapi snippet

```
model_name = "meta-llama/Llama-3.2-1B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)

@app.post("/unfluff")
async def unfluff_llm(unfluff_request: UnfluffRequest):
    messages = [
        {"role": "system", "content": "Summarise the given service order ticket content, answer should include action performed by helpdesk, don't include the meta info of the ticket such as status, type and priority, don't include when techician is added. Answer directly without any introductory phrases, just answer should include all key points in a concise way"},
        {"role": "user", "content": unfluff_request.text}
    ]
    formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(formatted_prompt, return_tensors="pt").to("cuda")
    with torch.no_grad():
        output = model.generate(**inputs, temperature=0.5, top_p=0.9, max_new_tokens=150)
        # trim the user input from response
        new_tokens = output[0][inputs['input_ids'].shape[-1]:]
        response = tokenizer.decode(new_tokens, skip_special_tokens=True)
        return {'result': response}
```

Initially it started with 3GB (along with other unrelated things), after the LLM endpoint gets called, the VRAM started to climb, but it peaked at about 22GB-ish then OOM, given the video card only has 24GB. If no further calls to the endpoint the VRAM also does not get released.

### Who can help?

_No response_

### 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

keep calling the endpoint that serves the LLM query.

### Expected behavior

Expected the VRAM would not increase overtime.

Comments:
Hi @variable, my guess is that intermediate tensors are not being released here. Try:

1) Totally disabling grad in your application if you're not training, `torch.set_grad_enabled(False)`
2) Running `gc.collect()` to see if tensors are being held in memory even when they should be cleaned up
3) Running this code in a non-async environment to see if you can reproduce the bug

If you keep running it, do you get an OOM error?

Hi @Rocketknight1 

I have run 10000 queries and it is sitting at 21GB, right now and I got OOM

```
  File "/mnt/data/projects/play/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 190, in forward
    down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 4.00 GiB. GPU 0 has a total capacity of 23.53 GiB of which 2.51 GiB is free. Including non-PyTorch memory, this process has 21.00 GiB memory in use. Of the allocated memory 15.10 GiB is allocated by PyTorch, and 5.45 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
```

Will try other stuff you mentioned tonight.

No probs! I've seen other issues from users who are trying to run PyTorch code in async contexts, so my guess is that this shouldn't happen if you disable grad and run this in a normal loop instead of an async coroutine, but let me know how it goes!

Hello @Rocketknight1 

So I have tried, `gc.collect()` and `torch.set_grad_enabled(False)`, even with running in Flask, still hogging memory and not releasing it.

What I have observed is the memory expands when processing larger queries (high water mark) and does not go back down.

Flask code

```python
import gc
import torch
from datasets import load_dataset
from transformers import AutoTokenizer, DataCollatorWithPadding
from transformers import AutoModelForSequenceClassification, Trainer
from transformers import AutoModelForCausalLM
from flask import Flask, request

app = Flask(__name__)

model_name = "meta-llama/Llama-3.2-1B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)

torch.set_grad_enabled(False)

@app.route("/unfluff", methods=['POST'])
def hello_world():
    data = request.get_json()
    messages = [
        {"role": "system", "content": "Summarise the given service order ticket content, answer should include action performed by helpdesk, don't include the meta info of the ticket such as status, type and priority, don't include when techician is added. Answer directly without any introductory phrases, just answer should include all key points in a concise way"},
        {"role": "user", "content": data['text']}
    ]
    formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(formatted_prompt, return_tensors="pt").to("cuda")
    with torch.no_grad():
        output = model.generate(**inputs, temperature=0.5, top_p=0.9, max_new_tokens=150)
        # trim the user input from response
        new_tokens = output[0][inputs['input_ids'].shape[-1]:]
        response = tokenizer.decode(new_tokens, skip_special_tokens=True)
        gc.collect()
        return {'result': response}
```

Good news, after I kept on searching on google, I managed to find someone said need to run `torch.cuda.empty_cache()` and it indeed cleared the memory after heavy query

@variable interesting, and thanks for the diagnosis! This means it's probably related to Torch's memory management, which means there's not much we can do in Transformers, as we operate on a higher level and don't really directly control that. Still, the thread will probably be useful to other people encountering the same issue.

I'm going to close for now since there's not much we can do, but thanks for letting us know your solution!