Por qué no pude entrenar Gemma 3n localmente (y por qué estoy usando Vertex AI en su lugar)

He estado trabajando en una aplicación de entrenamiento de pronunciación que necesita una IA multimodal capaz de entender tanto texto como audio. Actualmente, la aplicación utiliza los modelos de audio nativos de Google Gemini, pero quería ajustar un modelo dedicado para el entrenamiento de pronunciación. Gemma 3n de Google parecía perfecto con su soporte nativo de audio a través del codificador Universal Speech Model. Tengo un MacBook Pro M4 Pro con 48 GB de RAM. Seguramente es suficiente para entrenar un modelo de 5B parámetros localmente, ¿verdad?

Spoiler: no lo es. Esto es lo que aprendí al fracasar.

Por qué intenté el entrenamiento local

El atractivo era obvio:

• Sin coste: El entrenamiento en GPU en la nube cuesta entre 2 y 5 dólares por intento en servicios como RunPod o Vast.ai
• Privacidad: Mis datos de entrenamiento se quedan locales
• Iteración rápida: Sin necesidad de subir conjuntos de datos o configurar entornos en la nube
• Aprendizaje práctico: Quería experimentar con scripts de entrenamiento de PyTorch, comprender las técnicas de optimización de memoria y aprender mediante prueba y error. Localmente, puedo modificar el código y reintentar en segundos, en lugar de desplegar en la nube y esperar a que se ejecuten los scripts

Tenía 988 ejemplos de entrenamiento que cubrían errores de pronunciación, análisis de la entrega al hablar y conversaciones de entrenamiento. Parecía un candidato perfecto para el entrenamiento local durante la noche.

¿Qué significa realmente entrenar un modelo?

Aclaración rápida: cuando la gente dice "entrenar", pueden referirse a dos cosas muy diferentes.

Entrenamiento desde cero es lo que hacen Google y OpenAI. Empiezas con números aleatorios y enseñas a una red neuronal todo desde cero. Esto requiere miles de millones de muestras de texto, cientos de GPU y meses de tiempo de cómputo. Es caro y lento.

El ajuste fino (Fine-tuning) es tomar un modelo ya entrenado y enseñarle tu tarea específica. El trabajo duro ya está hecho. El modelo conoce el lenguaje y el razonamiento. Solo le estás mostrando tu formato y dominio específicos. Es como contratar a alguien que ya sabe programar y enseñarle tu base de código, en lugar de enseñarle a programar desde cero.

Para mi caso de uso, quiero el ajuste fino. El modelo ya entiende el inglés y la pronunciación. Solo necesito que dé su opinión en mi formato.

Cómo hacer ajuste fino con Hugging Face

Hugging Face hace que todo este proceso sea mucho más fácil de lo que debería ser. Alojan miles de modelos preentrenados y te proporcionan bibliotecas de Python para ajustarlos. No necesitas un doctorado para hacer esto.

Este es el flujo de trabajo básico:

from transformers import AutoProcessor, AutoModelForImageTextToText

# Descargar el modelo y el procesador
processor = AutoProcessor.from_pretrained("google/gemma-3n-E2B-it")
model = AutoModelForImageTextToText.from_pretrained("google/gemma-3n-E2B-it")

Dos líneas de código. Eso descarga un modelo de IA de 5.400 millones de parámetros en tu portátil. El processor convierte tus datos al formato correcto, y model es la red neuronal en sí.

Luego usas la clase Trainer para ejecutar realmente el entrenamiento:

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="outputs/gemma3n",
    num_train_epochs=3,
    per_device_train_batch_size=1,
    learning_rate=3e-5,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=my_dataset,
)

trainer.train()  # ¡Empezar a entrenar!

Esto es genial porque te saltas todo lo de bajo nivel, como escribir bucles de descenso de gradiente y gestionar puntos de control. ¿La desventaja? No sabes realmente lo que está pasando bajo el capó, especialmente cuánta memoria usa cada cosa. Que es exactamente el problema con el que me encontré.

El problema de la memoria

Entrenar un modelo de lenguaje grande requiere que varias cosas quepan en la memoria simultáneamente:

1. Pesos del modelo: Para Gemma 3n E2B (5.4B parámetros), son ~11 GB en float16
2. Gradientes: Otros ~11 GB (mismo tamaño que los pesos)
3. Estados del optimizador: Adam guarda 2 estados por parámetro, añadiendo ~22 GB
4. Activaciones: Resultados intermedios de la computación
5. Datos del lote (Batch data): Los ejemplos de entrenamiento que se están procesando

Total: 50-60 GB para el modelo E2B más pequeño, 70-80 GB para E4B. Mis 48 GB de RAM no iban a ser suficientes.

Lo que intenté

Intento 1: Ajuste fino completo con E4B
Sin memoria durante la inicialización del optimizador. El modelo se cargó (16 GB), pero el optimizador no pudo asignar su estado.

Intento 2: Tamaños de lote más pequeños
Reducir batch_size de 2 a 1 solo retrasó el error OOM hasta el pase hacia atrás (backward pass).

Intento 3: LoRA con E4B
LoRA (Low-Rank Adaptation) solo entrena capas adaptadoras pequeñas en lugar del modelo completo, reduciendo los parámetros entrenables de 7.8B a 40M (0.5%). Pero el modelo base congelado todavía necesita caber en la memoria: 59 GB solo para E4B.

Intento 4: LoRA con E2B
Esto casi funcionó. El modelo E2B más pequeño con LoRA usó 40-45 GB, dejando 3-8 GB de margen. El entrenamiento comenzó y se ejecutó durante unas horas antes de encontrar errores de memoria durante lotes específicos con secuencias más largas.

Por qué falló finalmente

Incluso con optimizaciones agresivas (LoRA, batch_size=1, sin punto de control de gradiente), 48 GB no fueron suficientes para un entrenamiento constante. El uso de memoria variaba según el ejemplo:

• Ejemplos cortos: 38-42 GB (bien)
• Ejemplos largos: 46-50 GB (fallos)

Probablemente podría hacerlo funcionar filtrando los ejemplos más largos o reduciendo max_length a 1024 tokens, pero en ese momento estoy comprometiendo las capacidades del modelo para adaptarme a las limitaciones del hardware.

Pasando a Vertex AI

Después de tres días de intentos fallidos, estoy cambiando al entrenamiento en Vertex AI de Google Cloud:

• GPU A100 de 40 GB: Suficiente memoria para el modelo E4B con LoRA
• Infraestructura gestionada: Sin depuración de memoria, sin fallos
• Coste: Aproximadamente 3-5 dólares por una ejecución de entrenamiento completa (3 épocas)
• Tiempo: 2-3 horas frente a las 6-9 horas que intenté localmente

Escribiré una publicación de blog separada sobre el proceso de entrenamiento de Vertex AI, que incluirá:

• Configuración de trabajos de entrenamiento personalizados con Gemma 3n
• Subida de conjuntos de datos a Cloud Storage
• Monitorización del entrenamiento con TensorBoard
• Despliegue del modelo entrenado en los endpoints de Vertex AI

¿Cuándo puedes entrenar localmente?

El entrenamiento local en Apple Silicon es viable si:

• Tu modelo es más pequeño: Los modelos de 1-3B parámetros con LoRA caben cómodamente en 48 GB
• Tienes más RAM: El M4 Max con 128 GB manejaría Gemma 3n E2B fácilmente
• Usas cuantización: La cuantización de 4 o 8 bits reduce aún más la memoria, pero añade complejidad

Lecciones clave

• La memoria es el cuello de botella: 48 GB de memoria unificada suenan a mucho hasta que intentas entrenar modelos de 5B+ parámetros. Solo los estados del optimizador pueden duplicar o triplicar tus requisitos de memoria
• LoRA ayuda pero no es magia: Reduce los parámetros entrenables en un 99%, pero los pesos del modelo base congelados todavía necesitan caber en la memoria durante los pases hacia adelante y hacia atrás
• El entrenamiento en la nube es rentable: 3-5 dólares por una ejecución de entrenamiento es razonable en comparación con días de depuración de errores OOM y experimentos fallidos
• Conoce los límites de tu hardware: La memoria unificada de Apple Silicon es excelente para la inferencia y el entrenamiento de modelos más pequeños (1-3B parámetros), pero entrenar modelos de 5B+ de forma fiable requiere VRAM de GPU dedicada

Recursos

• Gemma 3n E2B en HuggingFace
• Descripción general del entrenamiento de Vertex AI
• PEFT: Parameter-Efficient Fine-Tuning