# Albert Pre-training with Batch size 8 is throwing OOM

**URL:** <https://discuss.huggingface.co/t/albert-pre-training-with-batch-size-8-is-throwing-oom/13598>\
**Category:** 🤗Transformers\
**Created:** [January 12, 2022, 7:17am UTC](https://discuss.huggingface.co/t/albert-pre-training-with-batch-size-8-is-throwing-oom/13598 "2022-01-12T07:17:17Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![Kushwanth](https://avatars.discourse-cdn.com/v4/letter/k/278dde/32.png) [@Kushwanth](https://discuss.huggingface.co/u/Kushwanth)\
**Post date:** [January 12, 2022, 7:17am UTC](https://discuss.huggingface.co/t/albert-pre-training-with-batch-size-8-is-throwing-oom/13598/1 "2022-01-12T07:17:18Z")

</div>

## Environment info

- `transformers` version: 4.15.0
- Platform: Ubuntu 18.1
- Python version: 3.7
- PyTorch version (GPU?): 1.6+cuda 10.1
- Tensorflow version (GPU?): None
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: yes

### Who can help

The tasks I am working on is:  
Masked Language Modelling (using AlbertForMaskedLM)  
Training Albert model from scratch

python run\_mlm.py --model\_type albert --num\_train\_epochs 300 --train\_file /home/kushwanth/write\_chunks/sample.txt --validation\_file /home/kushwanth/write\_chunks/sample.txt --tokenizer\_name albert --do\_train=yes --output\_dir=/home/kushwanth/model --per\_device\_train\_batch\_size 8 --per\_device\_eval\_batch\_size 8 --save\_steps 3000 --logging\_steps 500 --report\_to tensorboard --preprocessing\_num\_workers 10 2\>&1 | tee /home/kushwanth/log\_1.txt

GPU spects:  
V-100  
Number of GPUs 8 with 16GB RAM each

Trying to allocate batch size of 8 and we are getting Out of memory error and GPU utilisation is 50% on avg

Model Config:  
“attention\_probs\_dropout\_prob”: 0,  
“bos\_token\_id”: 2,  
“classifier\_dropout\_prob”: 0.1,  
“embedding\_size”: 128,  
“eos\_token\_id”: 3,  
“hidden\_act”: “gelu\_new”,  
“hidden\_dropout\_prob”: 0,  
“hidden\_size”: 768,  
“initializer\_range”: 0.02,  
“inner\_group\_num”: 1,  
“intermediate\_size”: 3072,  
“layer\_norm\_eps”: 1e-12,  
“max\_position\_embeddings”: 512,  
“model\_type”: “albert”,  
“num\_attention\_heads”: 12,  
“num\_hidden\_groups”: 1,  
“num\_hidden\_layers”: 12,  
“pad\_token\_id”: 0,  
“position\_embedding\_type”: “absolute”,  
“torch\_dtype”: “float32”,  
“transformers\_version”: “4.15.0”,  
“type\_vocab\_size”: 1,  
“vocab\_size”: 40000

```auto
File "run_mlm.py", line 442, in main                                                                                      
    train_result = trainer.train(resume_from_checkpoint=checkpoint)                                                               
  File "/home/kushwanth/anaconda3/envs/py37/lib/python3.7/site-packages/transformers/trainer.py", line 1332, in train       
    tr_loss_step = self.training_step(model, inputs)                                                                              
  File "/home/kushwanth/anaconda3/envs/py37/lib/python3.7/site-packages/transformers/trainer.py", line 1891, in training_step loss = self.compute_loss(model, inputs)                                                                                       
  File "/home/kushwanth/anaconda3/envs/py37/lib/python3.7/site-packages/transformers/trainer.py", line 1923, in compute_loss      
    outputs = model(**inputs)                                                                                                     
  File "/home/kushwanth/anaconda3/envs/py37/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl result = self.forward(*input, **kwargs)                                                                                       
  File "/home/kushwanth/anaconda3/envs/py37/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py", line 156, in forward 
    return self.gather(outputs, self.output_device)                                                                               
  File "/home/kushwanth/anaconda3/envs/py37/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py", line 168, in gather  
    return gather(outputs, output_device, dim=self.dim)
  File "/home/kushwanth/anaconda3/envs/py37/lib/python3.7/site-packages/torch/nn/parallel/scatter_gather.py", line 68, in gather
    res = gather_map(outputs)
  File "/home/kushwanth/anaconda3/envs/py37/lib/python3.7/site-packages/torch/nn/parallel/scatter_gather.py", line 62, in gather_m
ap
    for k in out))
  File "<string>", line 7, in __init__
  File "/home/kushwanth/anaconda3/envs/py37/lib/python3.7/site-packages/transformers/file_utils.py", line 2294, in __post_init__
    for element in iterator:
  File "/home/kushwanth/anaconda3/envs/py37/lib/python3.7/site-packages/torch/nn/parallel/scatter_gather.py", line 62, in <genexpr>
    for k in out))
  File "/home/kushwanth/anaconda3/envs/py37/lib/python3.7/site-packages/torch/nn/parallel/scatter_gather.py", line 55, in gather_m
ap
    return Gather.apply(target_device, dim, *outputs)
  File "/home/kushwanth/anaconda3/envs/py37/lib/python3.7/site-packages/torch/nn/parallel/_functions.py", line 68, in forward
    return comm.gather(inputs, ctx.dim, ctx.target_device)
  File "/home/kushwanth/anaconda3/envs/py37/lib/python3.7/site-packages/torch/cuda/comm.py", line 166, in gather
    return torch._C._gather(tensors, dim, destination)
RuntimeError: CUDA out of memory. Tried to allocate 4.88 GiB (GPU 0; 15.78 GiB total capacity; 7.63 GiB already allocated; 2.08 Gi
B free; 12.50 GiB reserved in total by PyTorch)

```