Llama-3.2-1B微调耗尽128GB RAM致内核崩溃求助

架构老林

2026-6-16

问题描述

在配备128GB内存和RTX 4090 GPU的设备上微调meta-llama/Llama-3.2-1B模型，数据集仅10GB共7400行，但初始化SFTTrainer时内存被完全耗尽，内核直接终止。尝试多种公开方案未解决，相关代码和错误日志如下：

import os                       
import torch                    
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    pipeline,
    logging,
)
from peft import LoraConfig
from trl import SFTTrainer


dataset = load_dataset("csv", data_files="dataset/data.csv")

base_model = "meta-llama/Llama-3.2-1B"
compute_dtype = getattr(torch, "float16")

# Configure memory-efficient quantization
compute_dtype = getattr(torch, "float16")
quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=compute_dtype,
    bnb_4bit_use_double_quant=True,  # Enable double quantization
)

model = AutoModelForCausalLM.from_pretrained(
    base_model,
    quantization_config=quant_config,
    device_map="auto",  # Let transformers handle device mapping
    torch_dtype=torch.float16,  # Use fp16 for model weights
    low_cpu_mem_usage=True,    # Enable memory optimization
)

torch.cuda.empty_cache()
model.config.use_cache = False
model.config.pretraining_tp = 1


# Configure PEFT using LoRA for efficient fine-tuning of the model.
peft_params = LoraConfig(
    lora_alpha=16,    
    lora_dropout=0.1, 
    r=8,              
    bias="none",       
    task_type="CAUSAL_LM",
    target_modules="all-linear",
)

training_params = TrainingArguments(
    output_dir="./results",
    num_train_epochs=1,
    per_device_train_batch_size=2,  
    gradient_accumulation_steps=2,   
    optim="paged_adamw_8bit",    
    save_steps=50,
    logging_steps=50,
    learning_rate=2e-4,
    weight_decay=0.001,
    fp16=True,                   
    bf16=False,
    max_grad_norm=0.3,
    max_steps=-1,
    warmup_ratio=0.03,
    group_by_length=True,
    lr_scheduler_type="constant",
    report_to="tensorboard",
    gradient_checkpointing=True,   
)


tokenizer = AutoTokenizer.from_pretrained(
    base_model,
    padding_side="right",
    truncation_side="right",
)
tokenizer.pad_token = tokenizer.eos_token

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset['train'],
    peft_config=peft_params,
    dataset_text_field="input_text",
    max_seq_length=512,
    tokenizer=tokenizer,
    args=training_params,
    packing=False,
)

trainer.train()

错误日志（翻译后）：

SFTTrainer使用了已弃用的位置参数，请改用SFTConfig设置这些参数。
  warnings.warn(message, FutureWarning)
/home/.../python3.10/site-packages/trl/trainer/sft_trainer.py:300: UserWarning: 你向SFTTrainer传递了`max_seq_length`参数，该值将覆盖SFTConfig中的设置。
  warnings.warn(
/home/.../python3.10/site-packages/trl/trainer/sft_trainer.py:328: UserWarning: 你向SFTTrainer传递了`dataset_text_field`参数，该值将覆盖SFTConfig中的设置。
  warnings.warn(
Map:  14%|██████████████████▍                                                                                                                    | 1000/7346 [02:59&lt;18:56,  5.58 examples/s]
Killed

解决方案

优化数据集加载逻辑
避免一次性加载全量10GB数据，改用流式加载或分批映射：

# 方案1：流式加载数据集，按需读取
dataset = load_dataset("csv", data_files="dataset/data.csv", streaming=True)

# 方案2：非流式加载时，分批处理数据映射
dataset = dataset.map(
    lambda examples: tokenizer(examples["input_text"], truncation=True, max_length=512),
    batched=True,
    batch_size=100  # 减小批量大小降低内存占用
)

适配新版TRL的参数规范
新版TRL推荐用SFTConfig替代直接传参数，避免冗余配置冲突：

from trl import SFTConfig

sft_config = SFTConfig(
    dataset_text_field="input_text",
    max_seq_length=512,
    packing=False,
    args=training_params
)

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset['train'],
    peft_config=peft_params,
    tokenizer=tokenizer,
    config=sft_config
)

进一步压缩内存占用

把per_device_train_batch_size从2降到1，减少单批数据内存占用
启用CUDA内存优化：torch.backends.cuda.enable_mem_efficient_sdp(True)

若硬件支持，加载模型时使用FlashAttention2替代默认注意力：

model = AutoModelForCausalLM.from_pretrained(
    base_model,
    quantization_config=quant_config,
    device_map="auto",
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
    attn_implementation="flash_attention_2"  # 需要提前安装flash-attn库
)

手动清理内存碎片
在关键节点清理CUDA缓存和Python垃圾，避免内存泄漏：

import gc

# 加载模型后清理
torch.cuda.empty_cache()
gc.collect()

# 初始化trainer前再次清理
torch.cuda.empty_cache()
gc.collect()

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