Уровень бесплатного пользования Google Colab: код останавливается на 51000 примеров при тонкой настройке LLAMA 2 с помощью пользовательского набора данных

Question

Я столкнулся с проблемой при точной настройке Llama 2 в Google Colab с использованием специального набора данных. В процессе обучения код останавливается ровно на 51000 примерах, хотя мой набор данных содержит 61609 примеров. Странно то, что когда я тестировал код с еще большими наборами данных, он работал совершенно нормально. Я следил за руководством на YouTube, чтобы точно настроить Llama 2, и вы можете найти исходную ссылку на Colab и руководство ниже:

Ссылка на руководство: YouTube Tutorial Original Colab: Google Colab

Ссылка на набор данных: Мой пользовательский набор данных

Код:

      !pip install -q -U trl transformers accelerate git+https://github.com/huggingface/peft.git
!pip install -q datasets bitsandbytes einops wandb

from datasets import load_dataset
from transformers import AutoTokenizer, TrainingArguments
from peft import LoraConfig, get_peft_model
from trl import SFTTrainer

# Load dataset
dataset_name = 'harpyerr/merged-pf'
dataset = load_dataset(dataset_name, split="train")

# Define model_name, lora_alpha, lora_dropout, lora_r, and other configurations
model_name = "your_pretrained_model_name"  # Replace with the name of your pretrained model
lora_alpha = 16
lora_dropout = 0.1
lora_r = 64

# Initialize tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token

# Define LoraConfig
peft_config = LoraConfig(
    lora_alpha=lora_alpha,
    lora_dropout=lora_dropout,
    r=lora_r,
    bias="none",
    task_type="CAUSAL_LM"
)

# Define training arguments
output_dir = "./results"
per_device_train_batch_size = 4
gradient_accumulation_steps = 4
optim = "paged_adamw_32bit"
save_steps = 100
logging_steps = 10
learning_rate = 2e-4
max_grad_norm = 0.3
max_steps = 100
warmup_ratio = 0.03
lr_scheduler_type = "constant"

training_arguments = TrainingArguments(
    output_dir=output_dir,
    per_device_train_batch_size=per_device_train_batch_size,
    gradient_accumulation_steps=gradient_accumulation_steps,
    optim=optim,
    save_steps=save_steps,
    logging_steps=logging_steps,
    learning_rate=learning_rate,
    fp16=True,
    max_grad_norm=max_grad_norm,
    max_steps=max_steps,
    warmup_ratio=warmup_ratio,
    group_by_length=True,
    lr_scheduler_type=lr_scheduler_type,
)

# Initialize SFTTrainer
max_seq_length = 512
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=peft_config,
    dataset_text_field="text",
    max_seq_length=max_seq_length,
    tokenizer=tokenizer,
    args=training_arguments,
)

# Convert all normalization layers to float32
import torch
for name, module in trainer.model.named_modules():
    if "norm" in name:
        module = module.to(torch.float32)

# Start training
trainer.train()

Я пробовал использовать разные наборы данных большего размера, чтобы проверить, связана ли проблема с моим пользовательским набором данных. Удивительно, но когда я использовал другие более крупные наборы данных, код работал отлично, без каких-либо проблем с остановкой. Итак, я пришел к выводу, что проблема не в коде или трейнере, а может быть связана с конкретными характеристиками моего пользовательского набора данных.