# GPU memory GPTJ inference

**URL:** <https://discuss.huggingface.co/t/gpu-memory-gptj-inference/43058>\
**Category:** 🤗Transformers\
**Created:** [June 13, 2023, 10:08am UTC](https://discuss.huggingface.co/t/gpu-memory-gptj-inference/43058 "2023-06-13T10:08:31Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![JacopoBandoni](https://avatars.discourse-cdn.com/v4/letter/j/c68b51/32.png) [@JacopoBandoni](https://discuss.huggingface.co/u/JacopoBandoni)\
**Post date:** [June 13, 2023, 10:08am UTC](https://discuss.huggingface.co/t/gpu-memory-gptj-inference/43058/1 "2023-06-13T10:08:31Z")

</div>

I am using a fine tuned GPTJ and I would like to deployed for inference by taking a batch of input\_ids at the same time. Nevertheless I don’t understand why does memory occupancy grow so quickly by increasing the batch size.

```auto
from transformers import AutoModelForCausalLM, AutoTokenizer
import tqdm as notebook_tqdm
import torch
import time

torch.set_grad_enabled(False)

model_load_name = "GPTJgse"
tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v1-6b")
model = AutoModelForCausalLM.from_pretrained(f"{model_load_name}/",torch_dtype=torch.float16, device_map='auto', load_in_8bit=True, low_cpu_mem_usage=True)
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu'); device
print(torch.cuda.memory_reserved())
model.eval()

# model input defined somewhere else
model_input = "some text of 1491 tokens ..."

```

Here I have correctly reserved ‘6721372160’ Byte for the model because it is a 8 bit quantized 6.7 B parameters model.  
When I start inference, with the following code, memory grow very quickly:

```auto
def inference(batch_size):    
    torch.cuda.empty_cache()
    inputs = [model_input]*batch_size
    input_ids = tokenizer(inputs,return_tensors="pt",max_length=2048,)['input_ids']
    output = model.forward(input_ids, return_dict=True, use_cache=True)
    print(torch.cuda.memory_reserved())

inference(1)
inference(2)
inference(3)
inference(4)

```

I got as results:  
8621391872  
10657726464  
12689866752  
14684258304

So for each element in the batch I get about 2GB of GPU memory occupied.  
Nevertheless if I consider GPTJ hyperparameters: d\_head = 256, n\_head = 16, n\_layers = 28.  
I should occupy, assuming 16 bit precision:  
2 (Byte) \* 3 (values: key, value, query) \* 256 \* 16 \* 28 \* 2000 (context length as upper bound) .  
That is less then 1.4 GB of memory.

My goal would be to have greater batch\_size to better scale inference.  
Is there a solution? How could one go about estimating memory requirements?
