Загадочная ошибка CUDA при точной настройке модели классификации последовательностей

Question

Я работаю над тонкой настройкой Llama 2 7B для классификации последовательностей с использованием QLoRA. Я использую один графический процессор A100 и получаю одну и ту же загадочную ошибку CUDA даже при увеличении количества графических процессоров, увеличении памяти ЦП и использовании размера пакета, равного 1.

Это ошибка:

RuntimeError: ошибка CUDA: CUBLAS_STATUS_EXECUTION_FAILED при вызовеcublasGemmStridedBatchedExFix(handle, opa, opb, (int)m, (int)n, (int)k, (void*)&falpha, a, CUDA_R_16BF, (int)lda, stridea, b, CUDA_R_16BF, (int)ldb, strideb, (void*)&fbeta, c, CUDA_R_16BF, (int)ldc, stridec, (int)num_batches, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)

Ни моя команда, ни исследовательская группа университета не могут понять, почему я получаю эту ошибку. Тот же код работает для точной настройки BERT, гораздо меньшей модели. Вот мой соответствующий код:

      from transformers import (
    AutoTokenizer,
    BitsAndBytesConfig,
    LlamaTokenizer
)
from datasets import load_dataset, Dataset
import evaluate

from peft import (
    prepare_model_for_kbit_training,
    LoraConfig,
    get_peft_model,
    PeftModel,
    TaskType
)
from peft.tuners.lora import LoraLayer

# Load the dataset
dataset = load_dataset('csv', data_files="https://raw.githubusercontent.com/bryanchrist/llama2-70b/main/app/feedback.csv")

# Get the total number of examples in the dataset
total_examples = len(dataset['train'])

# Calculate the sizes of the training, test, and validation sets
train_size = int(0.8 * total_examples)
test_size = int(0.1 * total_examples)
valid_size = total_examples - train_size - test_size

# Manually split the dataset into training, test, and validation sets
train_dataset = dataset['train'].shuffle(seed=42).select(range(train_size))
test_dataset = dataset['train'].shuffle(seed=42).select(range(train_size, train_size + test_size))
valid_dataset = dataset['train'].shuffle(seed=42).select(range(train_size + test_size, total_examples))

# Create a DatasetDict to hold the splits
train_test_valid_dataset = DatasetDict({
    'train': train_dataset,
    'test': test_dataset,
    'valid': valid_dataset
})

#Set up tokenizer
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", use_auth_token=True)
tokenizer.add_special_tokens({"pad_token":"[PAD]"})

#Preprocess and collate data
def preprocess_function(examples):
    return tokenizer(examples["text"], truncation=True)

tokenized_train_test_valid_dataset = train_test_valid_dataset.map(preprocess_function, batched=True)

from transformers import DataCollatorWithPadding
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

#Prepare evaluation function
import evaluate
accuracy = evaluate.load("accuracy")

import numpy as np
def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return accuracy.compute(predictions=predictions, references=labels)

#Training
id2label = {0: "NOT SOLVABLE", 1: "SOLVABLE"}
label2id = {"NOT SOLVABLE": 0, "SOLVABLE": 1}

#Import model
from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer

model = AutoModelForSequenceClassification.from_pretrained(
    "meta-llama/Llama-2-7b-hf", num_labels=2, id2label=id2label, label2id=label2id,
        use_auth_token=True,  
        device_map = 'auto',
        quantization_config=BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type='nf4'))

model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)

peft_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    r=16,
    lora_alpha=32,
    lora_dropout=0.1,
)      

model = get_peft_model(model, peft_config)

training_args = TrainingArguments(
    output_dir="text_classifier_llama",
    learning_rate=2e-5,
    per_device_train_batch_size=1,
    num_train_epochs=8,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True, 
    bf16=True)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train_test_valid_dataset["train"],
    eval_dataset=tokenized_train_test_valid_dataset["valid"],
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()

Кто-нибудь знает, почему я могу получить эту ошибку?