Perché non sono riuscito ad addestrare Gemma 3n in locale (e perché sto usando Vertex AI)
Stavo lavorando a un'app di coaching sulla pronuncia che necessita di un'IA multimodale capace di comprendere sia testo che audio. L'app utilizza attualmente i modelli audio nativi di Google Gemini, ma volevo ottimizzare (fare il fine-tuning) un modello dedicato per il coaching della pronuncia. Gemma 3n di Google sembrava perfetto grazie al suo supporto audio nativo tramite l'encoder Universal Speech Model. Ho un MacBook Pro M4 Pro con 48 GB di RAM. Sicuramente sono sufficienti per addestrare un modello da 5 miliardi di parametri in locale, giusto?
Spoiler: non lo sono. Ecco cosa ho imparato dai miei fallimenti.
Perché ho provato l'addestramento locale
Il fascino era evidente:
- Nessun costo: L'addestramento su GPU cloud costa 2-5 $ per tentativo su servizi come RunPod o Vast.ai
- Privacy: I miei dati di addestramento rimangono locali
- Iterazione rapida: Niente caricamento di dataset o configurazione di ambienti cloud
- Apprendimento pratico: Volevo sperimentare con gli script di addestramento PyTorch, comprendere le tecniche di ottimizzazione della memoria e imparare per tentativi ed errori. In locale, posso modificare il codice e riprovare in pochi secondi, rispetto al deployment su cloud e all'attesa che gli script vengano eseguiti.
Avevo 988 esempi di addestramento che coprivano errori di pronuncia, analisi dell'esecuzione del parlato e conversazioni di coaching. Sembrava un candidato perfetto per un addestramento locale notturno.
Cosa significa realmente addestrare un modello?
Breve chiarimento: quando le persone dicono "addestramento" (training), potrebbero riferirsi a due cose molto diverse.
L'addestramento da zero (Training from scratch) è ciò che fanno Google e OpenAI. Si parte da numeri casuali e si insegna a una rete neurale tutto da zero. Ciò richiede miliardi di campioni di testo, centinaia di GPU e mesi di tempo di calcolo. È costoso e lento.
Il fine-tuning consiste nel prendere un modello già addestrato e insegnargli il tuo compito specifico. Il lavoro più duro è già stato fatto. Il modello conosce la lingua e il ragionamento. Gli stai solo mostrando il tuo formato e dominio specifici. È come assumere qualcuno che sa già programmare e insegnargli la tua codebase, invece di insegnare a qualcuno a programmare da zero.
Per il mio caso d'uso, voglio il fine-tuning. Il modello capisce già l'inglese e la pronuncia. Ho solo bisogno che dia un feedback nel mio formato.
Come fare il fine-tuning con Hugging Face
Hugging Face rende tutto questo processo molto più semplice di quanto dovrebbe essere. Ospitano migliaia di modelli pre-addestrati e forniscono librerie Python per ottimizzarli. Non serve un dottorato per farlo.
Ecco il flusso di lavoro di base:
from transformers import AutoProcessor, AutoModelForImageTextToText
# Scarica il modello e il processore
processor = AutoProcessor.from_pretrained("google/gemma-3n-E2B-it")
model = AutoModelForImageTextToText.from_pretrained("google/gemma-3n-E2B-it")
Due righe di codice. Questo scarica un modello di IA da 5,4 miliardi di parametri sul tuo laptop. Il processor converte i tuoi dati nel formato corretto e model è la rete neurale stessa.
Quindi si utilizza la classe Trainer per eseguire effettivamente l'addestramento:
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() # Avvia l'addestramento!
Questo è ottimo perché si saltano tutte le operazioni di basso livello come la scrittura di cicli di discesa del gradiente e la gestione dei checkpoint. Lo svantaggio? Non si capisce bene cosa succede sotto il cofano, specialmente quanta memoria utilizza ogni cosa. Che è esattamente il problema che ho riscontrato.
Il problema della memoria
L'addestramento di un modello linguistico di grandi dimensioni richiede che diverse cose stiano in memoria contemporaneamente:
- Pesi del modello: Per Gemma 3n E2B (5,4B parametri), sono circa 11 GB in float16
- Gradienti: Altri circa 11 GB (stessa dimensione dei pesi)
- Stati dell'ottimizzatore: Adam mantiene 2 stati per parametro, aggiungendo circa 22 GB
- Attivazioni: Risultati intermedi dei calcoli
- Dati del batch: Gli esempi di addestramento in fase di elaborazione
Totale: 50-60 GB per il modello E2B più piccolo, 70-80 GB per E4B. I miei 48 GB di RAM non sarebbero bastati.
Cosa ho provato
Tentativo 1: Fine-tuning completo con E4B
Memoria esaurita durante l'inizializzazione dell'ottimizzatore. Il modello è stato caricato (16 GB), ma l'ottimizzatore non è riuscito ad allocare il suo stato.
Tentativo 2: Dimensioni del batch più piccole
Ridurre batch_size da 2 a 1 ha solo ritardato l'errore OOM (Out Of Memory) al passaggio all'indietro (backward pass).
Tentativo 3: LoRA con E4B
LoRA (Low-Rank Adaptation) addestra solo piccoli livelli adattatori invece dell'intero modello, riducendo i parametri addestrabili da 7,8B a 40M (0,5%). Ma il modello base congelato deve comunque stare in memoria: 59 GB solo per E4B.
Tentativo 4: LoRA con E2B
Questo è quasi riuscito. Il modello E2B più piccolo con LoRA utilizzava 40-45 GB, lasciando 3-8 GB di margine. L'addestramento è iniziato ed è proseguito per alcune ore prima di incontrare errori di memoria durante batch specifici con sequenze più lunghe.
Perché alla fine è fallito
Anche con ottimizzazioni aggressive (LoRA, batch_size=1, nessuna checkpointing dei gradienti), 48 GB non erano sufficienti per un addestramento costante. L'utilizzo della memoria variava in base all'esempio:
- Esempi brevi: 38-42 GB (OK)
- Esempi lunghi: 46-50 GB (crash)
Probabilmente avrei potuto farlo funzionare filtrando gli esempi più lunghi o riducendo max_length a 1024 token, ma a quel punto sto compromettendo le capacità del modello per adattarmi ai vincoli hardware.
Passaggio a Vertex AI
Dopo tre giorni di tentativi falliti, sono passato a Vertex AI di Google Cloud per l'addestramento:
- GPU A100 40GB: Memoria più che sufficiente per il modello E4B con LoRA
- Infrastruttura gestita: Niente debug della memoria, niente crash
- Costo: Circa 3-5 $ per un ciclo di addestramento completo (3 epoche)
- Tempo: 2-3 ore rispetto alle 6-9 ore che ho tentato in locale
Scriverò un post separato sul processo di addestramento di Vertex AI, che includerà:
- Impostazione di job di addestramento personalizzati con Gemma 3n
- Caricamento dei dataset su Cloud Storage
- Monitoraggio dell'addestramento con TensorBoard
- Deployment del modello addestrato sugli endpoint di Vertex AI
Quando è possibile addestrare in locale?
L'addestramento locale su Apple Silicon è fattibile se:
- Il tuo modello è più piccolo: I modelli da 1-3B parametri con LoRA rientrano comodamente nei 48 GB
- Hai più RAM: L'M4 Max con 128 GB gestirebbe facilmente Gemma 3n E2B
- Utilizzi la quantizzazione: La quantizzazione a 4 o 8 bit riduce ulteriormente la memoria ma aggiunge complessità
Lezioni chiave
- La memoria è il collo di bottiglia: 48 GB di memoria unificata sembrano tanti finché non si tenta di addestrare modelli con più di 5B parametri. Solo gli stati dell'ottimizzatore possono raddoppiare o triplicare i requisiti di memoria
- LoRA aiuta ma non è una bacchetta magica: Riduce i parametri addestrabili del 99%, ma i pesi del modello base congelato devono comunque stare in memoria durante i passaggi in avanti e all'indietro
- L'addestramento su cloud è conveniente: 3-5 $ per un ciclo di addestramento sono ragionevoli rispetto a giorni di debug di errori OOM e esperimenti falliti
- Conosci i limiti del tuo hardware: La memoria unificata di Apple Silicon è eccellente per l'inferenza e l'addestramento di modelli più piccoli (1-3B parametri), ma l'addestramento affidabile di modelli superiori a 5B parametri richiede VRAM dedicata della GPU.