Capítulo 3: Entrenamiento eficiente en memoria
Optimización del entrenamiento para hardware limitado
Este capítulo cubre todas las técnicas para maximizar la eficiencia del entrenamiento con memoria GPU limitada.
Técnicas de optimización de memoria
1. Gradient Checkpointing
Intercambie cómputo por memoria recomputando activaciones:
# Enable via training args
training_args = TrainingArguments(
gradient_checkpointing=True,
...
)
# Or directly on model
model.gradient_checkpointing_enable()
Ahorro: reducción de memoria de aproximadamente 30%
Costo: entrenamiento aproximadamente 20% más lento
2. Gradient Accumulation
Simule lotes más grandes con memoria limitada:
training_args = TrainingArguments(
per_device_train_batch_size=1, # Lote pequeño cabe en memoria
gradient_accumulation_steps=16, # Acumula gradientes
# Tamaño de lote efectivo = 1 * 16 = 16
)
3. Optimizadores eficientes
Adam de 8 bits
training_args = TrainingArguments(
optim="adamw_8bit", # Requiere bitsandbytes
)
Optimizadores paginados
training_args = TrainingArguments(
optim="paged_adamw_8bit", # Pagina a CPU cuando es necesario
)
Script completo de entrenamiento eficiente en memoria
import torch
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
BitsAndBytesConfig,
TrainingArguments,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer
from datasets import load_dataset
# Quantization
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
)
# Load model
model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-3.1-8B",
quantization_config=bnb_config,
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B")
tokenizer.pad_token = tokenizer.eos_token
# Prepare for training
model = prepare_model_for_kbit_training(model)
# LoRA config
lora_config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
lora_dropout=0.05,
task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)
# Dataset
dataset = load_dataset("databricks/dolly-15k", split="train[:1000]")
# Memory-efficient training args
training_args = TrainingArguments(
output_dir="./qlora-output",
num_train_epochs=1,
per_device_train_batch_size=1,
gradient_accumulation_steps=16,
learning_rate=2e-4,
bf16=True,
gradient_checkpointing=True,
optim="paged_adamw_8bit",
logging_steps=10,
save_strategy="epoch",
)
# Train
trainer = SFTTrainer(
model=model,
train_dataset=dataset,
args=training_args,
dataset_text_field="text",
max_seq_length=512,
tokenizer=tokenizer,
)
trainer.train()
model.save_pretrained("./qlora-adapter")
Lista de verificación de memoria
□ Use QLoRA (cuantización de 4 bits)
□ Active gradient checkpointing
□ Use gradient accumulation
□ Use optimizador de 8 bits o paginado
□ Use precisión bfloat16
□ Reduzca max_seq_length si es posible
□ Use tamaños de lote más pequeños
Uso de memoria esperado
| Modelo | Sin optimización | Optimización completa |
|---|---|---|
| 7B | ~14GB | ~4-6GB |
| 13B | ~26GB | ~8-10GB |
| 70B | ~140GB | ~40-48GB |
Puntos clave
✅ Combine múltiples técnicas de optimización
✅ Gradient checkpointing intercambia cómputo por memoria
✅ Gradient accumulation simula lotes más grandes
✅ Los optimizadores paginados manejan picos de memoria
✅ QLoRA + todas las optimizaciones permiten 7B en 6GB