Saltar a contenido

Unsloth + LoRA en M4 Pro

Receta para LoRA fine-tuning de Gemma 4 E2B / E4B en M4 Pro 24 GB.


Instalación

uv pip install unsloth
# o con pip:
pip install "unsloth[mps]>=2026.5"  # versión con soporte Apple Silicon

Requisitos: - Python 3.11+. - macOS 14+ con MPS habilitado. - ~15 GB libres en disco para checkpoints.


Script base (finetune.py)

"""LoRA fine-tune de Gemma 4 E4B con Unsloth en M4 Pro."""

from unsloth import FastLanguageModel
from datasets import load_dataset
from trl import SFTTrainer
from transformers import TrainingArguments

# -----------------------------------------------------------------------------
# Config
# -----------------------------------------------------------------------------
MODEL_NAME = "google/gemma-4-e4b"
DATASET_PATH = "data/train.jsonl"
OUTPUT_DIR = "outputs/gemma4-e4b-lora-v1"
MAX_SEQ_LENGTH = 4096
BATCH_SIZE = 2
GRAD_ACCUM = 4
LEARNING_RATE = 2e-4
NUM_EPOCHS = 2

# -----------------------------------------------------------------------------
# Load model + tokenizer
# -----------------------------------------------------------------------------
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=MODEL_NAME,
    max_seq_length=MAX_SEQ_LENGTH,
    dtype=None,                  # auto (bfloat16 on Apple Silicon)
    load_in_4bit=True,           # QLoRA (4-bit base + LoRA adapters)
)

# -----------------------------------------------------------------------------
# Add LoRA adapters
# -----------------------------------------------------------------------------
model = FastLanguageModel.get_peft_model(
    model,
    r=16,                        # rank
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    lora_alpha=16,
    lora_dropout=0,
    bias="none",
    use_gradient_checkpointing="unsloth",
    random_state=42,
)

# -----------------------------------------------------------------------------
# Dataset
# -----------------------------------------------------------------------------
dataset = load_dataset("json", data_files=DATASET_PATH, split="train")

def formatting_func(example):
    """Format messages as Gemma 4 chat template."""
    return tokenizer.apply_chat_template(
        example["messages"],
        tokenize=False,
        add_generation_prompt=False,
    )

# -----------------------------------------------------------------------------
# Trainer
# -----------------------------------------------------------------------------
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    formatting_func=formatting_func,
    max_seq_length=MAX_SEQ_LENGTH,
    args=TrainingArguments(
        output_dir=OUTPUT_DIR,
        per_device_train_batch_size=BATCH_SIZE,
        gradient_accumulation_steps=GRAD_ACCUM,
        learning_rate=LEARNING_RATE,
        num_train_epochs=NUM_EPOCHS,
        warmup_steps=10,
        logging_steps=10,
        save_steps=200,
        bf16=True,                   # bfloat16 en M-series
        optim="adamw_torch",
        seed=42,
    ),
)

# -----------------------------------------------------------------------------
# Train
# -----------------------------------------------------------------------------
trainer.train()

# -----------------------------------------------------------------------------
# Save LoRA adapter only (small ~150 MB)
# -----------------------------------------------------------------------------
model.save_pretrained(f"{OUTPUT_DIR}/adapter")
tokenizer.save_pretrained(f"{OUTPUT_DIR}/adapter")

# -----------------------------------------------------------------------------
# Merge + export GGUF for Ollama (optional, ~5.5 GB output)
# -----------------------------------------------------------------------------
model.save_pretrained_gguf(
    f"{OUTPUT_DIR}/gguf",
    tokenizer,
    quantization_method="q4_k_m",
)

Correr

cd fine-tuning
uv run python finetune.py

Tiempo esperado en M4 Pro 24 GB: - 1K ejemplos / 2 epochs: ~20-40 min - 5K ejemplos / 2 epochs: ~2-4 h - 10K ejemplos / 2 epochs: ~4-8 h


Memory budget M4 Pro 24 GB

Con E4B + QLoRA 4-bit + max_seq 4096: - Modelo: ~3.5 GB - LoRA adapters: ~0.2 GB - Optimizer states: ~1 GB - Gradients: ~0.5 GB - Activations + batch: ~6-8 GB - Total: ~12-14 GB → cabe holgado en 24 GB

Para E4B con max_seq 8192: sube a ~18 GB, justo en el límite.


Cargar adapter de vuelta a Ollama

# 1. Crea Modelfile que apunta al GGUF
cat > Modelfile <<EOF
FROM ./outputs/gemma4-e4b-lora-v1/gguf/gemma4-e4b-lora-q4_k_m.gguf

TEMPLATE """{{ .System }}
<start_of_turn>user
{{ .Prompt }}<end_of_turn>
<start_of_turn>model
"""

SYSTEM """You are a helpful assistant."""

PARAMETER temperature 0.3
PARAMETER num_ctx 8192
EOF

# 2. Import a Ollama
ollama create my-gemma4-finetuned -f Modelfile

# 3. Run
ollama run my-gemma4-finetuned "test prompt"

Trucos y gotchas

Aprende con learning rate más bajo si tu dataset es pequeño

  • 500-2K ejemplos: learning_rate=1e-4
  • 2K-10K ejemplos: learning_rate=2e-4 (default)
  • 10K+ ejemplos: learning_rate=3e-4

Si OOM

  • Reduce MAX_SEQ_LENGTH a 2048.
  • Reduce BATCH_SIZE a 1, sube GRAD_ACCUM a 8.
  • Cambia lora_r=8 (menos parámetros).

Si calidad no mejora

  • Verifica eval set: ¿es representativo?
  • Mira ejemplos del dataset: ¿hay ruido / labels malas?
  • Sube r a 32 (más expressivo).
  • Más epochs (3-4 max, después overfitting).

Si el modelo "olvida" capacidades generales (catastrophic forgetting)

  • Mezcla 10-20% de tu dataset con instrucciones generales (ej: subset de tulu).
  • Reduce learning rate a 1e-4.
  • Menos epochs (1-2).

Eval después del fine-tune

# Eval con eval set (no visto en entreno)
uv run python eval.py \
  --model-base google/gemma-4-e4b \
  --model-finetuned outputs/gemma4-e4b-lora-v1/adapter \
  --eval-set data/eval.jsonl \
  --output reports/eval-v1.json

Target: >5% improvement absoluto sobre baseline en métricas relevantes (F1, exact match, BLEU, etc., según tarea).

Si <5%, no deployes — vuelve a iterar dataset.