DeepSpeed, la libreria di Microsoft per il deep learning, introduce due aggiornamenti significativi focalizzati sull'ottimizzazione del training di modelli multimodali e sulla riduzione del consumo di memoria.

API Backward PyTorch-Identica

Il primo aggiornamento riguarda un'API backward identica a quella di PyTorch. Questa nuova API semplifica la scrittura di cicli di training per modelli multimodali, che spesso combinano encoder di visione e LLM, consentendo schemi di parallelismo sofisticati con codice più pulito. DeepSpeed gestisce in modo trasparente le ottimizzazioni delle performance. Un esempio è l'abilitazione del parallelismo ibrido disaggregato, che ha portato a un miglioramento del 30% nella velocità di training per modelli di AI multimodali.

L'API backward originale di DeepSpeed richiedeva l'uso di model_engine.backward(loss) anziché il consueto loss.backward() di PyTorch, limitando la flessibilità. La nuova API permette di combinare modelli multipli, definire funzioni di loss separate e gestire tensori non scalari con gradienti personalizzati, mantenendo le ottimizzazioni di DeepSpeed come ZeRO e l'offloading.

Training a Bassa Precisione per Ridurre l'Utilizzo di Memoria

Il secondo aggiornamento introduce un'opzione per mantenere tutti gli stati del modello (parametri, gradienti e stati dell'ottimizzatore) in bassa precisione, come BF16 o FP16. Questo riduce drasticamente l'impronta di memoria, consentendo ai ricercatori di effettuare il fine-tuning di modelli più grandi su hardware con risorse limitate. L'integrazione con torch.autocast garantisce la stabilità numerica durante il training.

Un test ha mostrato una riduzione del 40% della memoria di picco mantenendo la stabilità numerica. I risultati mostrano che il training a bassa precisione con BF16 mantiene una convergenza paragonabile al training con FP32, ma con un notevole risparmio di memoria.

Per chi valuta deployment on-premise, esistono trade-off tra performance e consumo di risorse. AI-RADAR offre framework analitici su /llm-onpremise per valutare questi aspetti.