Google Colab免费版微调Llama 2时，代码在51000条样本处中断

在Google Colab中使用自定义数据集微调Llama 2时，训练到第51000条样本就停止了，但数据集总共有61609条样本。测试更大的数据集时代码运行完全正常，因此推测问题出在自定义数据集的特性上。

相关背景 #

• 基于公开教程进行Llama 2微调操作
• 自定义数据集包含61609条训练样本
• 替换为其他更大数据集时无异常中断情况

复现代码 #

!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

# 加载数据集
dataset_name = 'harpyerr/merged-pf'
dataset = load_dataset(dataset_name, split="train")

# 模型及LoRA配置
model_name = "your_pretrained_model_name"  # 替换为实际预训练模型名
lora_alpha = 16
lora_dropout = 0.1
lora_r = 64

# 初始化Tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token

# 定义LoRA配置
peft_config = LoraConfig(
    lora_alpha=lora_alpha,
    lora_dropout=lora_dropout,
    r=lora_r,
    bias="none",
    task_type="CAUSAL_LM"
)

# 训练参数设置
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,
)

# 初始化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,
)

# 归一化层转float32
import torch
for name, module in trainer.model.named_modules():
    if "norm" in name:
        module = module.to(torch.float32)

# 启动训练
trainer.train()

可能的原因及解决方案 #

1. 目标样本存在格式异常

   • 定位并检查第51000条样本（数据集索引从0开始，实际取dataset[50999]），查看是否包含特殊字符、空值、乱码或超长内容：

     sample = dataset[50999]
     print(sample["text"])
     

   • 若样本存在问题，直接删除或修复后重新训练。

2. 样本长度超出阈值导致处理失败

   • 即使开启了group_by_length，个别超长样本在tokenization时可能触发错误。添加过滤逻辑提前剔除超长文本：

     def filter_long_samples(example):
         token_len = len(tokenizer(example["text"], truncation=False)["input_ids"])
         return token_len <= max_seq_length
     
     dataset = dataset.filter(filter_long_samples)
     

3. 数据集加载时的隐性损坏

   • 重新加载数据集并验证完整性，遍历到目标样本检查是否能正常访问：

     dataset = load_dataset(dataset_name, split="train")
     print(f"验证样本总数：{len(dataset)}")
     for i in range(50999, 51001):
         print(f"样本{i}内容：{dataset[i]}")
     

   • 若加载出错，重新上传数据集到Hugging Face Hub，确保文件无损坏。

4. Colab资源触发中断

   • 目标样本可能导致显存占用突增，触发Colab资源回收：
     • 降低per_device_train_batch_size至2
     • 在TrainingArguments中添加gradient_checkpointing=True
     • 切换到更高显存的Colab运行时（如A100）

内容的提问来源于stack exchange，提问作者CreekSi0