Llama 3: получение неизвестной ошибки CUDA при точной настройке Llama 3 в викитексте ⇐ Python

[Anonymous]

Я новичок в больших языковых моделях и API Hugging Face. В качестве практики я пытался настроить модель Llama 3.1 8b на наборе данных викитекста.
Когда я пытаюсь запустить следующий скрипт, я получаю неизвестную ошибку CUDA.

Код: Выделить всё

CUDA error: unknown error
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
File "/root/mesh_LLM.py", line 84, in 
trainer.train()
RuntimeError: CUDA error: unknown error
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Моя машина серии Dell Precision с 16 ГБ видеопамяти NVIDIA RTX A5000, поэтому я надеюсь, что это не проблема с памятью, учитывая, что я загружаю и настраиваю модель с точностью до 8 бит.
Вот код:

Код: Выделить всё

import torch,os
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import Trainer, TrainingArguments, BitsAndBytesConfig
from datasets import load_dataset
from peft import LoraConfig, get_peft_model

os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
model_name = "meta-llama/Llama-3.1-8b"

tokenizer = AutoTokenizer.from_pretrained(base_model_id, token='hf-****')
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(base_model_id, device_map='auto', load_in_8bit=True)
model.resize_token_embeddings(len(tokenizer))

peft_config = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, bias='none', task_type="CAUSAL_LM")
model = get_peft_model(model, peft_config)

model.print_trainable_parameters()

# Load a dataset
dataset = load_dataset("wikitext", "wikitext-2-raw-v1")

# Tokenize the dataset
def tokenize_function(examples):
return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)

tokenized_datasets = dataset.map(tokenize_function, batched=True)

training_args = TrainingArguments(
output_dir="./llama3_finetuned",  # Where to save the model
evaluation_strategy="steps",     # Evaluate during training
save_strategy="steps",           # Save checkpoints
learning_rate=2e-5,              # A good starting point for fine-tuning
per_device_train_batch_size=4,   # Adjust based on GPU memory
gradient_accumulation_steps=8,   # Simulates a larger batch size
num_train_epochs=3,              # Experiment with more epochs for small datasets
logging_steps=100,               # Log training progress
save_steps=500,                  # Save model every 500 steps
push_to_hub=False                # Skip pushing to Hugging Face Hub for now
)

trainer = Trainer(
model=model,
train_dataset=tokenized_datasets['train'],
eval_dataset=tokenized_datasets['validation'],
tokenizer=tokenizer,
args=training_args
)

trainer.train()
trainer.save_model('model_ft/fine_tuned_llama3-8B')

Любые предложения будут очень полезны!

Подробнее здесь: https://stackoverflow.com/questions/793 ... n-wikitext