OutOfMemoryError: CUDA out of memory in Google Colab

Raj_Kumar2: <p>I am trying to fine tune “Llama-2-7b-chat-hf” Model with “mlabonne/guanaco-llama2-1k” in Google Colab with T4 runtime environment. I am using Qlora technique to fine tune this model. Below is the code I am using.</p>
<blockquote>
<pre><code class="lang-auto">import os
import torch
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    HfArgumentParser,
    TrainingArguments,
    pipeline,
    logging,
)
from peft import LoraConfig, PeftModel
from trl import SFTTrainer

# Model from Hugging Face hub
base_model = "NousResearch/Llama-2-7b-chat-hf"

# New instruction dataset
guanaco_dataset = "mlabonne/guanaco-llama2-1k"

# Fine-tuned model
new_model = "llama-2-7b-chat-guanaco"

# Load Dataset
dataset = load_dataset(guanaco_dataset, split="train")


compute_dtype = getattr(torch, "float16")

quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=compute_dtype,
    bnb_4bit_use_double_quant=False,
)

# Load base model
model = AutoModelForCausalLM.from_pretrained(
    base_model,
    quantization_config=quant_config,
    device_map={"": 0}
)
model.config.use_cache = False
model.config.pretraining_tp = 1

# Load LLaMA tokenizer
tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right" 


# Load LoRA configuration
peft_args = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.1,
    r=64,
    bias="none",
    task_type="CAUSAL_LM",
)

# Set training parameters
training_params = TrainingArguments(
    output_dir="./results",
    num_train_epochs=1,
    per_device_train_batch_size=3,
    gradient_accumulation_steps=1,
    optim="paged_adamw_32bit",
    save_steps=25,
    logging_steps=25,
    learning_rate=2e-4,
    weight_decay=0.001,
    fp16=False,
    bf16=False,
    max_grad_norm=0.3,
    max_steps=-1,
    warmup_ratio=0.03,
    group_by_length=True,
    lr_scheduler_type="constant",
    report_to="tensorboard"
)

# Set supervised fine-tuning parameters
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=peft_args,
    dataset_text_field="text",
    max_seq_length=None,
    tokenizer=tokenizer,
    args=training_params,
    packing=False,
)

# Train model
trainer.train()
</code></pre>
</blockquote>
<p>After running the last code, I was getting the below error.</p>
<p>When “per_device_train_batch_size=4”, I got this error.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/7/2758770330ba263c7c6a612ba9e0e3249db3952e.png" data-download-href="https://discuss.pytorch.org/uploads/default/2758770330ba263c7c6a612ba9e0e3249db3952e" title="memory error"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/2758770330ba263c7c6a612ba9e0e3249db3952e_2_690x201.png" alt="memory error" data-base62-sha1="5C49HBCER6XRfhCeCsBIDvpWi18" width="690" height="201" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/2758770330ba263c7c6a612ba9e0e3249db3952e_2_690x201.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/2758770330ba263c7c6a612ba9e0e3249db3952e_2_1035x301.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/2758770330ba263c7c6a612ba9e0e3249db3952e_2_1380x402.png 2x" data-dominant-color="393A3A"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">memory error</span><span class="informations">1831×535 33.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>And, “When per_device_train_batch_size=3”, I got this error.</p>
<p>same error occurred with different memory utilization.</p>
<p>I have tried with changing <code>per_device_train_batch_size=</code> to 3,4,5,6 but not working for me.</p>
<p>Also tried</p>
<pre><code class="lang-auto"># Set the environment variable
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

# Verify that the environment variable is set
print(os.environ["PYTORCH_CUDA_ALLOC_CONF"])
</code></pre>
<p>And lastly</p>
<pre><code class="lang-auto">torch.cuda.empty_cache()
</code></pre>
<p>Nothing worked for me.</p>

drisspg: <p>Have you tried profiling the memory usage following techniques mentioned here: <a href="https://pytorch.org/blog/understanding-gpu-memory-1/" class="inline-onebox" rel="noopener nofollow ugc">Understanding GPU Memory 1: Visualizing All Allocations over Time | PyTorch</a></p>

Raj_Kumar2: <p>Tried: optimizer = torch.optim.AdamW(model.parameters(), lr=training_params.learning_rate)<br>
and <code>optimizer.zero_grad(set_to_none=True)</code></p>
<p>Still getting thre same error.</p>