DeepSeek-V3 da zero : costruire l’architettura in PyTorch

Il sito PyImageSearch pubblica un tutorial dedicato alla costruzione dell’architettura di DeepSeek-V3 con PyTorch, una libreria di deep learning. Quinto episodio di una serie di sei, questo tutorial, firmato Puneet Mangla, assembla i mattoni presentati nelle lezioni precedenti per formare un modello linguistico completo, che trasforma i token, le unità di testo elaborate dal modello, in logit, i punteggi assegnati a ogni token possibile.

Le lezioni precedenti hanno coperto la configurazione del modello, i Rotary Position Embeddings (RoPE), che codificano la posizione dei token, la Multihead Latent Attention (MLA), un meccanismo di attenzione, e la Mixture of Experts (MoE), o miscela di esperti, che ripartisce il calcolo tra diverse sottoreti specializzate. La Multi-Token Prediction (MTP), infine, permette di predire più token contemporaneamente. Secondo il tutorial, migliora l’efficienza dell’addestramento e può accelerare l’inferenza grazie al decoding speculativo.

Ogni blocco del trasformatore combina la MLA e la MoE secondo uno schema detto « pre-norm », in cui la normalizzazione interviene prima di ogni sottostrato. Connessioni residue fanno transitare i segnali direttamente da uno strato all’altro, il che, precisa l’autore, aiuta a risolvere il problema della scomparsa del gradiente nelle reti profonde. Il modello completo incatena poi gli embedding dei token e delle posizioni, uno strato di dropout, i blocchi, una normalizzazione finale, una testa di predizione del vocabolario e, opzionalmente, delle teste di MTP.

Condivisione dei pesi e predizione di più token

Il tutorial presenta anche la condivisione dei pesi tra la matrice di embedding dei token e la testa di predizione. Per un vocabolario di 50.259 token e una dimensione di embedding di 256, questa tecnica permette di risparmiare circa 12,9 milioni di parametri. Secondo l’autore, migliora spesso le prestazioni, poiché gli spazi di input e di output rimangono allineati.

Durante l’addestramento, le uscite delle teste di MTP sono calcolate per diversi orizzonti futuri, poi le loro perdite sono combinate con la perdita principale e con una perdita ausiliaria della MoE, destinata a bilanciare il carico tra gli esperti. Nel codice presentato, la perdita media della MTP è ponderata di 0,3 rispetto alla perdita principale. Le teste di MTP non sono utilizzate durante l’inferenza, cioè la generazione di testo.

Dimensione del modello e impronta di memoria

Secondo l’autore, il modello conta circa 30-35 milioni di parametri ripartiti su sei strati, ovvero 2-3 milioni per strato, ai quali si aggiungono 2-3 milioni di parametri per testa di MTP. Alla precisione FP32, su 32 bit, i soli parametri rappresentano 120-140 megabyte. L’addestramento, con i gradienti e gli stati dell’ottimizzatore AdamW, necessita di circa 500-600 MB, ai quali si aggiungono le attivazioni. Per un batch di 4 esempi di 1.024 token, queste pesano 200-300 MB, ovvero quasi 800 MB in totale, una quantità che il tutorial giudica compatibile con le attuali GPU.

In inferenza, la precisione FP16 riduce la memoria di circa la metà. L’implementazione non ricorre a una cache chiave-valore e ricalcola quindi il contesto a ogni passo di generazione. La prossima lezione della serie riguarderà l’addestramento di questo modello su dati reali.