I’m encountering a persistent issue while fine-tuning a Llama3 model using a Hugging Face dataset in a Windows Subsystem for Linux (WSL) Ubuntu 22.04 environment on Windows 11. Despite having ample GPU memory available (detailed specs provided below), I repeatedly face the error: “torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 224.00 MiB. GPU”. This occurs despite having two GPUs with sufficient memory.

System Specifications:

Processor: AMD Ryzen Threadripper PRO 5955WX 16-Cores @ 4.00 GHz
Installed RAM: 128 GB (128 GB usable)
Operating System: Windows 11 Pro Version 23H2, OS build 22631.3296
CUDA Version: 11.8
Available GPU devices: 2
GPU 1: NVIDIA RTX A4000
CUDA Version: 11.8
Software Versions:

PyTorch version: 2.3.0+cu118
Transformers version: 4.40.1
scikit-learn version: 1.4.2
Pandas version: 2.2.2
CUDA Version in WSL: CUDA compilation tools, release 11.8
Steps Taken:

Cleared GPU cache.
Checked CUDA version, GPU memory, and availability.
Ensured GPU 1 is utilized.
Loaded and preprocessed dataset.
Loaded model and tokenizer with authentication token.
Handled special tokens.
Setup mixed precision.
Defined training setup.
Initialized trainer.
Defined training function with autocast and scaler.
Initiated training.
Despite these efforts, the “CUDA out of memory” error persists.

Full Code:

from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
from datasets import load_dataset, Dataset
import pandas as pd
from sklearn.model_selection import train_test_split
import os

import torch
import transformers
from torch.cuda.amp import autocast, GradScaler

import GPUtil
from numba import cuda


import torch
print(f'PyTorch version: {torch.__version__}')
print('*'*10)
print(f'_CUDA version: ')
print('*'*10)
print(f'CUDNN version: {torch.backends.cudnn.version()}')
print(f'Available GPU devices: {torch.cuda.device_count()}')
print(f'Device Name: {torch.cuda.get_device_name()}')

# Free GPU cache
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
from datasets import load_dataset, Dataset
import pandas as pd
from sklearn.model_selection import train_test_split
import os

import torch
import transformers
from torch.cuda.amp import autocast, GradScaler

import GPUtil
from numba import cuda


import torch
print(f'PyTorch version: {torch.__version__}')
print('*'*10)
print(f'_CUDA version: ')
print('*'*10)
print(f'CUDNN version: {torch.backends.cudnn.version()}')
print(f'Available GPU devices: {torch.cuda.device_count()}')
print(f'Device Name: {torch.cuda.get_device_name()}')

# Free GPU cache
def free_gpu_cache():
    print("Initial GPU Usage:")
    GPUtil.showUtilization()

    torch.cuda.empty_cache()  # Clear PyTorch's cache
    cuda.select_device(1)
    cuda.close()  # Clear Numba cache
    cuda.select_device(1)

    print("GPU Usage after clearing cache:")
    GPUtil.showUtilization()

# Clear memory and show initial status
free_gpu_cache()

# Check CUDA version and GPU memory
print("CUDA Version:", torch.version.cuda)
print(f"PyTorch version: {torch.__version__}")
print(f"Transformers version: {transformers.__version__}")
print(f"CUDA Available: {torch.cuda.is_available()}")

# Ensure GPU 1 is used
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

# Check current GPU memory
print(torch.cuda.memory_summary())

# Load dataset
print("Loading dataset...")
dataset = load_dataset("naklecha/minecraft-question-answer-700k", token="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")
print("Dataset loaded successfully.")

# Convert to pandas DataFrame
data_df = pd.DataFrame(dataset['train'])

# Split the dataset into training and test sets
train_df, test_df = train_test_split(data_df, test_size=0.1)

# Convert back to Hugging Face datasets
train_dataset = Dataset.from_pandas(train_df)
test_dataset = Dataset.from_pandas(test_df)

print(f"Train set length: {len(train_dataset)}")
print(f"Test set length: {len(test_dataset)}")

# Load model and tokenizer with an authentication token
model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
auth_token = "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"

print("Loading the tokenizer and model...")
tokenizer = AutoTokenizer.from_pretrained(model_id, token=auth_token)
model = AutoModelForCausalLM.from_pretrained(model_id, token=auth_token)
print("Tokenizer and model loaded successfully.")

# Handle special tokens potentially added to the tokenizer
special_tokens_dict = {'additional_special_tokens': ['[USR]', '[SYS]']}
num_added_toks = tokenizer.add_special_tokens(special_tokens_dict)
model.resize_token_embeddings(len(tokenizer))
print(f"Added {num_added_toks} special tokens.")

# Mixed Precision setup
scaler = GradScaler()

# Training setup
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=1,  # Decrease batch size
    per_device_eval_batch_size=1,  # Reduce eval batch size
    gradient_accumulation_steps=8,  # Increase accumulation steps
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=10,
    report_to="none",
    load_best_model_at_end=True,
    metric_for_best_model='loss',
    greater_is_better=False,
    evaluation_strategy="steps",
    save_strategy="steps"
)

# Initialize trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset
)

# Define training function with autocast and scaler
def training_step(trainer, model, criterion, optimizer, data_loader):
    for i, data in enumerate(data_loader, 1):
        with autocast():
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            scaler.scale(loss).backward()

            # Accumulate gradients or step optimizer
            if i % training_args.gradient_accumulation_steps == 0:
                scaler.step(optimizer)
                scaler.update()

# Start training
print("Starting training...")
training_step(trainer, model, criterion, optimizer, train_loader)
print("Training completed.")

# Save model and tokenizer
model.save_pretrained('./fine_tuned_model')
tokenizer.save_pretrained('./fine_tuned_model')
print("Model and tokenizer saved successfully.")

[enter image description here](https://i.sstatic.net/8tQjRxTK.png)

I want to solve this error and finetune a llama3 model