Ottimizzazione Fine-Tuning dei Modelli Linguistici Utilizzando Solo Passaggi in Avanti

Abstract

Il fine-tuning dei modelli linguistici (LM) ha ottenuto successo in vari task downstream, ma man mano che le dimensioni dei LM aumentano, la retropropagazione richiede una quantità di memoria proibitivamente elevata. I metodi di ordine zero (ZO) possono in linea di principio stimare i gradienti utilizzando solo due passaggi in avanti, ma si teorizza che siano catastroficamente lenti per l'ottimizzazione di modelli di grandi dimensioni. In questo lavoro, proponiamo un ottimizzatore di ordine zero efficiente in termini di memoria (MeZO), adattando il classico metodo ZO-SGD per operare in-place, consentendo così il fine-tuning dei LM con lo stesso footprint di memoria dell'inferenza. Ad esempio, con una singola GPU A100 80GB, MeZO può addestrare un modello con 30 miliardi di parametri, mentre il fine-tuning con retropropagazione può addestrare solo un LM da 2,7B con lo stesso budget. Abbiamo condotto esperimenti completi su diversi tipi di modelli (LM mascherati e autoregressivi), scale di modelli (fino a 66B) e task downstream (classificazione, scelta multipla e generazione). I nostri risultati dimostrano che (1) MeZO supera significativamente l'apprendimento in-context e il probing lineare; (2) MeZO raggiunge prestazioni comparabili al fine-tuning con retropropagazione su più task, con una riduzione della memoria fino a 12x; (3) MeZO è compatibile sia con tecniche di tuning a parametri completi che con tecniche efficienti come LoRA e prefix tuning; (4) MeZO può ottimizzare efficacemente obiettivi non differenziabili (ad esempio, massimizzare l'accuratezza o l'F1). Supportiamo i nostri risultati empirici con intuizioni teoriche, evidenziando come un adeguato pre-training e prompt di task consentano a MeZO di effettuare il fine-tuning di modelli enormi, nonostante le analisi classiche ZO suggeriscano il contrario.

English

Fine-tuning language models (LMs) has yielded success on diverse downstream tasks, but as LMs grow in size, backpropagation requires a prohibitively large amount of memory. Zeroth-order (ZO) methods can in principle estimate gradients using only two forward passes but are theorized to be catastrophically slow for optimizing large models. In this work, we propose a memory-efficient zerothorder optimizer (MeZO), adapting the classical ZO-SGD method to operate in-place, thereby fine-tuning LMs with the same memory footprint as inference. For example, with a single A100 80GB GPU, MeZO can train a 30-billion parameter model, whereas fine-tuning with backpropagation can train only a 2.7B LM with the same budget. We conduct comprehensive experiments across model types (masked and autoregressive LMs), model scales (up to 66B), and downstream tasks (classification, multiple-choice, and generation). Our results demonstrate that (1) MeZO significantly outperforms in-context learning and linear probing; (2) MeZO achieves comparable performance to fine-tuning with backpropagation across multiple tasks, with up to 12x memory reduction; (3) MeZO is compatible with both full-parameter and parameter-efficient tuning techniques such as LoRA and prefix tuning; (4) MeZO can effectively optimize non-differentiable objectives (e.g., maximizing accuracy or F1). We support our empirical findings with theoretical insights, highlighting how adequate pre-training and task prompts enable MeZO to fine-tune huge models, despite classical ZO analyses suggesting otherwise.