In Huggingface transformers, resuming training with the same parameters as before fails with a CUDA out of memory error

username: YISTANFORD
cooked: <p>Hello,</p>
<p>I am using my university’s HPC cluster and there is a time limit per job. So I ran the train method of the Trainer class with resume_from_checkpoint=MODEL and resumed the training. The following is the code for resuming. To prevent CUDA out of memory errors, we set param.requires_grad = False in the model as before resuming.</p>
<pre><code class="lang-auto">training_args = TrainingArguments(
    logging_steps=500,
    save_steps=500,
    eval_steps=500,
    output_dir='/home/groups/group_a/yuspqr/result/'+NAME+'.global',
    num_train_epochs=3,              # total # of training epochs
    adam_epsilon=1e-6,
    per_device_train_batch_size=REAL_BS,  # batch size per device during training
    per_device_eval_batch_size=REAL_BS,   # batch size for evaluation
    warmup_steps=500,                # number of warmup steps for learning rate scheduler
    weight_decay=0.01,               # strength of weight decay
    gradient_accumulation_steps=ACCUM_NUM, # directory for storing logs
    #gradient_checkpointing=True,
    dataloader_num_workers=4,
    fp16=True,
    learning_rate=RL,
    do_train=True,
    do_eval=True,
    logging_dir='/home/users/yuspqr/runs/'+NAME+'/',
    report_to='tensorboard',
    evaluation_strategy='steps',
    logging_strategy='steps',
    save_strategy='steps',
)

import glob
def latest_checkpoint(path):
    cl=glob.glob(path)
    max_digit=max([len(i) for i in cl])
    return max([i for i in cl if len(i)==max_digit])

MODEL=latest_checkpoint('/home/groups/group_a/yuspqr/result/'+NAME+'.global/*')
model = RobertaLongForSequenceClassification.from_pretrained(MODEL)

for param in model.roberta.embeddings.parameters():
    param.requires_grad = False
for layer in model.roberta.encoder.layer[:8]:
    for param in layer.parameters():
        param.requires_grad = False 

trainer = Trainer(
    model=model,                         # the instantiated 🤗 Transformers model to be trained
    args=training_args,                  # training arguments, defined above
    train_dataset=train_dataset,         # training dataset
    eval_dataset=val_dataset,            # evaluation dataset
)
trainer.train(resume_from_checkpoint=MODEL)
</code></pre>
<p>Here I got the following error. This happened after skipping the first batches.</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "longformer_processing_final2.py", line 209, in &lt;module&gt;
    trainer.train(resume_from_checkpoint=MODEL)
  File "/home/users/yuspqr/.local/lib/python3.6/site-packages/transformers/trainer.py", line 1332, in train
    tr_loss_step = self.training_step(model, inputs)
  File "/home/users/yuspqr/.local/lib/python3.6/site-packages/transformers/trainer.py", line 1891, in training_step
    loss = self.compute_loss(model, inputs)
  File "/home/users/yuspqr/.local/lib/python3.6/site-packages/transformers/trainer.py", line 1923, in compute_loss
    outputs = model(**inputs)
  File "/home/users/yuspqr/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/users/yuspqr/.local/lib/python3.6/site-packages/transformers/models/roberta/modeling_roberta.py", line 1212, in forward
    return_dict=return_dict,
  File "/home/users/yuspqr/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/users/yuspqr/.local/lib/python3.6/site-packages/transformers/models/roberta/modeling_roberta.py", line 861, in forward
    return_dict=return_dict,
  File "/home/users/yuspqr/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/users/yuspqr/.local/lib/python3.6/site-packages/transformers/models/roberta/modeling_roberta.py", line 533, in forward
    output_attentions,
  File "/home/users/yuspqr/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/users/yuspqr/.local/lib/python3.6/site-packages/transformers/models/roberta/modeling_roberta.py", line 417, in forward
    past_key_value=self_attn_past_key_value,
  File "/home/users/yuspqr/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/users/yuspqr/.local/lib/python3.6/site-packages/transformers/models/roberta/modeling_roberta.py", line 346, in forward
    output_attentions,
  File "/home/users/yuspqr/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "longformer_processing_final2.py", line 132, in forward
    output_attentions=output_attentions)
  File "/home/users/yuspqr/.local/lib/python3.6/site-packages/transformers/models/longformer/modeling_longformer.py", line 632, in forward
    attn_probs = torch.masked_fill(attn_probs, is_index_masked[:, :, None, None], 0.0)
RuntimeError: CUDA out of memory. Tried to allocate 1.84 GiB (GPU 0; 31.75 GiB total capacity; 29.37 GiB already allocated; 616.19 MiB free; 29.95 GiB reserved in total by PyTorch)
  0%|          | 0/21072 [01:04&lt;?, ?it/s]
</code></pre>
<p>My research has stopped and I am really in trouble. Your help is sincerely appreciated.</p>

username: eqy
cooked: <p>This is an interesting scenario; can you reproduce it via either a pretrained roberta from huggingface or provide a repro script that e.g., does a few iterations on dummy data and reproduces this OOM when resuming training?</p>
<p>Also out of curiosity, what is the purpose of <code>:8</code> in <code>for layer in model.roberta.encoder.layer[:8]:</code>?</p>
<p>It would also be helpful to understand more about the setup such as the GPU that is being used and the versions of various libraries:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/master/torch/utils/collect_env.py">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/master/torch/utils/collect_env.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/torch/utils/collect_env.py" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/master/torch/utils/collect_env.py</a></h4>


      <pre><code class="lang-py">from __future__ import print_function

# Unlike the rest of the PyTorch this file must be python2 compliant.
# This script outputs relevant system environment info
# Run it with `python collect_env.py`.
import datetime
import locale
import re
import subprocess
import sys
import os
from collections import namedtuple


try:
    import torch
    TORCH_AVAILABLE = True
except (ImportError, NameError, AttributeError, OSError):
    TORCH_AVAILABLE = False

</code></pre>



  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/master/torch/utils/collect_env.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


username: YISTANFORD
cooked: <p>Thanks so much for replying, I appreciate it!<br>
Below are two scripts, which I used for my training.<br>
First, I used training_script_before_resuming to train up to the server time limit, then resumed training with training_script_afrer_resuming and got “CUDA out of memory error”.</p>
<p><a href="https://github.com/huggingface/transformers/files/8646199/longformer_training.txt" rel="noopener nofollow ugc">training_script_before_resuming</a><br>
<a href="https://github.com/huggingface/transformers/files/8646200/longformer_resume_training.txt" rel="noopener nofollow ugc">training_script_after_resuming</a></p>
<p>The purpose of the purpose of <code>:8</code> in <code>for layer in model.roberta.encoder.layer[:8]:</code> is to layer-freeze up to the eighth layer to save GPU memory.</p>
<p>The GPU and other setups are as follows.</p>
<ul>
<li>GPU Tesla V100 32GB × 1</li>
<li>
<code>transformers</code> version: 4.15.0</li>
<li>Platform: Linux-3.10.0-1160.59.1.el7.x86_64-x86_64-with-centos-7.9.2009-Core</li>
<li>Python version: 3.6.8</li>
<li>PyTorch version (GPU?): 1.7.0 (False)</li>
<li>Tensorflow version (GPU?): 2.6.2 (False)</li>
<li>Flax version (CPU?/GPU?/TPU?): not installed (NA)</li>
<li>Jax version: not installed</li>
<li>JaxLib version: not installed</li>
<li>Using GPU in script?: Yes</li>
<li>Using distributed or parallel set-up in script?: No</li>
</ul>
<p>Again, thank you for your help!</p>