I am reaching out to seek assistance regarding a persistent issue I am facing 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 two GPUs with sufficient memory (NVIDIA RTX A4000), I continually encounter the error message: “torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 224.00 MiB.” This issue persists even though GPU memory usage indicates available space on both GPUs.
Below are the system and software specifications for your reference:
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 in WSL: CUDA compilation tools, release 11.8
Library Versions:
PyTorch version: 2.3.0+cu118
Transformers version: 4.40.1
scikit-learn version: 1.4.2
Pandas version: 2.2.2
I have attempted various troubleshooting steps, including clearing GPU cache, ensuring proper CUDA and GPU utilization, and optimizing code for memory usage. However, the issue persists.
Is there anyone who has a solution for this issue when finetuning an llm
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
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="xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx")
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 = "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
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.")
