Domyn Logo Glyph
Domyn Research
Technical Team
Pubblicato il
12
November
2026
|
30 min
|

Addestrare un MoE da 1.000 miliardi di parametri sull’infrastruttura pubblica europea con PyTorch e TorchTitan

In questo articolo
Consigliati
No items found.

In Domyn abbiamo avviato di recente l’addestramento di un modello di frontiera nell’ambito della European Frontier AI Grand Challenge, insieme a Fraunhofer IAIS e Fraunhofer IIS. La prima sfida è stata capire se e come fosse possibile addestrare un modello da mille miliardi di parametri sui supercomputer pubblici europei, generalmente ottimizzati per i carichi di lavoro scientifici.

Questo articolo descrive nel dettaglio come ci siamo riusciti. Riteniamo che il risultato rappresenti, di per sé, una svolta per la posizione dell’Europa nella corsa all’AI. Mentre si costruiscono AI factory dedicate, dimostriamo che i supercomputer pubblici europei possono già oggi raggiungere la scala di frontiera.

Nel giugno 2026, il consorzio EUROPA, guidato da Domyn con Fraunhofer IAIS e Fraunhofer IIS, si è aggiudicato la Frontier AI Grand Challenge. L’obiettivo del progetto è creare un modello di frontiera aperto, addestrato su infrastruttura pubblica europea e rilasciato insieme ai pesi e alle metodologie di addestramento. Il consorzio ha ottenuto l’accesso al 2,5% della capacità di calcolo di EuroHPC sui supercomputer pubblici europei.

Per allinearci all’attuale scala di frontiera, abbiamo deciso di addestrare un modello linguistico Mixture-of-Experts (MoE) da circa 1.000 miliardi di parametri totali. Tra i supercomputer europei di EuroHPC, per il pretraining abbiamo scelto di concentrarci su Jupiter, il supercomputer exascale dell’EuroHPC Joint Undertaking, ospitato presso il Forschungszentrum Jülich.

Finora modelli di queste dimensioni sono stati addestrati soprattutto su cluster dedicati di generazione Blackwell. Questo hardware offre una banda più elevata tra le GPU, più memoria per GPU e maggiore libertà nella scelta dello schema di parallelismo. Il Booster di Jupiter, invece, è un sistema di generazione Hopper condiviso da molti utenti. Accediamo al sistema come semplici utenti, quindi non possiamo modificare il kernel, l’interconnessione, la configurazione della memoria dei nodi né il software di sistema. I nodi sono stati configurati per la simulazione scientifica, e il loro equilibrio tra memoria e rete è diverso da quello di un cluster progettato per l’addestramento. Sulla stessa macchina operano centinaia di progetti diversi.

Di seguito descriviamo la configurazione di addestramento che usiamo in queste condizioni e le modifiche che hanno portato il throughput da 90 a circa 1.300 token/s per GPU su 384 GPU GH200 (96 nodi). Forniamo dettagli sufficienti perché altri utenti di Jupiter e di sistemi EuroHPC simili possano riutilizzarla. Quando abbiamo riscontrato problemi nel software condiviso, li abbiamo corretti direttamente in PyTorch e TorchTitan, e abbiamo segnalato i problemi di sistema allo Jülich Supercomputing Centre (JSC). L’elenco completo è nella sezione Contributi open source.

‍

Ambito
Questo articolo descrive la configurazione di addestramento attuale, nel momento in cui scriviamo. La configurazione è ancora in evoluzione, e alcuni dettagli potrebbero non essere più attuali al momento della lettura.

La Figura 1 mostra il miglioramento del throughput ottenuto con ciascuna delle modifiche descritte nell’articolo, a partire dalla configurazione di partenza. Passando il cursore su una riga se ne visualizzano tempo per step, memoria e tentativi di riallocazione dell’allocator.

Throughput per modifica, da 90 a 1.293 token/s per GPU
Throughput a regime (mediana degli step 4-10, rank 0) con l’introduzione progressiva delle modifiche. La configurazione di partenza usa sequenze da 8.192 token. La lunghezza delle sequenze arriva a 65.536 token grazie all’adozione dell’ottimizzatore Muon. Le esecuzioni rappresentate in trasparenza si sono interrotte in anticipo.

Date le caratteristiche dell’hardware, descritte più avanti, la maggior parte delle modifiche si è concentrata sulla riduzione del tempo che le GPU trascorrono in attesa della rete. Lo mostra la figura seguente, che per ogni esecuzione analizzata con il profiler ripartisce lo step tra calcolo, attesa dello scambio di dati e inattività.

Tempo per step per attività
Ripartizione dello step del rank 0 per ogni esecuzione analizzata con il profiler: calcolo sovrapposto allo scambio di dati, solo calcolo, attesa dello scambio di dati e inattività.

Sistema e modello

Hardware

Jupiter è il primo supercomputer europeo di classe exascale. Il suo modulo Booster conta 5.884 nodi di calcolo, ciascuno con quattro superchip NVIDIA GH200 Grace Hopper. Ogni superchip comprende i seguenti componenti:

  • Grace: una CPU Arm con 72 core e 120 GB di LPDDR5X a 512 GB/s.‍
  • Hopper: una GPU di classe H100 con 96 GB di HBM3 a 4 TB/s.‍
  • NVLink-C2C: un collegamento a 900 GB/s tra Grace e Hopper.

All’interno di un nodo, le coppie di GPU sono collegate tramite NVLink 4 a 300 GB/s (150 GB/s per direzione), e le quattro CPU Grace sono collegate a 100 GB/s per direzione. Ogni nodo dispone di quattro adattatori InfiniBand NDR200 per il traffico tra nodi, ciascuno con una velocità di 200 Gbit/s. Complessivamente offrono 800 Gbit/s (100 GB/s) per nodo, pari a circa 25 GB/s per GPU.

Nel corso dell’articolo faremo riferimento alle seguenti larghezze di banda:

LinkBandwidthRelative to InfiniBand
Grace ↔ Hopper (NVLink-C2C)900 GB/s36×
GPU ↔ GPU, same node (NVLink 4)150 GB/s per direction6×
GPU ↔ outside the node (InfiniBand)~25 GB/s per GPU1×
Larghezze di banda di un nodo Jupiter, in scala
Lo spessore delle linee è proporzionale alla larghezza di banda. Le linee sottili che escono dal nodo sono i collegamenti InfiniBand, su cui passa tutto il traffico tra nodi.

Software

Per il modello EUROPA abbiamo scelto di adottare PyTorch e TorchTitan (Liang et al., 2025), la libreria nativa di PyTorch per il pretraining di LLM su larga scala. Abbiamo scelto TorchTitan per i seguenti motivi:

  • È il framework di addestramento che il nostro team conosce meglio.
  • Il codice è compatto e semplice da adattare alla nostra configurazione.
  • Il progetto ha una comunità di sviluppo attiva, e le pull request vengono revisionate e integrate rapidamente.

Modello

L’architettura si basa su MiMo-V2.6 (Core Team et al., 2026), che estende l’architettura MiMo-V2-Flash (Core Team, Xiao, et al., 2026). La tabella riporta la configurazione che usiamo attualmente per gli esperimenti. Alcuni valori potrebbero cambiare prima dell’avvio dell’addestramento vero e proprio.

PropertyValue
Total parameters951.9B (32.8B active per token)
Layers74
Hidden size6,144
Experts384 routed, top-8
Expert hidden size1,792
Vocabulary256,000
Sequence length65,536 tokens
Hardware96 nodes, 384 GH200 superchips
Tokens per step384 × 65,536 ≈ 25M
OptimizerMuon (all weight matrices except the router), AdamW (embeddings, LM head, router, norms)
  1. Liang, W., Liu, T., Wright, L., Constable, W., Gu, A., Huang, C.-C., Zhang, I., Feng, W., Huang, H., Wang, J., Purandare, S., Nadathur, G., & Idreos, S. (2025). TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training. https://arxiv.org/abs/2410.06511

‍

Vincoli e schema di parallelismo

Vincoli

Come accennato, la potenza di calcolo grezza delle GPU non era il fattore limitante. Nel pretraining di un MoE da mille miliardi di parametri, le difficoltà riguardavano altri due elementi del sistema.

Banda tra nodi

Il traffico InfiniBand tra nodi è 6 volte più lento di quello tra GPU dello stesso nodo, e 36 volte più lento di quello tra GPU e CPU. In uno step di addestramento distribuito su 384 GPU per replica data-parallel, come nella nostra configurazione, questo collegamento viene attraversato molte volte. Il traffico su InfiniBand deve quindi essere ridotto oppure avvenire in parallelo al calcolo delle GPU.

Capacità della memoria host

Nei nodi di login di Jupiter, GH200 dispone di fino a 480 GB di LPDDR5X a 384 GB/s per Grace. I nodi di calcolo, invece, hanno 120 GB a 512 GB/s: una configurazione che privilegia la banda rispetto alla capacità. È adatta ai carichi di lavoro di simulazione, limitati dalla banda di memoria, ma per l’addestramento limita la quantità di dati che ogni nodo può elaborare e conservare tra HBM e memoria di Grace.

La memoria effettivamente utilizzabile, infatti, è inferiore a 120 GB. Su GH200, Linux espone l’HBM di ogni Hopper come un nodo NUMA senza CPU. Ogni nodo riporta 858 GiB di memoria, di cui 456 GiB di memoria Grace. Del totale, il kernel ne riserva circa 24 GiB, e restano così circa 108 GiB di memoria host utilizzabile per GPU. Altre perdite sono descritte nella sezione Perdite di memoria host.

Schema di parallelismo

Per un MoE da mille miliardi di parametri, scegliere lo schema di parallelismo più adatto all’architettura hardware è decisivo per raggiungere il throughput necessario in produzione. Inizialmente abbiamo valutato il pipeline parallelism, ma alla luce delle considerazioni che seguono ci siamo concentrati su FSDP ed expert parallelism, dedicando il nostro lavoro alla loro ottimizzazione: un impegno che alla fine ha dato i suoi frutti. Poiché lo schema adottato determina tutte le ottimizzazioni successive, non possiamo dire se uno schema basato sul pipeline parallelism sarebbe stato più veloce. La decisione si è basata su tre principi:

  1. Expert parallelism all’interno del nodo. L’all-to-all dei token di un layer MoE si trova sul percorso critico tra il router e gli esperti ed è difficile da sovrapporre al calcolo. Su NVLink il tempo che richiede è accettabile, su InfiniBand è troppo alto. Abbiamo quindi limitato l’expert parallelism a 4, con un gruppo per nodo.
  2. Step brevi. Un pipeline parallelism profondo richiede molti microbatch in ogni step per contenere la bubble, e questo allunga gli step. Con step più lunghi diventa più difficile individuare opportunità di ottimizzazione e problemi che emergono su larga scala, e i tempi di riscontro si allungano.
  3. Batch ampio. Un batch più ampio riduce il rumore dei gradienti e assegna a ogni GPU più lavoro per step, ammortizzando costi fissi come il lancio dei kernel e lo step dell’ottimizzatore.

Sulla base di queste considerazioni, abbiamo optato per FSDP (fully sharded data parallelism) a 384 vie, con expert parallelism a 4 vie all’interno di ciascun nodo. Tutte le esecuzioni descritte in questo articolo usano una sola replica data-parallel (DP = 1), quindi un unico gruppo FSDP copre l’intero job. In pratica, possiamo eseguire fino a 4 repliche data-parallel su 1.536 GPU senza effetti apprezzabili sul throughput.

DPNodesGPUsBatch size (sequences of 65,536 tokens)Throughput vs DP = 1
196384384 (≈25M tokens)100%
2192768768 (≈50M tokens)99.3%
43841,5361,536 (≈101M tokens)89.4%
Schema di parallelismo
96 nodi × 4 GPU. Sono rappresentati solo i primi nodi e l’ultimo. Il colore indica il blocco di 96 esperti che risiede su ciascuna GPU. Passando il cursore su una GPU se ne visualizzano i gruppi. Il nodo di appartenenza è il suo gruppo di expert parallelism (all-to-all su NVLink). Le 96 GPU dello stesso colore formano il gruppo FSDP tra cui gli esperti vengono suddivisi tramite InfiniBand. I parametri densi sono suddivisi su tutte le 384 GPU. La riga con il contorno è il gruppo FSDP. Eventuali repliche data-parallel (DP) in più comporterebbero altre righe come quella in trasparenza, ciascuna una copia completa del gruppo FSDP. Tutte le esecuzioni descritte in questo articolo usano una sola replica (DP = 1).

Lunghezza delle sequenze, ottimizzatore e kernel MoE

Configurazione di partenza

FSDP riusciva a svolgere in parallelo al calcolo solo una parte dello scambio di dati. Spesso il calcolo di un layer terminava prima che arrivassero i pesi del successivo, e la GPU restava inattiva. Nella configurazione di partenza (8.192 token per sequenza), uno step profilato sul rank 0 durava in totale 83,6 s: 14,3 s di calcolo e 68,3 s trascorsi in attesa della rete.

Timeline di GPU e CPU nella configurazione di partenza
Attività di GPU e CPU sul rank 0 durante lo step profilato della configurazione di partenza. L’all-gather è in esecuzione quasi senza sosta, mentre il calcolo occupa brevi intervalli. Ingrandendo il forward pass si vedono i calcoli di ciascun layer.

Poiché in FSDP lo scambio di dati riguarda solo i pesi, e il volume di pesi da raccogliere per ogni layer non dipende dal numero di token, con sequenze più lunghe la GPU svolge più calcolo a parità di dati trasmessi. Abbiamo quindi allungato le sequenze fin dove la memoria lo permetteva. Tuttavia, anche con l’activation checkpointing della configurazione di partenza, già una sequenza da 12.288 token superava la capacità dell’HBM.

Activation checkpointing e offload su Grace

Sequenze più lunghe richiedono più memoria per le attivazioni. In un primo momento abbiamo salvato solo le attivazioni in uscita da ciascun layer, ricalcolando durante il backward pass tutto ciò che avviene al suo interno. Il ricalcolo completo aumenta il carico di lavoro, ma la GPU restava inattiva per gran parte dello step, quindi il lavoro aggiuntivo non ne ha allungato la durata. Con questa modifica, la lunghezza delle sequenze è arrivata a 32.768 token e il throughput a 370 token/s per GPU, 4 volte quello della configurazione di partenza.

Con sequenze da 65.536 token, le sole attivazioni in uscita dai layer occupano 74 layers × 65,536 tokens × 6,144 × 2 bytes ≈ 55 GiB per GPU: troppo per l’HBM, che deve ospitare anche i pesi e le altre attivazioni. Abbiamo quindi deciso di spostarle su Grace tramite il collegamento NVLink-C2C da 900 GB/s. Ciascuna di queste attivazioni viene copiata su Grace dopo il forward pass, e riportata sulla GPU prima del ricalcolo durante il backward pass. Abbiamo impostato a 0,86 il parametro offload_fraction, che stabilisce la quota di layer da trasferire, e quindi il compromesso tra memoria HBM e memoria di Grace. Con l’offload, la lunghezza delle sequenze è arrivata a 49.152 token e il throughput a 508 token/s per GPU.

Schema di offload delle attivazioni
Schema semplificato a 16 layer anziché 74. Nel forward pass ogni attivazione in uscita viene copiata su Grace; nel backward pass torna su Hopper un layer prima di essere ricalcolata. I layer esclusi dalla quota di offload mantengono la propria attivazione in uscita nell’HBM.

Abbiamo sviluppato e integrato in TorchTitan gli hook per l’offload degli activation checkpoint, insieme a un offloader su CPU minimale fornito come esempio documentato. Su questi hook abbiamo poi costruito la strategia di offload descritta qui.

Muon con Newton-Schulz suddiviso per esperti

Lo stato dell’ottimizzatore occupa una parte significativa della memoria. Usiamo Muon (Jordan et al., 2024) per tutte le matrici dei pesi tranne il router, circa il 99,5% dei parametri, e AdamW per embedding, LM head, router e layer di normalizzazione. Muon memorizza un solo stato in fp32 per parametro (il momentum), mentre Adam ne memorizza due. Fare a meno del secondo stato ha permesso di far rientrare in memoria sequenze da 65.536 token. Muon è risultato più efficiente di AdamW nell’addestramento di modelli di grandi dimensioni (Liu et al., 2025), ed è stato usato per modelli MoE da mille miliardi di parametri (Kimi Team, 2025; Kimi Team et al., 2026).

Con Muon, ogni aggiornamento viene ortogonalizzato con alcune iterazioni di Newton-Schulz, che operano sulla matrice intera. Nella nostra configurazione FSDP, però, le matrici sono divise in 384 shard. Un’implementazione di base ricomporrebbe ciascuna matrice su un unico rank, applicherebbe Newton-Schulz e poi redistribuirebbe il risultato, con un aumento significativo dello scambio di dati.

FlexShard, ora integrato in TorchTitan, ci permette di stabilire quale rank possiede quali parametri. Grazie a questa possibilità, abbiamo distribuito gli esperti interi, senza suddividerli: ogni rank possiede matrici di esperti complete ed esegue le relative iterazioni di Newton-Schulz senza alcuno scambio di dati. Le matrici dense, come le proiezioni dell’attention, restano invece suddivise e vengono raccolte sul rank a cui sono assegnate. Questa configurazione consente di sfruttare ancora di più la sovrapposizione tra calcolo e scambio di dati: ogni rank può elaborare i propri esperti mentre le matrici dense sono in transito, così lo scambio di dati si sovrappone completamente al calcolo.

Un passaggio chiave è stato spostare l’esecuzione dell’ottimizzatore sul processore Grace. A differenza delle altre operazioni dell’ottimizzatore, però, Newton-Schulz è computazionalmente oneroso e trae grande vantaggio dall’accelerazione su GPU. La figura seguente confronta Newton-Schulz su Grace e su Hopper, con le dimensioni dei nostri esperti.

Tempo di Newton-Schulz su Grace e su Hopper
Tempo mediano per ortogonalizzare la matrice di un esperto (scala logaritmica) e tempo risultante per rank per step. Grace impiega circa 35 s per step usando tutti i 72 core, Hopper 0,2 s.

Durante l’addestramento, con Newton-Schulz su Grace, lo step dell’ottimizzatore richiedeva complessivamente 41,7 s. Abbiamo quindi sviluppato e integrato in TorchTitan il supporto per eseguire Newton-Schulz sulla GPU quando i parametri sono in offload. Con questa modifica, lo step dell’ottimizzatore è sceso a 2,2 s, e il throughput è salito da 441 a 597 token/s per GPU (+35%).

Su un dispositivo CPU come Grace, stream ed eventi di PyTorch non hanno alcun effetto. Di conseguenza, la sovrapposizione che FlexShard garantisce su CUDA va persa senza alcun avviso. Abbiamo modificato il codice lato CPU per eseguire in modo asincrono gli all-to-all delle matrici dense, anziché bloccare il thread principale. Così lo step dell’ottimizzatore è sceso a 1,6 s.

Step dell’ottimizzatore con Newton-Schulz su Grace e su Hopper
Step profilati sul rank 0 da 80 s in poi, su un asse temporale comune. Reimpostando lo zoom si vedono gli step completi. Con Newton-Schulz su Grace, la GPU resta inattiva per circa 34 s dopo il backward pass, come mostra la riga della CPU. La riga gloo mostra gli ultimi 10 s circa: un all-to-all delle matrici dense lato CPU, che deve attendere i rank più lenti. Con Newton-Schulz su Hopper, l’ottimizzatore occupa una piccola frazione dello step.

Scheduling dinamico dei tile in SonicMoE

Per gli esperti usiamo SonicMoE (Guo et al., 2026), un kernel MoE fuso, basato sui kernel CuTe DSL di QuACK. Esegue le GEMM degli esperti direttamente sul layout di dispatch dei token, senza doverli prima copiare e riordinare. Su una singola GH200, con dimensioni di test, ha ridotto il picco di memoria del layer da 12,5 a 5,4 GiB e il tempo di forward e backward da 213 a 153 ms.

I kernel di comunicazione vengono eseguiti sugli stessi SM della GPU dei kernel di calcolo. Di conseguenza, finché è in esecuzione, un all-gather di FSDP occupa alcuni SM con i thread block di NCCL. Le GEMM di SonicMoE usano uno scheduler persistente statico, che al lancio assegna a ciascuno dei 132 thread block persistenti un elenco fisso di tile. Se l’SM di un blocco è occupato dall’all-gather, quel blocco inizia a elaborare i propri tile solo quando l’all-gather termina, e la GEMM si conclude soltanto quando finisce l’ultimo blocco. In un’esecuzione, la stessa up-projection richiedeva 55-60 ms senza altre operazioni in corso e circa 100 ms in concomitanza con un all-gather.

QuACK offre anche uno scheduler dinamico, in cui ogni blocco, quando è libero, prende in carico il tile successivo tramite un contatore atomico condiviso: così i blocchi già in esecuzione elaborano anche i tile di quelli non ancora partiti. In SonicMoE, però, lo scheduler statico è fisso nel codice, quindi abbiamo racchiuso le sue chiamate GEMM in un wrapper che attiva quello dinamico. Abbiamo poi misurato le prestazioni di un layer MoE, con le dimensioni usate nel nostro addestramento, su una singola GH200, con e senza un kernel fittizio che occupa 40 SM.

Scheduling statico e dinamico dei tile in SonicMoE
Tempo mediano per ogni GEMM degli esperti e per forward e backward dell’intero layer. Con 40 SM occupati, lo scheduler dinamico riduce il tempo del layer da 304 a 217 ms.

Con 40 SM occupati, lo scheduler dinamico rende le GEMM più veloci del 17-45%, e il tempo del layer scende da 304 a 217 ms. Su una GPU libera è più veloce del 3-32% per GEMM e dell’11% sul layer. Gli output sono identici bit per bit, salvo gli arrotondamenti in fp32 nel gradiente dei punteggi di routing. In addestramento lo abbiamo attivato sulle quattro GEMM degli esperti in cui la dimensione dei token varia da un esperto all’altro. Il tempo per step è sceso da 108,7 a 103,8 s, con un ulteriore guadagno di throughput del 5%.

Abbiamo chiesto agli autori di SonicMoE perché lo scheduler predefinito sia quello statico. Limitare gli SM usati da NCCL affronta la stessa contesa, riducendo il numero di SM occupati dalla comunicazione.

‍

Bilanciamento del carico tra i rank

Le 384 GPU procedono in sincronia, quindi è il rank più lento a determinare la durata dello step. Due fonti di squilibrio rendevano alcuni rank più lenti di altri.

Bilanciamento dei dati in base al costo dell’attention

Il costo dell’attention cresce con il quadrato della lunghezza del documento. Con righe da 65.536 token, un rank potrebbe ricevere un unico documento da 60.000 token e un altro cento documenti brevi. Entrambi elaborano lo stesso numero di token, ma il primo svolge molto più lavoro di attention, e il suo calcolo richiede quindi più tempo.

Per attenuare questo effetto, bilanciamo i dati offline. Pre-tokenizziamo il corpus e raggruppiamo i documenti, in ordine casuale, in righe da 65.536 token. Un documento che non entra in una riga prosegue nella successiva; per ogni riga calcoliamo poi il costo dell’attention come somma dei quadrati delle lunghezze dei segmenti, dove un segmento è la parte di un documento contenuta nella riga. Ordiniamo le righe in fasce di costo ristrette, le mescoliamo all’interno di ogni fascia e riempiamo ogni step con righe consecutive secondo quest’ordine, una per rank. Le righe di uno step hanno quindi un costo quasi uguale. Infine, mescoliamo l’ordine degli step.

Sui 1.057 step, il rank più oneroso di uno step svolgeva in mediana 7,9 volte il lavoro di attention medio prima del bilanciamento, e 1,15 volte dopo.

Costo massimo dell’attention per rank, step per step
Per ogni step, il costo dell’attention più alto tra i rank (somma dei quadrati delle lunghezze dei segmenti) diviso per la media dello step. L’ordine di arrivo usa le righe nel loro ordine originale. Il piano bilanciato riunisce nello stesso step righe di costo simile.

In addestramento, il bilanciamento dei dati ha ridotto il tempo per step da 103,8 a 91,7 s, con un guadagno di throughput del 13%.

Il bilanciamento modifica la distribuzione dei dati all’interno di uno step. I documenti lunghi vengono raggruppati negli stessi step, invece di arrivare mescolati come accadrebbe naturalmente. Se un rank riceve una riga con un unico documento da 65.000 token, ogni altro rank in quello step riceve una riga simile, e questo può influire sulla dinamica dell’addestramento. Ci aspettiamo che l’effetto sia limitato ai documenti lunghi, e prevediamo di compensarlo con una fase dedicata al contesto lungo.

Routing degli esperti limitato per gruppo

La seconda fonte di squilibrio è il router del MoE. Quando invia più token agli esperti di una GPU, quella GPU diventa la più lenta del suo nodo. In esecuzioni precedenti senza bilanciamento dei dati, lo squilibrio tra esperti costava circa 16-17 s per step, circa il 20% dello step.

Abbiamo provato a rendere il carico degli esperti perfettamente uniforme. I 384 esperti sono divisi in quattro gruppi da 96, uno per rank di expert parallelism. Il router seleziona lo stesso numero di esperti da ogni gruppo (due per gruppo con routing top-8). Ogni GPU riceve così lo stesso numero di coppie token-esperto in ogni layer. Sommata al bilanciamento dei dati, questa modifica ha ridotto il tempo per step da 91,7 a 85,9 s, con un guadagno di throughput del 7%. Inoltre, il picco di HBM è sceso da 87,2 a 81,3 GiB, perché nessuna GPU deve più riservare memoria per un eccesso di token.

Il vincolo, però, riduce la libertà di routing, e questo può incidere sull’efficacia dell’addestramento. Ad esempio, un token i cui tre esperti con il punteggio più alto appartengono allo stesso gruppo può essere instradato solo verso due di essi. Selezionare più esperti per gruppo aumenta la probabilità che ogni token raggiunga i suoi esperti migliori mantenendo il carico uniforme. Le nostre prime esecuzioni indicano che il vincolo può effettivamente penalizzare l’addestramento, e stiamo ancora valutando se mantenerlo o allentarlo con una strategia intermedia.

‍

Ridurre il volume dello scambio di dati

Oltre una certa lunghezza, sequenze più lunghe non miglioravano più il throughput. Quando le attivazioni mettono sotto pressione la memoria, l’allocator CUDA inizia a ritentare le allocazioni e a frammentare la memoria, e la frammentazione finisce per costare più tempo di quanto se ne risparmi con sequenze più lunghe. Da quel punto, l’unica opzione rimasta era inviare meno dati. Prima delle modifiche descritte in questa sezione, uno step profilato passava 39,5 dei suoi 98,7 s ad attendere la rete (vedi la ripartizione dello step nel grafico Tempo per step per attività, e all-gather e reduce-scatter occupavano la GPU per oltre 50 s.

All-gather dei pesi degli esperti in FP8

Dall’analisi con il profiler, l’all-gather dei pesi degli esperti risulta la voce più onerosa dello scambio di dati in uno step: circa 7 GB per rank in bf16, cioè 28 GB per nodo. Raccogliere i pesi in FP8 dimezzerebbe questo volume.

A questo scopo, i pesi vengono divisi in blocchi da 128 valori, e ogni blocco viene scalato all’intervallo dell’FP8 (E4M3) e convertito. Ogni rank invia i valori FP8 più un fattore di scala in fp32 per blocco; i riceventi li riconvertono poi in bf16 prima del calcolo degli esperti. Il forward pass usa quindi i pesi a precisione ridotta, con un errore per valore fino a circa il 6% del massimo del blocco. I master weight, i gradienti e lo stato dell’ottimizzatore restano a precisione piena. L’approccio è simile a quello usato per DeepSeek-V3 (DeepSeek-AI, 2024).

Con questa strategia, il tempo di all-gather sulla GPU è sceso da 53,6 a 35,1 s per step, e il throughput è salito da 763 a 898 token/s per GPU (+18%). I primi risultati non mostrano effetti apprezzabili sull’addestramento.

Poiché i pesi arrivano già in FP8, anche le GEMM degli esperti potrebbero essere eseguite in FP8, con un’ulteriore piccola perdita di precisione, come nell’addestramento MoE in FP8 di DeepSeek-V3. Potremmo adottare questa soluzione come ottimizzazione aggiuntiva.

All-to-all dei gradienti in BF16 con accumulo in FP32

Con gli all-gather in FP8, lo scambio di dati del forward pass si sovrappone completamente al calcolo. Il backward pass, invece, restava limitato dallo scambio di dati a causa della riduzione dei gradienti.

Nella configurazione standard di FSDP a precisione mista, i gradienti vengono calcolati in bf16 e ridotti in fp32. Prima del reduce-scatter vengono convertiti in fp32, quindi l’anello accumula in fp32 e anche i gradienti vengono trasmessi in fp32.

Solo la somma, però, richiede l’fp32; il trasferimento può avvenire in bf16. Partendo da questa osservazione, abbiamo deciso di inviare i gradienti in bf16 tramite un all-to-all, in cui ogni rank invia a ciascun altro rank la porzione del proprio gradiente di cui quel rank è proprietario. Il proprietario converte le porzioni in fp32 e le somma localmente. Il risultato coincide con quello del reduce-scatter in fp32 a meno dell’ordine di somma, e nei nostri test su 4 rank era identico bit per bit. L’all-to-all invia metà dei byte e dimezza il buffer dei gradienti mantenuto tra un layer e l’altro.

Riduzione dei gradienti: reduce-scatter in fp32 e all-to-all in bf16
Quattro secondi del backward pass sul rank 0. Nell’esecuzione in alto, l’anello del reduce-scatter in fp32 è occupato quasi senza sosta. In quella in basso, i gradienti degli esperti vengono inviati con un all-to-all in bf16, e il reduce-scatter rimanente (parametri densi) è breve.

Il tempo della GPU nelle operazioni collettive sui gradienti è sceso da 50,7 a 33,2 s per step (10,4 s di reduce-scatter e 22,8 s di all-to-all). Il tempo trascorso in attesa della rete si è ridotto da 27,3 a 13,9 s, e il throughput è salito da 898 a 1.022 token/s per GPU (+14%). Con tutte le altre modifiche di questa sezione applicate, disattivare l’all-to-all in bf16 riduce il throughput da 1.210 a 1.004 token/s per GPU.

Né gli all-gather in FP8 né i gradienti in bf16 modificano la loss in modo misurabile. In una serie di esecuzioni di confronto, con lo stesso ordine dei dati e la stessa inizializzazione, dopo dieci step le loss differiscono di meno di 0,01.

Resharding locale al nodo dei pesi densi

Con FSDP, i pesi di ogni layer vengono liberati dopo il forward pass e raccolti di nuovo nel backward pass per il ricalcolo. Ogni nuovo gather nel backward pass compete per la banda InfiniBand con il reduce-scatter dei gradienti.

Gli esperti costituiscono la maggior parte del modello e non possono restare in memoria in quella fase. I parametri densi (proiezioni dell’attention, router e layer di normalizzazione), invece, possono restare sulla GPU se suddivisi. Dopo il forward pass, suddividiamo i pesi densi di ogni layer in quattro parti all’interno del nodo, anziché in 384 parti su tutto il job, così il nuovo gather nel backward pass avviene su NVLink anziché su InfiniBand. Questo occupa circa 55 MiB di HBM per layer, circa 4 GiB in totale. Al contrario, mantenere i pesi densi interi (circa 16 GiB) provocava tentativi di riallocazione con la nostra lunghezza di sequenza, con una conseguente perdita di throughput.

La modifica ha richiesto poco codice nuovo. TorchTitan registra i parametri su due mesh: una mesh densa che copre tutti i rank e una mesh sparsa per gli esperti, la cui dimensione di expert parallelism corrisponde al nodo. Abbiamo impostato il resharding dei parametri densi dopo il forward pass (post_forward_mesh_info) sul gruppo locale al nodo, e allentato alcuni controlli di FSDP più rigidi del necessario.

Da sola, questa modifica ha migliorato il throughput del 2% (da 1.022 a 1.038 token/s per GPU). I log dell’esecuzione mostrano che i 4 GiB aggiuntivi provocavano tentativi di riallocazione su tutti i 384 rank.

I pool dell’allocator sugli stream secondari di FSDP

Dopo un aggiornamento del nostro stack software, l’allocator CUDA ha iniziato a ritentare le allocazioni su GPU con un utilizzo della memoria intorno al 70%. Ogni nuovo tentativo libera l’intera cache dell’allocator e richiede una sincronizzazione completa del dispositivo, un’operazione lenta. Al momento dei tentativi, 25-30 GiB di HBM erano liberi, mappati e presenti nella cache dell’allocator, ma l’allocator non riusciva a usarli.

La causa è nel modo in cui FSDP sovrappone lo scambio di dati al calcolo.Una GPU esegue il lavoro in stream, cioè code che procedono in parallelo. FSDP usa uno stream di calcolo e quattro stream secondari:

StreamWork
ComputeForward and backward pass, and the copy of gathered weights into each parameter
All-gather copy-inPacking of this rank’s weight shards into one flat buffer
All-gatherNCCL all-gather of that buffer
Reduce-scatterGradient reduce-scatter
All-reduceAll-reduce across replicas, and hooks such as our gradient norm

Il caching allocator di PyTorch mantiene un pool di memoria per ogni stream. La memoria liberata torna al pool dello stream che l’ha allocata, e solo quello stream può riutilizzarla, perché un altro stream avrebbe bisogno di una sincronizzazione per sapere che il precedente utilizzatore ha finito. FSDP alloca ogni buffer sullo stream che lo utilizza per scrivere i dati. Il nostro all-gather in FP8 e il nostro all-to-all in bf16 allocano i propri buffer temporanei all’interno delle chiamate di FSDP, quindi anche questi buffer finivano sugli stream secondari. Ogni buffer viene usato brevemente una volta per layer, ma la cache di ciascuno stream finisce per trattenere i blocchi più grandi che ha allocato.

Abbiamo ricostruito il log degli eventi dell’allocator a partire da uno snapshot della memoria del rank 0. Al picco di memoria attiva, gli stream secondari trattenevano 27,3 GiB, di cui solo 6,4 effettivamente in uso. I restanti 20,9 GiB erano inaccessibili allo stream di calcolo, che disponeva di 4,5 GiB di cache libera. A ogni nuovo tentativo, lo stream di calcolo non riusciva a ottenere un blocco libero per un buffer di SonicMoE da 4 GiB. Ridurre il numero di layer non risolveva il problema, perché diminuiva la memoria attiva ma lasciava invariata la cache inutilizzata.

HBM per stream dell’allocator
Ricostruzione a partire dalla traccia degli eventi del caching allocator CUDA sul rank 0, al picco di memoria attiva. Pieno: memoria occupata da tensori attivi. In trasparenza: memoria in cache dello stream ma inutilizzata, che non può essere riutilizzata dagli altri stream.

Abbiamo introdotto tre modifiche perché questa memoria venga allocata da pool riutilizzabili dallo stream di calcolo:

  1. Rilascio anticipato dei gradienti. Il nostro hook per la norma dei gradienti (vedi il paragrafo Modifiche minori a memoria e latenza viene eseguito su ogni layer. Quando un hook è registrato, FSDP mantiene in memoria il gradiente ridotto di ogni layer fino alla fine del backward pass, perché lo stream del reduce-scatter non lo riutilizzi mentre l’hook lo legge. Con un solo hook non è un problema, ma con 74 hook il costo è di 9 GiB. Noi rilasciamo ogni gradiente due layer più tardi, dopo aver fatto in modo che lo stream del reduce-scatter attenda la fine dell’hook.
  2. All-gather sullo stream di calcolo. FSDP2 prevedeva un flag privato e non documentato, unshard_async_op, che esegue il setup dell’all-gather sullo stream di calcolo e lascia che NCCL lavori sul proprio stream interno. I buffer vengono così presi dal pool dello stream di calcolo e vi ritornano. Per farlo è servita una modalità asincrona del nostro gather in FP8, che ora restituisce un handle che esegue la riconversione quando FSDP lo attende: lo stream di calcolo attende la rete solo quando ha bisogno dei pesi. Abbiamo proposto a TorchTitan un’opzione pubblica per attivare questo flag. Su un modello di test più piccolo, ha ridotto il picco di memoria riservata di circa l’11%, al costo di circa lo 0,7% di throughput.‍
  3. Buffer di ricezione dei gradienti sullo stream di calcolo. Il buffer di ricezione dell’all-to-all in bf16, grande quanto il gradiente, viene allocato dal pool dello stream di calcolo.

Sull’intero job da 384 GPU, con le modifiche descritte nel paragrafo Modifiche minori a memoria e latenza già applicate, il numero di rank con tentativi di riallocazione è sceso da 384 a zero. Il picco di memoria riservata è passato da 88,1 a 74,6 GiB e il throughput da 996 a 1.210 token/s per GPU (+21%). Il volume di dati inviati resta invariato: cambiano solo i pool da cui viene allocata la memoria.

Dopo la revisione della nostra PR su TorchTitan, il flag unshard_async_op sta per diventare un’API pubblica di FSDP (PyTorch #202343).

  1. DeepSeek-AI. (2024). DeepSeek-V3 Technical Report. arXiv Preprint arXiv:2412.19437.

‍

Modifiche minori a memoria e latenza

Singolarmente, le modifiche di questa sezione hanno un effetto ridotto sul tempo per step. Le abbiamo applicate mentre tutti i rank ritentavano le allocazioni I pool dell’allocator sugli stream secondari di FSDP, e il throughput è passato da 1.038 a 996 token/s per GPU. Non abbiamo misurato il loro effetto individuale sul throughput, ma hanno ridotto l’uso della memoria.

Norma dei gradienti durante il backward pass

Il gradient clipping richiede la norma globale dei gradienti. Con l’offload su CPU, TorchTitan la calcola dopo il backward pass sui gradienti già trasferiti su Grace, con circa 0,43 s di _foreach_norm sulla CPU e circa 0,19 s di all-reduce gloo.

Ogni shard dei gradienti si trova sulla GPU quando il suo reduce-scatter si conclude. In quel momento aggiungiamo il quadrato della sua norma a un accumulatore scalare. Dopo il backward pass, un solo all-reduce NCCL dello scalare fornisce la norma globale, e il clipping procede come prima.

Nessun all-gather degli embedding nel backward pass

Di default, FSDP raccoglie di nuovo la tabella degli embedding nel backward pass. Il gradiente degli embedding dipende solo dagli ID dei token in input, che ogni rank possiede già, quindi il gather è superfluo. Nel nostro communicator, inoltre, restava in attesa dietro il reduce-scatter del layer 1, esponendo fino a circa 126 ms per step. La correzione richiede una sola riga:

model.tok_embeddings.set_unshard_in_backward(False)

L’abbiamo proposta per l’integrazione in TorchTitan.

Cache RoPE condivisa

TorchTitan costruisce una cache RoPE per ogni layer, anche se layer con la stessa configurazione ottengono cache identiche. Con 65.536 token, i nostri 37 layer RoPE occupavano 2,75 GiB con copie di due sole cache distinte, abbastanza da provocare un errore di memoria esaurita. Condividere le cache riduce questa occupazione a 0,19 GiB, con un risparmio di 2,56 GiB per GPU.

La condivisione di un unico buffer era già stata proposta, ma la proposta si è arenata per timori legati alla fragilità della soluzione e al pipeline parallelism. Abbiamo presentato una versione più contenuta, che condivide la cache nello stesso punto in cui TorchTitan condivide gli embedding, ed è quindi compatibile con il pipeline parallelism.

‍

Limitare gli SM usati da NCCL

Sovrapporre comunicazione e calcolo significa eseguirli entrambi sulla stessa GPU nello stesso momento. Questa sezione riguarda gli SM che i kernel di comunicazione sottraggono ai kernel di calcolo.

NCCL esegue ogni operazione collettiva come un kernel GPU con un cooperative thread array (CTA, il termine di NVIDIA per un thread block) per canale. Ogni CTA occupa un SM per tutta la durata della collettiva e passa gran parte di quel tempo ad attendere la rete. I nostri kernel GEMM e di attention richiedono quasi tutta la shared memory di un SM, quindi non possono essere eseguiti su un SM che ospita un CTA di NCCL. NCCL sceglie il numero di canali necessario a sfruttare i collegamenti alla massima velocità. Nel nostro job usava 40 CTA per l’all-gather degli esperti e 64 per l’all-to-all dei gradienti in bf16. Hopper ha 132 SM. Quando queste due collettive si sovrapponevano, NCCL occupava fino a 88 SM, lasciandone 44 alle GEMM eseguite in parallelo. Nello step profilato, NCCL occupava in media 50 SM mentre era in corso una qualsiasi collettiva, e 64 o più per 20 s.

Lo scheduler dinamico riduce l’effetto degli SM persi sulle GEMM degli esperti. Questa sezione riduce invece il numero di SM che NCCL occupa. Queste collettive attraversano InfiniBand a circa 25 GB/s per GPU, e questo ne limita la velocità. Oltre un certo numero di CTA, aggiungerne altri non aumenta il throughput. NCCL_MAX_CTAS limita il numero di CTA per collettiva, e noi lo abbiamo impostato a 16.

Thread block NCCL per collettiva: valori predefiniti e NCCL_MAX_CTAS=16
Step profilato del rank 0 con il numero di thread block predefinito di NCCL e con NCCL_MAX_CTAS=16. Ogni thread block occupa uno dei 132 SM mentre la sua collettiva è in esecuzione. Il selettore delle metriche mostra la durata di ogni collettiva e il tempo GPU dei kernel di calcolo.

Le due collettive principali interessate dal limite hanno rallentato appena. L’all-gather degli esperti ha richiesto 155 ms sia con 16 sia con 40 CTA, e l’all-to-all è stato più lento del 3%. Le altre collettive sono cambiate di più, in entrambe le direzioni, come descritto più avanti. Le due esecuzioni elaborano gli stessi dati nello stesso ordine, quindi ogni lancio di un kernel di calcolo in un’esecuzione svolge lo stesso lavoro del corrispondente nell’altra. La figura seguente mostra il tempo dei principali kernel di calcolo nell’arco dello step in ciascuna esecuzione, e il numero di SM lasciati liberi da NCCL durante la loro esecuzione.

Tempo dei kernel di calcolo e SM liberi: valori predefiniti e NCCL_MAX_CTAS=16
Kernel di calcolo dello step profilato del rank 0, senza (magenta) e con (azzurro) NCCL_MAX_CTAS=16. La riga sotto ogni nome indica gli SM lasciati liberi da NCCL durante l’esecuzione di quel kernel, in media sul suo tempo di esecuzione, prima e dopo il limite.

Tutti questi kernel sono diventati più veloci, in misura circa proporzionale al numero di SM recuperati. I kernel di backward di FA3 venivano eseguiti per lo più quando non era in corso alcuna collettiva: disponevano di 126 SM con i valori predefiniti e di 130 con il limite, e hanno richiesto il 3% di tempo in meno. Le GEMM di QuACK con attivazione fusa sono passate da 81 a 118 SM e hanno richiesto il 22% di tempo in meno. Le GEMM di cuBLAS con layout TN sono passate da 82 a 112 SM e hanno richiesto il 20% di tempo in meno.

Complessivamente, questi kernel sono passati da 52,5 a 44,8 s, e le sole GEMM degli esperti da 23,5 a 19,2 s. Alcuni kernel di comunicazione interni al nodo (un all-gather su NVLink e la sincronizzazione di HybridEP) hanno trascorso più tempo in attesa, quindi il tempo GPU totale dello step è sceso di 6 s, da 60,6 a 54,7 s.

La comunicazione esposta è aumentata di 1,6 s, meno del tempo di calcolo risparmiato. Le collettive dense hanno mantenuto i loro 8 CTA e si sono concluse prima (da 119 a 40 ms per l’all-gather e da 131 a 78 ms per il reduce-scatter). Una spiegazione probabile è che una collettiva attende anche il rank più lento, e ora ogni rank raggiunge la collettiva prima. Con questa sola variabile d’ambiente, il throughput è salito da 1.210 a 1.293 token/s per GPU (+7%).

Abbiamo provato solo il valore 16, e altri valori potrebbero dare risultati migliori. Il limite si applica anche alle collettive che richiedono più SM. Il nuovo gather dei pesi densi su NVLink è stato eseguito con 16 CTA anziché 24, e il suo tempo è passato da 1,3 a 17 ms. Per questo job la differenza è modesta, ma un job con più traffico interno al nodo richiederebbe un limite più alto.

‍

Perdite di memoria host

I problemi che hanno richiesto più tempo in questo lavoro sono state le differenze tra la memoria host prevista e quella effettivamente disponibile sui nodi. Grace ha circa 108 GiB di memoria utilizzabile per GPU, e la configurazione di addestramento è impostata per usarne quasi tutta su ogni rank. La memoria persa su un qualsiasi rank abbassa il limite per l’intero job. Questa sezione descrive quattro cause di perdita di memoria, all’incirca nell’ordine in cui le abbiamo individuate.

Allocazioni host finite sull’HBM

Poiché l’HBM è esposta come nodo NUMA, quando la memoria di Grace è piena Linux vi colloca le normali allocazioni host. L’allocazione riesce senza errori, e in seguito CUDA esaurisce la memoria con gigabyte occupati che non risultano attribuiti ad alcun processo. numactl --membind impedisce al kernel di farlo, ma non copre le allocazioni avviate dalla GPU. Un buffer host paginabile la cui prima scrittura è una copia dalla GPU alla CPU ha le sue pagine collocate sull’HBM della GPU, anche con un binding rigoroso. Abbiamo scoperto che la nostra arena di memoria per l’offload delle attivazioni si comportava in questo modo, e occupava 13-17 GiB di HBM su ogni GPU. Abbiamo risolto scrivendo ogni pagina dell’arena dalla CPU una volta all’avvio.

Arrotondamento della memoria pinned

L’offload su CPU di FSDP blocca in memoria (pinning) ogni parametro e gradiente suddiviso. L’allocator di memoria pinned di PyTorch arrotonda ogni allocazione alla potenza di due successiva, quindi un tensore da 70 MiB occupa 128 MiB. Su tutto il modello, questo spazio aggiuntivo ammontava a 14-28 GiB per rank, e non era visibile in /dev/shm, ipcs o smaps.

Impostare pinned_max_round_threshold_mb:1 in PYTORCH_CUDA_ALLOC_CONF disattiva l’arrotondamento per le allocazioni superiori a 1 MiB. Su quattro allocazioni di test, l’impostazione riduce la memoria occupata da 1.728 MiB ai 1.077 MiB richiesti.

Allocazioni pinned: memoria richiesta e occupata
HMemoria host mappata per allocazioni pinned di quattro dimensioni, con l’arrotondamento predefinito di PyTorch e con pinned_max_round_threshold_mb:1.

Peso QKV replicato nel checkpointing

Un hook di state_dict che divide il peso QKV fuso per il checkpointing raccoglieva il peso completo su ogni rank, e il sistema di checkpoint conservava il risultato in cache per tutta l’esecuzione. Questo occupava 28 GiB per rank, indipendentemente dal numero di nodi. Ora i nostri checkpoint memorizzano il layout fuso. Nel frattempo TorchTitan ha spostato la conversione QKV fuori dai checkpoint finiva per conservare il risultato in cache per tutta l’esecuzione. Questo occupava 28 GiB per rank, indipendentemente dal numero di nodi. Abbiamo quindi ripiegato sulla memorizzazione dei checkpoint con il layout fuso. Nel frattempo TorchTitan ha spostato la conversione QKV fuori dai checkpoint di addestramento, negli adapter per Hugging Face (#4815), eliminando la causa alla radice.

Memoria non disponibile all’avvio dei job

Dopo queste correzioni, alcuni rank esaurivano ancora la memoria di Grace. Curiosamente, gli errori di memoria esaurita venivano segnalati come segmentation fault, probabilmente a causa della minore maturità dello stack sottostante.

Abbiamo scoperto che i nodi non avviavano i nostri job con la memoria libera. Abbiamo registrato lo stato di ogni socket Grace all’avvio di 62 dei nostri job, per un totale di 20.436 osservazioni. Un socket ha 119,7 GiB. Il socket mediano aveva 100,7 GiB liberi all’avvio del job, il peggiore 72,7 GiB. Gran parte della differenza è memoria che il kernel non attribuisce ad alcun processo né ad alcuna categoria di meminfo. Questa memoria non attribuita aveva una mediana di 10,9 GiB per socket, 17,6 GiB al 90° percentile e un massimo di 39,9 GiB. La maggior parte dei socket la rilascia quando la memoria scarseggia, ma per alcuni socket questo non avveniva, con conseguenze sull’intera esecuzione.

Memoria Grace libera all’avvio dei job
Ogni socket Grace all’avvio di 62 job (20.436 osservazioni). Il selettore delle metriche mostra la memoria che il kernel non attribuisce ad alcun processo. Una lunga coda di socket parte con meno dei circa 108 GiB previsti.

Abbiamo trovato anche una perdita di memoria nel nostro ambiente. La libc10 con cui compiliamo usa mimalloc, che appoggia le sue arene su transparent huge page da 512 MiB. Impostare MIMALLOC_ALLOW_THP=0 ha eliminato la perdita.

Abbiamo segnalato la memoria non attribuita al supporto del JSC, che ha confermato il problema e ci ha consigliato di restare sotto i 75 GB circa per rank finché gli epilogue dei job non fossero stati corretti. Il problema riguardava tutti gli utenti della macchina, ed è stato poi risolto.

‍

Inizializzazione dei parametri con offload su CPU

L’offload su CPU attiva percorsi del codice generalmente poco utilizzati. Questa sezione descrive due problemi di inizializzazione che abbiamo riscontrato su questi percorsi.

Valori casuali identici su tutti i rank

Con l’offload su CPU, TorchTitan inizializza ogni shard dei parametri direttamente sulla CPU, risparmiando HBM. Sulla GPU, l’RNG distribuito di PyTorch applica a ogni rank un offset dello stato casuale, così che gli shard vengano generati in modo indipendente. Sulla CPU, invece, l’inizializzatore non rileva che il parametro si trova sulla CPU e segue il percorso del codice CUDA. Quel percorso attinge correttamente al generatore della CPU, ma al termine ne reimposta lo stato.

Di conseguenza, ogni chiamata all’inizializzatore parte dallo stesso stato casuale, e ogni tensore su ogni rank riceve gli stessi valori casuali. Con la normale troncata di TorchTitan, l’inizializzazione non termina mai. I valori fuori dall’intervallo di troncamento vengono ricampionati, ma poiché lo stato dell’RNG non avanza mai, a ogni iterazione vengono estratti gli stessi valori fuori intervallo.

Abbiamo risolto il problema in due punti:

  • PyTorch: una issue e una correzione che fa avanzare lo stato dell’RNG. La correzione garantisce che l’inizializzazione termini, ma non che i rank estraggano valori diversi.
  • TorchTitan: una PR che permette di scegliere gli assi della mesh lungo cui applicare l’offset del seed.

Con la patch di PyTorch applicata e dp_shard come asse per il seed, il modello si inizializza correttamente con l’offload attivo.

Inizializzazione lenta della normale troncata

L’inizializzazione richiedeva circa 10 minuti, a causa di un secondo problema della normale troncata. Il suo ciclo di rigetto usa .any() per verificare se qualche valore va ancora ricampionato. Su un DTensor suddiviso, questa riduzione completa ridistribuisce la maschera a ogni rank, quindi ogni iterazione esegue un all-gather di una maschera grande quanto il tensore per calcolare un singolo valore booleano. Per tensori piccoli o collegamenti veloci il costo è trascurabile. Per un grande tensore di esperti fuso su InfiniBand, però, ogni verifica può richiedere secondi, che sommati su tutti i layer diventano minuti.

La correzione riduce ogni shard localmente ed esegue l’all-reduce di un unico flag. Ogni rank ottiene lo stesso risultato, quindi gli stati dell’RNG di tutti i rank restano sincronizzati. Per i dettagli si veda la PR.

‍

Contributi open source

La configurazione descritta in questo articolo si basa su software open source e su risorse di calcolo finanziate con fondi pubblici. Proponiamo tutte le nostre correzioni ai progetti originali e segnaliamo i problemi di sistema ai gestori del supercomputer, perché ne beneficino anche gli altri utenti di PyTorch, TorchTitan e Jupiter. La tabella elenca i contributi descritti in questo articolo.

WhereWhat
TorchTitan #5090Hooks for activation-checkpoint offloading
TorchTitan #5352Muon Newton-Schulz on GPU with CPU offload
TorchTitan #5579Skip embedding re-gather in backward
TorchTitan #4600Configurable seed-offset mesh axes
PyTorch issue #196562, PR #196072CPU init resets the RNG state
TorchTitan #5583Tied RoPE cache
PyTorch #201840trunc_normal_ without the all-gather on DTensors
TorchTitan #5627Option to allocate FSDP’s all-gather buffers on the compute stream (unshard_async_op)
SonicMoE issue #75Dynamic scheduler disabled with no option to enable it
Jupiter (reported to JSC support)Unattributed host memory at job start

‍

Conclusioni e prossimi passi

La configurazione finale per il modello EUROPA raggiunge circa 1.300 token/s per GPU su 384 GPU GH200 per replica DP, 14 volte la configurazione di partenza. A questo ritmo, il job elabora circa 43 miliardi di token al giorno per replica. Con quattro repliche data-parallel su 1.536 GPU, e con l’efficienza di scalabilità dell’89% che abbiamo misurato, ne elaborerebbe circa 150 miliardi al giorno. Con questo lavoro abbiamo dimostrato che un sistema pubblico, condiviso e di generazione Hopper, a cui si accede come semplici utenti, può ospitare un addestramento su scala di mille miliardi di parametri: un risultato significativo per l’ecosistema europeo dell’AI.

La configurazione a cui siamo arrivati è costruita attorno ai due limiti principali che abbiamo individuato: la memoria host per GPU e la banda tra nodi.

Ognuno dei due limiti ha richiesto una serie di interventi specifici. Per la banda tra nodi, l’aumento della lunghezza delle sequenze e il passaggio agli all-gather dei pesi in FP8 e agli all-to-all dei gradienti in BF16 hanno aumentato il rapporto tra calcolo e scambio di dati. Il resharding locale al nodo ha spostato su NVLink il nuovo gather dei pesi densi nel backward pass, e NCCL_MAX_CTAS ha liberato SM per i kernel di calcolo. Per la memoria host, abbiamo spostato attivazioni e stato dell’ottimizzatore su Grace, mantenendo Newton-Schulz sulla GPU, e abbiamo individuato le cause della perdita di memoria host. Vogliamo sottolineare che diversi di questi problemi emergono solo con l’offload su CPU a queste dimensioni del modello, un percorso in generale meno battuto per PyTorch e TorchTitan. Ciononostante, le correzioni sono già state integrate in PyTorch e TorchTitan, a testimonianza della vitalità di questi progetti e dell’impegno verso la loro comunità.

Questo articolo lascia diversi aspetti ancora da esplorare. Non abbiamo documentato quanto il routing limitato per gruppo incida a lungo termine sulla qualità del modello, se le GEMM degli esperti possano essere eseguite in FP8 senza perdite misurabili, se valori di NCCL_MAX_CTAS diversi da 16 diano risultati migliori, né se uno schema basato sul pipeline parallelism sarebbe stato più veloce. Contiamo di riferire su questi punti man mano che l’addestramento procede.

Il modello e tutte le sue metodologie di addestramento saranno rilasciati in forma aperta con la pubblicazione di EUROPA.

‍

Ringraziamenti

Ringraziamo il Jülich Supercomputing Centre per il supporto, in particolare nell’analisi della memoria host, e i maintainer di PyTorch e TorchTitan per aver revisionato e integrato le nostre pull request. EUROPA è finanziato dall’EuroHPC Joint Undertaking nell’ambito della AI-BOOST Frontier AI Grand Challenge.

‍

Riferimenti bibliografici

  1. Guo, W., Mishra, M., Cheng, X., Stoica, I., & Dao, T. (2026). SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations. https://arxiv.org/abs/2512.14080
  2. Jordan, K., Jin, Y., Boza, V., You, J., Cesista, F., Newhouse, L., & Bernstein, J. (2024). Muon: An optimizer for hidden layers in neural networks. https://kellerjordan.github.io/posts/muon/
  3. Kimi Team. (2025). Kimi K2: Open Agentic Intelligence. arXiv Preprint arXiv:2507.20534.
  4. Kimi Team, Bai, T., Bai, Y., Bao, Y., C., M., Cai, J., Cai, X., Cao, P., Cao, Y., Chai, Z., Charles, Y., Che, H. S., Chen, G., Chen, G., Chen, G., Chen, H., Chen, J., Chen, J., Chen, J., … Zu, X. (2026). Kimi K3: Open Frontier Intelligence. https://arxiv.org/abs/2607.24653
  5. Liu, J., Su, J., Yao, X., Jiang, Z., Lai, G., Du, Y., Qin, Y., Xu, W., Lu, E., Yan, J., & others. (2025). Muon is Scalable for LLM Training. arXiv Preprint arXiv:2502.16982.

Domande frequenti

No items found.
en

Addestrare un MoE da 1.000 miliardi di parametri sull’infrastruttura pubblica europea con PyTorch e TorchTitan

iGenius
November 12, 2026
·
30 min

In Domyn abbiamo avviato di recente l’addestramento di un modello di frontiera nell’ambito della European Frontier AI Grand Challenge, insieme a Fraunhofer IAIS e Fraunhofer IIS. La prima sfida è stata capire se e come fosse possibile addestrare un modello da mille miliardi di parametri sui supercomputer pubblici europei, generalmente ottimizzati per i carichi di lavoro scientifici.

Questo articolo descrive nel dettaglio come ci siamo riusciti. Riteniamo che il risultato rappresenti, di per sé, una svolta per la posizione dell’Europa nella corsa all’AI. Mentre si costruiscono AI factory dedicate, dimostriamo che i supercomputer pubblici europei possono già oggi raggiungere la scala di frontiera.

Nel giugno 2026, il consorzio EUROPA, guidato da Domyn con Fraunhofer IAIS e Fraunhofer IIS, si è aggiudicato la Frontier AI Grand Challenge. L’obiettivo del progetto è creare un modello di frontiera aperto, addestrato su infrastruttura pubblica europea e rilasciato insieme ai pesi e alle metodologie di addestramento. Il consorzio ha ottenuto l’accesso al 2,5% della capacità di calcolo di EuroHPC sui supercomputer pubblici europei.

Per allinearci all’attuale scala di frontiera, abbiamo deciso di addestrare un modello linguistico Mixture-of-Experts (MoE) da circa 1.000 miliardi di parametri totali. Tra i supercomputer europei di EuroHPC, per il pretraining abbiamo scelto di concentrarci su Jupiter, il supercomputer exascale dell’EuroHPC Joint Undertaking, ospitato presso il Forschungszentrum Jülich.

Finora modelli di queste dimensioni sono stati addestrati soprattutto su cluster dedicati di generazione Blackwell. Questo hardware offre una banda più elevata tra le GPU, più memoria per GPU e maggiore libertà nella scelta dello schema di parallelismo. Il Booster di Jupiter, invece, è un sistema di generazione Hopper condiviso da molti utenti. Accediamo al sistema come semplici utenti, quindi non possiamo modificare il kernel, l’interconnessione, la configurazione della memoria dei nodi né il software di sistema. I nodi sono stati configurati per la simulazione scientifica, e il loro equilibrio tra memoria e rete è diverso da quello di un cluster progettato per l’addestramento. Sulla stessa macchina operano centinaia di progetti diversi.

Di seguito descriviamo la configurazione di addestramento che usiamo in queste condizioni e le modifiche che hanno portato il throughput da 90 a circa 1.300 token/s per GPU su 384 GPU GH200 (96 nodi). Forniamo dettagli sufficienti perché altri utenti di Jupiter e di sistemi EuroHPC simili possano riutilizzarla. Quando abbiamo riscontrato problemi nel software condiviso, li abbiamo corretti direttamente in PyTorch e TorchTitan, e abbiamo segnalato i problemi di sistema allo Jülich Supercomputing Centre (JSC). L’elenco completo è nella sezione Contributi open source.

‍

Ambito
Questo articolo descrive la configurazione di addestramento attuale, nel momento in cui scriviamo. La configurazione è ancora in evoluzione, e alcuni dettagli potrebbero non essere più attuali al momento della lettura.

La Figura 1 mostra il miglioramento del throughput ottenuto con ciascuna delle modifiche descritte nell’articolo, a partire dalla configurazione di partenza. Passando il cursore su una riga se ne visualizzano tempo per step, memoria e tentativi di riallocazione dell’allocator.

Throughput per modifica, da 90 a 1.293 token/s per GPU
Throughput a regime (mediana degli step 4-10, rank 0) con l’introduzione progressiva delle modifiche. La configurazione di partenza usa sequenze da 8.192 token. La lunghezza delle sequenze arriva a 65.536 token grazie all’adozione dell’ottimizzatore Muon. Le esecuzioni rappresentate in trasparenza si sono interrotte in anticipo.

Date le caratteristiche dell’hardware, descritte più avanti, la maggior parte delle modifiche si è concentrata sulla riduzione del tempo che le GPU trascorrono in attesa della rete. Lo mostra la figura seguente, che per ogni esecuzione analizzata con il profiler ripartisce lo step tra calcolo, attesa dello scambio di dati e inattività.

Tempo per step per attività
Ripartizione dello step del rank 0 per ogni esecuzione analizzata con il profiler: calcolo sovrapposto allo scambio di dati, solo calcolo, attesa dello scambio di dati e inattività.

Sistema e modello

Hardware

Jupiter è il primo supercomputer europeo di classe exascale. Il suo modulo Booster conta 5.884 nodi di calcolo, ciascuno con quattro superchip NVIDIA GH200 Grace Hopper. Ogni superchip comprende i seguenti componenti:

  • Grace: una CPU Arm con 72 core e 120 GB di LPDDR5X a 512 GB/s.‍
  • Hopper: una GPU di classe H100 con 96 GB di HBM3 a 4 TB/s.‍
  • NVLink-C2C: un collegamento a 900 GB/s tra Grace e Hopper.

All’interno di un nodo, le coppie di GPU sono collegate tramite NVLink 4 a 300 GB/s (150 GB/s per direzione), e le quattro CPU Grace sono collegate a 100 GB/s per direzione. Ogni nodo dispone di quattro adattatori InfiniBand NDR200 per il traffico tra nodi, ciascuno con una velocità di 200 Gbit/s. Complessivamente offrono 800 Gbit/s (100 GB/s) per nodo, pari a circa 25 GB/s per GPU.

Nel corso dell’articolo faremo riferimento alle seguenti larghezze di banda:

LinkBandwidthRelative to InfiniBand
Grace ↔ Hopper (NVLink-C2C)900 GB/s36×
GPU ↔ GPU, same node (NVLink 4)150 GB/s per direction6×
GPU ↔ outside the node (InfiniBand)~25 GB/s per GPU1×
Larghezze di banda di un nodo Jupiter, in scala
Lo spessore delle linee è proporzionale alla larghezza di banda. Le linee sottili che escono dal nodo sono i collegamenti InfiniBand, su cui passa tutto il traffico tra nodi.

Software

Per il modello EUROPA abbiamo scelto di adottare PyTorch e TorchTitan (Liang et al., 2025), la libreria nativa di PyTorch per il pretraining di LLM su larga scala. Abbiamo scelto TorchTitan per i seguenti motivi:

  • È il framework di addestramento che il nostro team conosce meglio.
  • Il codice è compatto e semplice da adattare alla nostra configurazione.
  • Il progetto ha una comunità di sviluppo attiva, e le pull request vengono revisionate e integrate rapidamente.

Modello

L’architettura si basa su MiMo-V2.6 (Core Team et al., 2026), che estende l’architettura MiMo-V2-Flash (Core Team, Xiao, et al., 2026). La tabella riporta la configurazione che usiamo attualmente per gli esperimenti. Alcuni valori potrebbero cambiare prima dell’avvio dell’addestramento vero e proprio.

PropertyValue
Total parameters951.9B (32.8B active per token)
Layers74
Hidden size6,144
Experts384 routed, top-8
Expert hidden size1,792
Vocabulary256,000
Sequence length65,536 tokens
Hardware96 nodes, 384 GH200 superchips
Tokens per step384 × 65,536 ≈ 25M
OptimizerMuon (all weight matrices except the router), AdamW (embeddings, LM head, router, norms)
  1. Liang, W., Liu, T., Wright, L., Constable, W., Gu, A., Huang, C.-C., Zhang, I., Feng, W., Huang, H., Wang, J., Purandare, S., Nadathur, G., & Idreos, S. (2025). TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training. https://arxiv.org/abs/2410.06511

‍

Vincoli e schema di parallelismo

Vincoli

Come accennato, la potenza di calcolo grezza delle GPU non era il fattore limitante. Nel pretraining di un MoE da mille miliardi di parametri, le difficoltà riguardavano altri due elementi del sistema.

Banda tra nodi

Il traffico InfiniBand tra nodi è 6 volte più lento di quello tra GPU dello stesso nodo, e 36 volte più lento di quello tra GPU e CPU. In uno step di addestramento distribuito su 384 GPU per replica data-parallel, come nella nostra configurazione, questo collegamento viene attraversato molte volte. Il traffico su InfiniBand deve quindi essere ridotto oppure avvenire in parallelo al calcolo delle GPU.

Capacità della memoria host

Nei nodi di login di Jupiter, GH200 dispone di fino a 480 GB di LPDDR5X a 384 GB/s per Grace. I nodi di calcolo, invece, hanno 120 GB a 512 GB/s: una configurazione che privilegia la banda rispetto alla capacità. È adatta ai carichi di lavoro di simulazione, limitati dalla banda di memoria, ma per l’addestramento limita la quantità di dati che ogni nodo può elaborare e conservare tra HBM e memoria di Grace.

La memoria effettivamente utilizzabile, infatti, è inferiore a 120 GB. Su GH200, Linux espone l’HBM di ogni Hopper come un nodo NUMA senza CPU. Ogni nodo riporta 858 GiB di memoria, di cui 456 GiB di memoria Grace. Del totale, il kernel ne riserva circa 24 GiB, e restano così circa 108 GiB di memoria host utilizzabile per GPU. Altre perdite sono descritte nella sezione Perdite di memoria host.

Schema di parallelismo

Per un MoE da mille miliardi di parametri, scegliere lo schema di parallelismo più adatto all’architettura hardware è decisivo per raggiungere il throughput necessario in produzione. Inizialmente abbiamo valutato il pipeline parallelism, ma alla luce delle considerazioni che seguono ci siamo concentrati su FSDP ed expert parallelism, dedicando il nostro lavoro alla loro ottimizzazione: un impegno che alla fine ha dato i suoi frutti. Poiché lo schema adottato determina tutte le ottimizzazioni successive, non possiamo dire se uno schema basato sul pipeline parallelism sarebbe stato più veloce. La decisione si è basata su tre principi:

  1. Expert parallelism all’interno del nodo. L’all-to-all dei token di un layer MoE si trova sul percorso critico tra il router e gli esperti ed è difficile da sovrapporre al calcolo. Su NVLink il tempo che richiede è accettabile, su InfiniBand è troppo alto. Abbiamo quindi limitato l’expert parallelism a 4, con un gruppo per nodo.
  2. Step brevi. Un pipeline parallelism profondo richiede molti microbatch in ogni step per contenere la bubble, e questo allunga gli step. Con step più lunghi diventa più difficile individuare opportunità di ottimizzazione e problemi che emergono su larga scala, e i tempi di riscontro si allungano.
  3. Batch ampio. Un batch più ampio riduce il rumore dei gradienti e assegna a ogni GPU più lavoro per step, ammortizzando costi fissi come il lancio dei kernel e lo step dell’ottimizzatore.

Sulla base di queste considerazioni, abbiamo optato per FSDP (fully sharded data parallelism) a 384 vie, con expert parallelism a 4 vie all’interno di ciascun nodo. Tutte le esecuzioni descritte in questo articolo usano una sola replica data-parallel (DP = 1), quindi un unico gruppo FSDP copre l’intero job. In pratica, possiamo eseguire fino a 4 repliche data-parallel su 1.536 GPU senza effetti apprezzabili sul throughput.

DPNodesGPUsBatch size (sequences of 65,536 tokens)Throughput vs DP = 1
196384384 (≈25M tokens)100%
2192768768 (≈50M tokens)99.3%
43841,5361,536 (≈101M tokens)89.4%
Schema di parallelismo
96 nodi × 4 GPU. Sono rappresentati solo i primi nodi e l’ultimo. Il colore indica il blocco di 96 esperti che risiede su ciascuna GPU. Passando il cursore su una GPU se ne visualizzano i gruppi. Il nodo di appartenenza è il suo gruppo di expert parallelism (all-to-all su NVLink). Le 96 GPU dello stesso colore formano il gruppo FSDP tra cui gli esperti vengono suddivisi tramite InfiniBand. I parametri densi sono suddivisi su tutte le 384 GPU. La riga con il contorno è il gruppo FSDP. Eventuali repliche data-parallel (DP) in più comporterebbero altre righe come quella in trasparenza, ciascuna una copia completa del gruppo FSDP. Tutte le esecuzioni descritte in questo articolo usano una sola replica (DP = 1).

Lunghezza delle sequenze, ottimizzatore e kernel MoE

Configurazione di partenza

FSDP riusciva a svolgere in parallelo al calcolo solo una parte dello scambio di dati. Spesso il calcolo di un layer terminava prima che arrivassero i pesi del successivo, e la GPU restava inattiva. Nella configurazione di partenza (8.192 token per sequenza), uno step profilato sul rank 0 durava in totale 83,6 s: 14,3 s di calcolo e 68,3 s trascorsi in attesa della rete.

Timeline di GPU e CPU nella configurazione di partenza
Attività di GPU e CPU sul rank 0 durante lo step profilato della configurazione di partenza. L’all-gather è in esecuzione quasi senza sosta, mentre il calcolo occupa brevi intervalli. Ingrandendo il forward pass si vedono i calcoli di ciascun layer.

Poiché in FSDP lo scambio di dati riguarda solo i pesi, e il volume di pesi da raccogliere per ogni layer non dipende dal numero di token, con sequenze più lunghe la GPU svolge più calcolo a parità di dati trasmessi. Abbiamo quindi allungato le sequenze fin dove la memoria lo permetteva. Tuttavia, anche con l’activation checkpointing della configurazione di partenza, già una sequenza da 12.288 token superava la capacità dell’HBM.

Activation checkpointing e offload su Grace

Sequenze più lunghe richiedono più memoria per le attivazioni. In un primo momento abbiamo salvato solo le attivazioni in uscita da ciascun layer, ricalcolando durante il backward pass tutto ciò che avviene al suo interno. Il ricalcolo completo aumenta il carico di lavoro, ma la GPU restava inattiva per gran parte dello step, quindi il lavoro aggiuntivo non ne ha allungato la durata. Con questa modifica, la lunghezza delle sequenze è arrivata a 32.768 token e il throughput a 370 token/s per GPU, 4 volte quello della configurazione di partenza.

Con sequenze da 65.536 token, le sole attivazioni in uscita dai layer occupano 74 layers × 65,536 tokens × 6,144 × 2 bytes ≈ 55 GiB per GPU: troppo per l’HBM, che deve ospitare anche i pesi e le altre attivazioni. Abbiamo quindi deciso di spostarle su Grace tramite il collegamento NVLink-C2C da 900 GB/s. Ciascuna di queste attivazioni viene copiata su Grace dopo il forward pass, e riportata sulla GPU prima del ricalcolo durante il backward pass. Abbiamo impostato a 0,86 il parametro offload_fraction, che stabilisce la quota di layer da trasferire, e quindi il compromesso tra memoria HBM e memoria di Grace. Con l’offload, la lunghezza delle sequenze è arrivata a 49.152 token e il throughput a 508 token/s per GPU.

Schema di offload delle attivazioni
Schema semplificato a 16 layer anziché 74. Nel forward pass ogni attivazione in uscita viene copiata su Grace; nel backward pass torna su Hopper un layer prima di essere ricalcolata. I layer esclusi dalla quota di offload mantengono la propria attivazione in uscita nell’HBM.

Abbiamo sviluppato e integrato in TorchTitan gli hook per l’offload degli activation checkpoint, insieme a un offloader su CPU minimale fornito come esempio documentato. Su questi hook abbiamo poi costruito la strategia di offload descritta qui.

Muon con Newton-Schulz suddiviso per esperti

Lo stato dell’ottimizzatore occupa una parte significativa della memoria. Usiamo Muon (Jordan et al., 2024) per tutte le matrici dei pesi tranne il router, circa il 99,5% dei parametri, e AdamW per embedding, LM head, router e layer di normalizzazione. Muon memorizza un solo stato in fp32 per parametro (il momentum), mentre Adam ne memorizza due. Fare a meno del secondo stato ha permesso di far rientrare in memoria sequenze da 65.536 token. Muon è risultato più efficiente di AdamW nell’addestramento di modelli di grandi dimensioni (Liu et al., 2025), ed è stato usato per modelli MoE da mille miliardi di parametri (Kimi Team, 2025; Kimi Team et al., 2026).

Con Muon, ogni aggiornamento viene ortogonalizzato con alcune iterazioni di Newton-Schulz, che operano sulla matrice intera. Nella nostra configurazione FSDP, però, le matrici sono divise in 384 shard. Un’implementazione di base ricomporrebbe ciascuna matrice su un unico rank, applicherebbe Newton-Schulz e poi redistribuirebbe il risultato, con un aumento significativo dello scambio di dati.

FlexShard, ora integrato in TorchTitan, ci permette di stabilire quale rank possiede quali parametri. Grazie a questa possibilità, abbiamo distribuito gli esperti interi, senza suddividerli: ogni rank possiede matrici di esperti complete ed esegue le relative iterazioni di Newton-Schulz senza alcuno scambio di dati. Le matrici dense, come le proiezioni dell’attention, restano invece suddivise e vengono raccolte sul rank a cui sono assegnate. Questa configurazione consente di sfruttare ancora di più la sovrapposizione tra calcolo e scambio di dati: ogni rank può elaborare i propri esperti mentre le matrici dense sono in transito, così lo scambio di dati si sovrappone completamente al calcolo.

Un passaggio chiave è stato spostare l’esecuzione dell’ottimizzatore sul processore Grace. A differenza delle altre operazioni dell’ottimizzatore, però, Newton-Schulz è computazionalmente oneroso e trae grande vantaggio dall’accelerazione su GPU. La figura seguente confronta Newton-Schulz su Grace e su Hopper, con le dimensioni dei nostri esperti.

Tempo di Newton-Schulz su Grace e su Hopper
Tempo mediano per ortogonalizzare la matrice di un esperto (scala logaritmica) e tempo risultante per rank per step. Grace impiega circa 35 s per step usando tutti i 72 core, Hopper 0,2 s.

Durante l’addestramento, con Newton-Schulz su Grace, lo step dell’ottimizzatore richiedeva complessivamente 41,7 s. Abbiamo quindi sviluppato e integrato in TorchTitan il supporto per eseguire Newton-Schulz sulla GPU quando i parametri sono in offload. Con questa modifica, lo step dell’ottimizzatore è sceso a 2,2 s, e il throughput è salito da 441 a 597 token/s per GPU (+35%).

Su un dispositivo CPU come Grace, stream ed eventi di PyTorch non hanno alcun effetto. Di conseguenza, la sovrapposizione che FlexShard garantisce su CUDA va persa senza alcun avviso. Abbiamo modificato il codice lato CPU per eseguire in modo asincrono gli all-to-all delle matrici dense, anziché bloccare il thread principale. Così lo step dell’ottimizzatore è sceso a 1,6 s.

Step dell’ottimizzatore con Newton-Schulz su Grace e su Hopper
Step profilati sul rank 0 da 80 s in poi, su un asse temporale comune. Reimpostando lo zoom si vedono gli step completi. Con Newton-Schulz su Grace, la GPU resta inattiva per circa 34 s dopo il backward pass, come mostra la riga della CPU. La riga gloo mostra gli ultimi 10 s circa: un all-to-all delle matrici dense lato CPU, che deve attendere i rank più lenti. Con Newton-Schulz su Hopper, l’ottimizzatore occupa una piccola frazione dello step.

Scheduling dinamico dei tile in SonicMoE

Per gli esperti usiamo SonicMoE (Guo et al., 2026), un kernel MoE fuso, basato sui kernel CuTe DSL di QuACK. Esegue le GEMM degli esperti direttamente sul layout di dispatch dei token, senza doverli prima copiare e riordinare. Su una singola GH200, con dimensioni di test, ha ridotto il picco di memoria del layer da 12,5 a 5,4 GiB e il tempo di forward e backward da 213 a 153 ms.

I kernel di comunicazione vengono eseguiti sugli stessi SM della GPU dei kernel di calcolo. Di conseguenza, finché è in esecuzione, un all-gather di FSDP occupa alcuni SM con i thread block di NCCL. Le GEMM di SonicMoE usano uno scheduler persistente statico, che al lancio assegna a ciascuno dei 132 thread block persistenti un elenco fisso di tile. Se l’SM di un blocco è occupato dall’all-gather, quel blocco inizia a elaborare i propri tile solo quando l’all-gather termina, e la GEMM si conclude soltanto quando finisce l’ultimo blocco. In un’esecuzione, la stessa up-projection richiedeva 55-60 ms senza altre operazioni in corso e circa 100 ms in concomitanza con un all-gather.

QuACK offre anche uno scheduler dinamico, in cui ogni blocco, quando è libero, prende in carico il tile successivo tramite un contatore atomico condiviso: così i blocchi già in esecuzione elaborano anche i tile di quelli non ancora partiti. In SonicMoE, però, lo scheduler statico è fisso nel codice, quindi abbiamo racchiuso le sue chiamate GEMM in un wrapper che attiva quello dinamico. Abbiamo poi misurato le prestazioni di un layer MoE, con le dimensioni usate nel nostro addestramento, su una singola GH200, con e senza un kernel fittizio che occupa 40 SM.

Scheduling statico e dinamico dei tile in SonicMoE
Tempo mediano per ogni GEMM degli esperti e per forward e backward dell’intero layer. Con 40 SM occupati, lo scheduler dinamico riduce il tempo del layer da 304 a 217 ms.

Con 40 SM occupati, lo scheduler dinamico rende le GEMM più veloci del 17-45%, e il tempo del layer scende da 304 a 217 ms. Su una GPU libera è più veloce del 3-32% per GEMM e dell’11% sul layer. Gli output sono identici bit per bit, salvo gli arrotondamenti in fp32 nel gradiente dei punteggi di routing. In addestramento lo abbiamo attivato sulle quattro GEMM degli esperti in cui la dimensione dei token varia da un esperto all’altro. Il tempo per step è sceso da 108,7 a 103,8 s, con un ulteriore guadagno di throughput del 5%.

Abbiamo chiesto agli autori di SonicMoE perché lo scheduler predefinito sia quello statico. Limitare gli SM usati da NCCL affronta la stessa contesa, riducendo il numero di SM occupati dalla comunicazione.

‍

Bilanciamento del carico tra i rank

Le 384 GPU procedono in sincronia, quindi è il rank più lento a determinare la durata dello step. Due fonti di squilibrio rendevano alcuni rank più lenti di altri.

Bilanciamento dei dati in base al costo dell’attention

Il costo dell’attention cresce con il quadrato della lunghezza del documento. Con righe da 65.536 token, un rank potrebbe ricevere un unico documento da 60.000 token e un altro cento documenti brevi. Entrambi elaborano lo stesso numero di token, ma il primo svolge molto più lavoro di attention, e il suo calcolo richiede quindi più tempo.

Per attenuare questo effetto, bilanciamo i dati offline. Pre-tokenizziamo il corpus e raggruppiamo i documenti, in ordine casuale, in righe da 65.536 token. Un documento che non entra in una riga prosegue nella successiva; per ogni riga calcoliamo poi il costo dell’attention come somma dei quadrati delle lunghezze dei segmenti, dove un segmento è la parte di un documento contenuta nella riga. Ordiniamo le righe in fasce di costo ristrette, le mescoliamo all’interno di ogni fascia e riempiamo ogni step con righe consecutive secondo quest’ordine, una per rank. Le righe di uno step hanno quindi un costo quasi uguale. Infine, mescoliamo l’ordine degli step.

Sui 1.057 step, il rank più oneroso di uno step svolgeva in mediana 7,9 volte il lavoro di attention medio prima del bilanciamento, e 1,15 volte dopo.

Costo massimo dell’attention per rank, step per step
Per ogni step, il costo dell’attention più alto tra i rank (somma dei quadrati delle lunghezze dei segmenti) diviso per la media dello step. L’ordine di arrivo usa le righe nel loro ordine originale. Il piano bilanciato riunisce nello stesso step righe di costo simile.

In addestramento, il bilanciamento dei dati ha ridotto il tempo per step da 103,8 a 91,7 s, con un guadagno di throughput del 13%.

Il bilanciamento modifica la distribuzione dei dati all’interno di uno step. I documenti lunghi vengono raggruppati negli stessi step, invece di arrivare mescolati come accadrebbe naturalmente. Se un rank riceve una riga con un unico documento da 65.000 token, ogni altro rank in quello step riceve una riga simile, e questo può influire sulla dinamica dell’addestramento. Ci aspettiamo che l’effetto sia limitato ai documenti lunghi, e prevediamo di compensarlo con una fase dedicata al contesto lungo.

Routing degli esperti limitato per gruppo

La seconda fonte di squilibrio è il router del MoE. Quando invia più token agli esperti di una GPU, quella GPU diventa la più lenta del suo nodo. In esecuzioni precedenti senza bilanciamento dei dati, lo squilibrio tra esperti costava circa 16-17 s per step, circa il 20% dello step.

Abbiamo provato a rendere il carico degli esperti perfettamente uniforme. I 384 esperti sono divisi in quattro gruppi da 96, uno per rank di expert parallelism. Il router seleziona lo stesso numero di esperti da ogni gruppo (due per gruppo con routing top-8). Ogni GPU riceve così lo stesso numero di coppie token-esperto in ogni layer. Sommata al bilanciamento dei dati, questa modifica ha ridotto il tempo per step da 91,7 a 85,9 s, con un guadagno di throughput del 7%. Inoltre, il picco di HBM è sceso da 87,2 a 81,3 GiB, perché nessuna GPU deve più riservare memoria per un eccesso di token.

Il vincolo, però, riduce la libertà di routing, e questo può incidere sull’efficacia dell’addestramento. Ad esempio, un token i cui tre esperti con il punteggio più alto appartengono allo stesso gruppo può essere instradato solo verso due di essi. Selezionare più esperti per gruppo aumenta la probabilità che ogni token raggiunga i suoi esperti migliori mantenendo il carico uniforme. Le nostre prime esecuzioni indicano che il vincolo può effettivamente penalizzare l’addestramento, e stiamo ancora valutando se mantenerlo o allentarlo con una strategia intermedia.

‍

Ridurre il volume dello scambio di dati

Oltre una certa lunghezza, sequenze più lunghe non miglioravano più il throughput. Quando le attivazioni mettono sotto pressione la memoria, l’allocator CUDA inizia a ritentare le allocazioni e a frammentare la memoria, e la frammentazione finisce per costare più tempo di quanto se ne risparmi con sequenze più lunghe. Da quel punto, l’unica opzione rimasta era inviare meno dati. Prima delle modifiche descritte in questa sezione, uno step profilato passava 39,5 dei suoi 98,7 s ad attendere la rete (vedi la ripartizione dello step nel grafico Tempo per step per attività, e all-gather e reduce-scatter occupavano la GPU per oltre 50 s.

All-gather dei pesi degli esperti in FP8

Dall’analisi con il profiler, l’all-gather dei pesi degli esperti risulta la voce più onerosa dello scambio di dati in uno step: circa 7 GB per rank in bf16, cioè 28 GB per nodo. Raccogliere i pesi in FP8 dimezzerebbe questo volume.

A questo scopo, i pesi vengono divisi in blocchi da 128 valori, e ogni blocco viene scalato all’intervallo dell’FP8 (E4M3) e convertito. Ogni rank invia i valori FP8 più un fattore di scala in fp32 per blocco; i riceventi li riconvertono poi in bf16 prima del calcolo degli esperti. Il forward pass usa quindi i pesi a precisione ridotta, con un errore per valore fino a circa il 6% del massimo del blocco. I master weight, i gradienti e lo stato dell’ottimizzatore restano a precisione piena. L’approccio è simile a quello usato per DeepSeek-V3 (DeepSeek-AI, 2024).

Con questa strategia, il tempo di all-gather sulla GPU è sceso da 53,6 a 35,1 s per step, e il throughput è salito da 763 a 898 token/s per GPU (+18%). I primi risultati non mostrano effetti apprezzabili sull’addestramento.

Poiché i pesi arrivano già in FP8, anche le GEMM degli esperti potrebbero essere eseguite in FP8, con un’ulteriore piccola perdita di precisione, come nell’addestramento MoE in FP8 di DeepSeek-V3. Potremmo adottare questa soluzione come ottimizzazione aggiuntiva.

All-to-all dei gradienti in BF16 con accumulo in FP32

Con gli all-gather in FP8, lo scambio di dati del forward pass si sovrappone completamente al calcolo. Il backward pass, invece, restava limitato dallo scambio di dati a causa della riduzione dei gradienti.

Nella configurazione standard di FSDP a precisione mista, i gradienti vengono calcolati in bf16 e ridotti in fp32. Prima del reduce-scatter vengono convertiti in fp32, quindi l’anello accumula in fp32 e anche i gradienti vengono trasmessi in fp32.

Solo la somma, però, richiede l’fp32; il trasferimento può avvenire in bf16. Partendo da questa osservazione, abbiamo deciso di inviare i gradienti in bf16 tramite un all-to-all, in cui ogni rank invia a ciascun altro rank la porzione del proprio gradiente di cui quel rank è proprietario. Il proprietario converte le porzioni in fp32 e le somma localmente. Il risultato coincide con quello del reduce-scatter in fp32 a meno dell’ordine di somma, e nei nostri test su 4 rank era identico bit per bit. L’all-to-all invia metà dei byte e dimezza il buffer dei gradienti mantenuto tra un layer e l’altro.

Riduzione dei gradienti: reduce-scatter in fp32 e all-to-all in bf16
Quattro secondi del backward pass sul rank 0. Nell’esecuzione in alto, l’anello del reduce-scatter in fp32 è occupato quasi senza sosta. In quella in basso, i gradienti degli esperti vengono inviati con un all-to-all in bf16, e il reduce-scatter rimanente (parametri densi) è breve.

Il tempo della GPU nelle operazioni collettive sui gradienti è sceso da 50,7 a 33,2 s per step (10,4 s di reduce-scatter e 22,8 s di all-to-all). Il tempo trascorso in attesa della rete si è ridotto da 27,3 a 13,9 s, e il throughput è salito da 898 a 1.022 token/s per GPU (+14%). Con tutte le altre modifiche di questa sezione applicate, disattivare l’all-to-all in bf16 riduce il throughput da 1.210 a 1.004 token/s per GPU.

Né gli all-gather in FP8 né i gradienti in bf16 modificano la loss in modo misurabile. In una serie di esecuzioni di confronto, con lo stesso ordine dei dati e la stessa inizializzazione, dopo dieci step le loss differiscono di meno di 0,01.

Resharding locale al nodo dei pesi densi

Con FSDP, i pesi di ogni layer vengono liberati dopo il forward pass e raccolti di nuovo nel backward pass per il ricalcolo. Ogni nuovo gather nel backward pass compete per la banda InfiniBand con il reduce-scatter dei gradienti.

Gli esperti costituiscono la maggior parte del modello e non possono restare in memoria in quella fase. I parametri densi (proiezioni dell’attention, router e layer di normalizzazione), invece, possono restare sulla GPU se suddivisi. Dopo il forward pass, suddividiamo i pesi densi di ogni layer in quattro parti all’interno del nodo, anziché in 384 parti su tutto il job, così il nuovo gather nel backward pass avviene su NVLink anziché su InfiniBand. Questo occupa circa 55 MiB di HBM per layer, circa 4 GiB in totale. Al contrario, mantenere i pesi densi interi (circa 16 GiB) provocava tentativi di riallocazione con la nostra lunghezza di sequenza, con una conseguente perdita di throughput.

La modifica ha richiesto poco codice nuovo. TorchTitan registra i parametri su due mesh: una mesh densa che copre tutti i rank e una mesh sparsa per gli esperti, la cui dimensione di expert parallelism corrisponde al nodo. Abbiamo impostato il resharding dei parametri densi dopo il forward pass (post_forward_mesh_info) sul gruppo locale al nodo, e allentato alcuni controlli di FSDP più rigidi del necessario.

Da sola, questa modifica ha migliorato il throughput del 2% (da 1.022 a 1.038 token/s per GPU). I log dell’esecuzione mostrano che i 4 GiB aggiuntivi provocavano tentativi di riallocazione su tutti i 384 rank.

I pool dell’allocator sugli stream secondari di FSDP

Dopo un aggiornamento del nostro stack software, l’allocator CUDA ha iniziato a ritentare le allocazioni su GPU con un utilizzo della memoria intorno al 70%. Ogni nuovo tentativo libera l’intera cache dell’allocator e richiede una sincronizzazione completa del dispositivo, un’operazione lenta. Al momento dei tentativi, 25-30 GiB di HBM erano liberi, mappati e presenti nella cache dell’allocator, ma l’allocator non riusciva a usarli.

La causa è nel modo in cui FSDP sovrappone lo scambio di dati al calcolo.Una GPU esegue il lavoro in stream, cioè code che procedono in parallelo. FSDP usa uno stream di calcolo e quattro stream secondari:

StreamWork
ComputeForward and backward pass, and the copy of gathered weights into each parameter
All-gather copy-inPacking of this rank’s weight shards into one flat buffer
All-gatherNCCL all-gather of that buffer
Reduce-scatterGradient reduce-scatter
All-reduceAll-reduce across replicas, and hooks such as our gradient norm

Il caching allocator di PyTorch mantiene un pool di memoria per ogni stream. La memoria liberata torna al pool dello stream che l’ha allocata, e solo quello stream può riutilizzarla, perché un altro stream avrebbe bisogno di una sincronizzazione per sapere che il precedente utilizzatore ha finito. FSDP alloca ogni buffer sullo stream che lo utilizza per scrivere i dati. Il nostro all-gather in FP8 e il nostro all-to-all in bf16 allocano i propri buffer temporanei all’interno delle chiamate di FSDP, quindi anche questi buffer finivano sugli stream secondari. Ogni buffer viene usato brevemente una volta per layer, ma la cache di ciascuno stream finisce per trattenere i blocchi più grandi che ha allocato.

Abbiamo ricostruito il log degli eventi dell’allocator a partire da uno snapshot della memoria del rank 0. Al picco di memoria attiva, gli stream secondari trattenevano 27,3 GiB, di cui solo 6,4 effettivamente in uso. I restanti 20,9 GiB erano inaccessibili allo stream di calcolo, che disponeva di 4,5 GiB di cache libera. A ogni nuovo tentativo, lo stream di calcolo non riusciva a ottenere un blocco libero per un buffer di SonicMoE da 4 GiB. Ridurre il numero di layer non risolveva il problema, perché diminuiva la memoria attiva ma lasciava invariata la cache inutilizzata.

HBM per stream dell’allocator
Ricostruzione a partire dalla traccia degli eventi del caching allocator CUDA sul rank 0, al picco di memoria attiva. Pieno: memoria occupata da tensori attivi. In trasparenza: memoria in cache dello stream ma inutilizzata, che non può essere riutilizzata dagli altri stream.

Abbiamo introdotto tre modifiche perché questa memoria venga allocata da pool riutilizzabili dallo stream di calcolo:

  1. Rilascio anticipato dei gradienti. Il nostro hook per la norma dei gradienti (vedi il paragrafo Modifiche minori a memoria e latenza viene eseguito su ogni layer. Quando un hook è registrato, FSDP mantiene in memoria il gradiente ridotto di ogni layer fino alla fine del backward pass, perché lo stream del reduce-scatter non lo riutilizzi mentre l’hook lo legge. Con un solo hook non è un problema, ma con 74 hook il costo è di 9 GiB. Noi rilasciamo ogni gradiente due layer più tardi, dopo aver fatto in modo che lo stream del reduce-scatter attenda la fine dell’hook.
  2. All-gather sullo stream di calcolo. FSDP2 prevedeva un flag privato e non documentato, unshard_async_op, che esegue il setup dell’all-gather sullo stream di calcolo e lascia che NCCL lavori sul proprio stream interno. I buffer vengono così presi dal pool dello stream di calcolo e vi ritornano. Per farlo è servita una modalità asincrona del nostro gather in FP8, che ora restituisce un handle che esegue la riconversione quando FSDP lo attende: lo stream di calcolo attende la rete solo quando ha bisogno dei pesi. Abbiamo proposto a TorchTitan un’opzione pubblica per attivare questo flag. Su un modello di test più piccolo, ha ridotto il picco di memoria riservata di circa l’11%, al costo di circa lo 0,7% di throughput.‍
  3. Buffer di ricezione dei gradienti sullo stream di calcolo. Il buffer di ricezione dell’all-to-all in bf16, grande quanto il gradiente, viene allocato dal pool dello stream di calcolo.

Sull’intero job da 384 GPU, con le modifiche descritte nel paragrafo Modifiche minori a memoria e latenza già applicate, il numero di rank con tentativi di riallocazione è sceso da 384 a zero. Il picco di memoria riservata è passato da 88,1 a 74,6 GiB e il throughput da 996 a 1.210 token/s per GPU (+21%). Il volume di dati inviati resta invariato: cambiano solo i pool da cui viene allocata la memoria.

Dopo la revisione della nostra PR su TorchTitan, il flag unshard_async_op sta per diventare un’API pubblica di FSDP (PyTorch #202343).

  1. DeepSeek-AI. (2024). DeepSeek-V3 Technical Report. arXiv Preprint arXiv:2412.19437.

‍

Modifiche minori a memoria e latenza

Singolarmente, le modifiche di questa sezione hanno un effetto ridotto sul tempo per step. Le abbiamo applicate mentre tutti i rank ritentavano le allocazioni I pool dell’allocator sugli stream secondari di FSDP, e il throughput è passato da 1.038 a 996 token/s per GPU. Non abbiamo misurato il loro effetto individuale sul throughput, ma hanno ridotto l’uso della memoria.

Norma dei gradienti durante il backward pass

Il gradient clipping richiede la norma globale dei gradienti. Con l’offload su CPU, TorchTitan la calcola dopo il backward pass sui gradienti già trasferiti su Grace, con circa 0,43 s di _foreach_norm sulla CPU e circa 0,19 s di all-reduce gloo.

Ogni shard dei gradienti si trova sulla GPU quando il suo reduce-scatter si conclude. In quel momento aggiungiamo il quadrato della sua norma a un accumulatore scalare. Dopo il backward pass, un solo all-reduce NCCL dello scalare fornisce la norma globale, e il clipping procede come prima.

Nessun all-gather degli embedding nel backward pass

Di default, FSDP raccoglie di nuovo la tabella degli embedding nel backward pass. Il gradiente degli embedding dipende solo dagli ID dei token in input, che ogni rank possiede già, quindi il gather è superfluo. Nel nostro communicator, inoltre, restava in attesa dietro il reduce-scatter del layer 1, esponendo fino a circa 126 ms per step. La correzione richiede una sola riga:

model.tok_embeddings.set_unshard_in_backward(False)

L’abbiamo proposta per l’integrazione in TorchTitan.

Cache RoPE condivisa

TorchTitan costruisce una cache RoPE per ogni layer, anche se layer con la stessa configurazione ottengono cache identiche. Con 65.536 token, i nostri 37 layer RoPE occupavano 2,75 GiB con copie di due sole cache distinte, abbastanza da provocare un errore di memoria esaurita. Condividere le cache riduce questa occupazione a 0,19 GiB, con un risparmio di 2,56 GiB per GPU.

La condivisione di un unico buffer era già stata proposta, ma la proposta si è arenata per timori legati alla fragilità della soluzione e al pipeline parallelism. Abbiamo presentato una versione più contenuta, che condivide la cache nello stesso punto in cui TorchTitan condivide gli embedding, ed è quindi compatibile con il pipeline parallelism.

‍

Limitare gli SM usati da NCCL

Sovrapporre comunicazione e calcolo significa eseguirli entrambi sulla stessa GPU nello stesso momento. Questa sezione riguarda gli SM che i kernel di comunicazione sottraggono ai kernel di calcolo.

NCCL esegue ogni operazione collettiva come un kernel GPU con un cooperative thread array (CTA, il termine di NVIDIA per un thread block) per canale. Ogni CTA occupa un SM per tutta la durata della collettiva e passa gran parte di quel tempo ad attendere la rete. I nostri kernel GEMM e di attention richiedono quasi tutta la shared memory di un SM, quindi non possono essere eseguiti su un SM che ospita un CTA di NCCL. NCCL sceglie il numero di canali necessario a sfruttare i collegamenti alla massima velocità. Nel nostro job usava 40 CTA per l’all-gather degli esperti e 64 per l’all-to-all dei gradienti in bf16. Hopper ha 132 SM. Quando queste due collettive si sovrapponevano, NCCL occupava fino a 88 SM, lasciandone 44 alle GEMM eseguite in parallelo. Nello step profilato, NCCL occupava in media 50 SM mentre era in corso una qualsiasi collettiva, e 64 o più per 20 s.

Lo scheduler dinamico riduce l’effetto degli SM persi sulle GEMM degli esperti. Questa sezione riduce invece il numero di SM che NCCL occupa. Queste collettive attraversano InfiniBand a circa 25 GB/s per GPU, e questo ne limita la velocità. Oltre un certo numero di CTA, aggiungerne altri non aumenta il throughput. NCCL_MAX_CTAS limita il numero di CTA per collettiva, e noi lo abbiamo impostato a 16.

Thread block NCCL per collettiva: valori predefiniti e NCCL_MAX_CTAS=16
Step profilato del rank 0 con il numero di thread block predefinito di NCCL e con NCCL_MAX_CTAS=16. Ogni thread block occupa uno dei 132 SM mentre la sua collettiva è in esecuzione. Il selettore delle metriche mostra la durata di ogni collettiva e il tempo GPU dei kernel di calcolo.

Le due collettive principali interessate dal limite hanno rallentato appena. L’all-gather degli esperti ha richiesto 155 ms sia con 16 sia con 40 CTA, e l’all-to-all è stato più lento del 3%. Le altre collettive sono cambiate di più, in entrambe le direzioni, come descritto più avanti. Le due esecuzioni elaborano gli stessi dati nello stesso ordine, quindi ogni lancio di un kernel di calcolo in un’esecuzione svolge lo stesso lavoro del corrispondente nell’altra. La figura seguente mostra il tempo dei principali kernel di calcolo nell’arco dello step in ciascuna esecuzione, e il numero di SM lasciati liberi da NCCL durante la loro esecuzione.

Tempo dei kernel di calcolo e SM liberi: valori predefiniti e NCCL_MAX_CTAS=16
Kernel di calcolo dello step profilato del rank 0, senza (magenta) e con (azzurro) NCCL_MAX_CTAS=16. La riga sotto ogni nome indica gli SM lasciati liberi da NCCL durante l’esecuzione di quel kernel, in media sul suo tempo di esecuzione, prima e dopo il limite.

Tutti questi kernel sono diventati più veloci, in misura circa proporzionale al numero di SM recuperati. I kernel di backward di FA3 venivano eseguiti per lo più quando non era in corso alcuna collettiva: disponevano di 126 SM con i valori predefiniti e di 130 con il limite, e hanno richiesto il 3% di tempo in meno. Le GEMM di QuACK con attivazione fusa sono passate da 81 a 118 SM e hanno richiesto il 22% di tempo in meno. Le GEMM di cuBLAS con layout TN sono passate da 82 a 112 SM e hanno richiesto il 20% di tempo in meno.

Complessivamente, questi kernel sono passati da 52,5 a 44,8 s, e le sole GEMM degli esperti da 23,5 a 19,2 s. Alcuni kernel di comunicazione interni al nodo (un all-gather su NVLink e la sincronizzazione di HybridEP) hanno trascorso più tempo in attesa, quindi il tempo GPU totale dello step è sceso di 6 s, da 60,6 a 54,7 s.

La comunicazione esposta è aumentata di 1,6 s, meno del tempo di calcolo risparmiato. Le collettive dense hanno mantenuto i loro 8 CTA e si sono concluse prima (da 119 a 40 ms per l’all-gather e da 131 a 78 ms per il reduce-scatter). Una spiegazione probabile è che una collettiva attende anche il rank più lento, e ora ogni rank raggiunge la collettiva prima. Con questa sola variabile d’ambiente, il throughput è salito da 1.210 a 1.293 token/s per GPU (+7%).

Abbiamo provato solo il valore 16, e altri valori potrebbero dare risultati migliori. Il limite si applica anche alle collettive che richiedono più SM. Il nuovo gather dei pesi densi su NVLink è stato eseguito con 16 CTA anziché 24, e il suo tempo è passato da 1,3 a 17 ms. Per questo job la differenza è modesta, ma un job con più traffico interno al nodo richiederebbe un limite più alto.

‍

Perdite di memoria host

I problemi che hanno richiesto più tempo in questo lavoro sono state le differenze tra la memoria host prevista e quella effettivamente disponibile sui nodi. Grace ha circa 108 GiB di memoria utilizzabile per GPU, e la configurazione di addestramento è impostata per usarne quasi tutta su ogni rank. La memoria persa su un qualsiasi rank abbassa il limite per l’intero job. Questa sezione descrive quattro cause di perdita di memoria, all’incirca nell’ordine in cui le abbiamo individuate.

Allocazioni host finite sull’HBM

Poiché l’HBM è esposta come nodo NUMA, quando la memoria di Grace è piena Linux vi colloca le normali allocazioni host. L’allocazione riesce senza errori, e in seguito CUDA esaurisce la memoria con gigabyte occupati che non risultano attribuiti ad alcun processo. numactl --membind impedisce al kernel di farlo, ma non copre le allocazioni avviate dalla GPU. Un buffer host paginabile la cui prima scrittura è una copia dalla GPU alla CPU ha le sue pagine collocate sull’HBM della GPU, anche con un binding rigoroso. Abbiamo scoperto che la nostra arena di memoria per l’offload delle attivazioni si comportava in questo modo, e occupava 13-17 GiB di HBM su ogni GPU. Abbiamo risolto scrivendo ogni pagina dell’arena dalla CPU una volta all’avvio.

Arrotondamento della memoria pinned

L’offload su CPU di FSDP blocca in memoria (pinning) ogni parametro e gradiente suddiviso. L’allocator di memoria pinned di PyTorch arrotonda ogni allocazione alla potenza di due successiva, quindi un tensore da 70 MiB occupa 128 MiB. Su tutto il modello, questo spazio aggiuntivo ammontava a 14-28 GiB per rank, e non era visibile in /dev/shm, ipcs o smaps.

Impostare pinned_max_round_threshold_mb:1 in PYTORCH_CUDA_ALLOC_CONF disattiva l’arrotondamento per le allocazioni superiori a 1 MiB. Su quattro allocazioni di test, l’impostazione riduce la memoria occupata da 1.728 MiB ai 1.077 MiB richiesti.

Allocazioni pinned: memoria richiesta e occupata
HMemoria host mappata per allocazioni pinned di quattro dimensioni, con l’arrotondamento predefinito di PyTorch e con pinned_max_round_threshold_mb:1.

Peso QKV replicato nel checkpointing

Un hook di state_dict che divide il peso QKV fuso per il checkpointing raccoglieva il peso completo su ogni rank, e il sistema di checkpoint conservava il risultato in cache per tutta l’esecuzione. Questo occupava 28 GiB per rank, indipendentemente dal numero di nodi. Ora i nostri checkpoint memorizzano il layout fuso. Nel frattempo TorchTitan ha spostato la conversione QKV fuori dai checkpoint finiva per conservare il risultato in cache per tutta l’esecuzione. Questo occupava 28 GiB per rank, indipendentemente dal numero di nodi. Abbiamo quindi ripiegato sulla memorizzazione dei checkpoint con il layout fuso. Nel frattempo TorchTitan ha spostato la conversione QKV fuori dai checkpoint di addestramento, negli adapter per Hugging Face (#4815), eliminando la causa alla radice.

Memoria non disponibile all’avvio dei job

Dopo queste correzioni, alcuni rank esaurivano ancora la memoria di Grace. Curiosamente, gli errori di memoria esaurita venivano segnalati come segmentation fault, probabilmente a causa della minore maturità dello stack sottostante.

Abbiamo scoperto che i nodi non avviavano i nostri job con la memoria libera. Abbiamo registrato lo stato di ogni socket Grace all’avvio di 62 dei nostri job, per un totale di 20.436 osservazioni. Un socket ha 119,7 GiB. Il socket mediano aveva 100,7 GiB liberi all’avvio del job, il peggiore 72,7 GiB. Gran parte della differenza è memoria che il kernel non attribuisce ad alcun processo né ad alcuna categoria di meminfo. Questa memoria non attribuita aveva una mediana di 10,9 GiB per socket, 17,6 GiB al 90° percentile e un massimo di 39,9 GiB. La maggior parte dei socket la rilascia quando la memoria scarseggia, ma per alcuni socket questo non avveniva, con conseguenze sull’intera esecuzione.

Memoria Grace libera all’avvio dei job
Ogni socket Grace all’avvio di 62 job (20.436 osservazioni). Il selettore delle metriche mostra la memoria che il kernel non attribuisce ad alcun processo. Una lunga coda di socket parte con meno dei circa 108 GiB previsti.

Abbiamo trovato anche una perdita di memoria nel nostro ambiente. La libc10 con cui compiliamo usa mimalloc, che appoggia le sue arene su transparent huge page da 512 MiB. Impostare MIMALLOC_ALLOW_THP=0 ha eliminato la perdita.

Abbiamo segnalato la memoria non attribuita al supporto del JSC, che ha confermato il problema e ci ha consigliato di restare sotto i 75 GB circa per rank finché gli epilogue dei job non fossero stati corretti. Il problema riguardava tutti gli utenti della macchina, ed è stato poi risolto.

‍

Inizializzazione dei parametri con offload su CPU

L’offload su CPU attiva percorsi del codice generalmente poco utilizzati. Questa sezione descrive due problemi di inizializzazione che abbiamo riscontrato su questi percorsi.

Valori casuali identici su tutti i rank

Con l’offload su CPU, TorchTitan inizializza ogni shard dei parametri direttamente sulla CPU, risparmiando HBM. Sulla GPU, l’RNG distribuito di PyTorch applica a ogni rank un offset dello stato casuale, così che gli shard vengano generati in modo indipendente. Sulla CPU, invece, l’inizializzatore non rileva che il parametro si trova sulla CPU e segue il percorso del codice CUDA. Quel percorso attinge correttamente al generatore della CPU, ma al termine ne reimposta lo stato.

Di conseguenza, ogni chiamata all’inizializzatore parte dallo stesso stato casuale, e ogni tensore su ogni rank riceve gli stessi valori casuali. Con la normale troncata di TorchTitan, l’inizializzazione non termina mai. I valori fuori dall’intervallo di troncamento vengono ricampionati, ma poiché lo stato dell’RNG non avanza mai, a ogni iterazione vengono estratti gli stessi valori fuori intervallo.

Abbiamo risolto il problema in due punti:

  • PyTorch: una issue e una correzione che fa avanzare lo stato dell’RNG. La correzione garantisce che l’inizializzazione termini, ma non che i rank estraggano valori diversi.
  • TorchTitan: una PR che permette di scegliere gli assi della mesh lungo cui applicare l’offset del seed.

Con la patch di PyTorch applicata e dp_shard come asse per il seed, il modello si inizializza correttamente con l’offload attivo.

Inizializzazione lenta della normale troncata

L’inizializzazione richiedeva circa 10 minuti, a causa di un secondo problema della normale troncata. Il suo ciclo di rigetto usa .any() per verificare se qualche valore va ancora ricampionato. Su un DTensor suddiviso, questa riduzione completa ridistribuisce la maschera a ogni rank, quindi ogni iterazione esegue un all-gather di una maschera grande quanto il tensore per calcolare un singolo valore booleano. Per tensori piccoli o collegamenti veloci il costo è trascurabile. Per un grande tensore di esperti fuso su InfiniBand, però, ogni verifica può richiedere secondi, che sommati su tutti i layer diventano minuti.

La correzione riduce ogni shard localmente ed esegue l’all-reduce di un unico flag. Ogni rank ottiene lo stesso risultato, quindi gli stati dell’RNG di tutti i rank restano sincronizzati. Per i dettagli si veda la PR.

‍

Contributi open source

La configurazione descritta in questo articolo si basa su software open source e su risorse di calcolo finanziate con fondi pubblici. Proponiamo tutte le nostre correzioni ai progetti originali e segnaliamo i problemi di sistema ai gestori del supercomputer, perché ne beneficino anche gli altri utenti di PyTorch, TorchTitan e Jupiter. La tabella elenca i contributi descritti in questo articolo.

WhereWhat
TorchTitan #5090Hooks for activation-checkpoint offloading
TorchTitan #5352Muon Newton-Schulz on GPU with CPU offload
TorchTitan #5579Skip embedding re-gather in backward
TorchTitan #4600Configurable seed-offset mesh axes
PyTorch issue #196562, PR #196072CPU init resets the RNG state
TorchTitan #5583Tied RoPE cache
PyTorch #201840trunc_normal_ without the all-gather on DTensors
TorchTitan #5627Option to allocate FSDP’s all-gather buffers on the compute stream (unshard_async_op)
SonicMoE issue #75Dynamic scheduler disabled with no option to enable it
Jupiter (reported to JSC support)Unattributed host memory at job start

‍

Conclusioni e prossimi passi

La configurazione finale per il modello EUROPA raggiunge circa 1.300 token/s per GPU su 384 GPU GH200 per replica DP, 14 volte la configurazione di partenza. A questo ritmo, il job elabora circa 43 miliardi di token al giorno per replica. Con quattro repliche data-parallel su 1.536 GPU, e con l’efficienza di scalabilità dell’89% che abbiamo misurato, ne elaborerebbe circa 150 miliardi al giorno. Con questo lavoro abbiamo dimostrato che un sistema pubblico, condiviso e di generazione Hopper, a cui si accede come semplici utenti, può ospitare un addestramento su scala di mille miliardi di parametri: un risultato significativo per l’ecosistema europeo dell’AI.

La configurazione a cui siamo arrivati è costruita attorno ai due limiti principali che abbiamo individuato: la memoria host per GPU e la banda tra nodi.

Ognuno dei due limiti ha richiesto una serie di interventi specifici. Per la banda tra nodi, l’aumento della lunghezza delle sequenze e il passaggio agli all-gather dei pesi in FP8 e agli all-to-all dei gradienti in BF16 hanno aumentato il rapporto tra calcolo e scambio di dati. Il resharding locale al nodo ha spostato su NVLink il nuovo gather dei pesi densi nel backward pass, e NCCL_MAX_CTAS ha liberato SM per i kernel di calcolo. Per la memoria host, abbiamo spostato attivazioni e stato dell’ottimizzatore su Grace, mantenendo Newton-Schulz sulla GPU, e abbiamo individuato le cause della perdita di memoria host. Vogliamo sottolineare che diversi di questi problemi emergono solo con l’offload su CPU a queste dimensioni del modello, un percorso in generale meno battuto per PyTorch e TorchTitan. Ciononostante, le correzioni sono già state integrate in PyTorch e TorchTitan, a testimonianza della vitalità di questi progetti e dell’impegno verso la loro comunità.

Questo articolo lascia diversi aspetti ancora da esplorare. Non abbiamo documentato quanto il routing limitato per gruppo incida a lungo termine sulla qualità del modello, se le GEMM degli esperti possano essere eseguite in FP8 senza perdite misurabili, se valori di NCCL_MAX_CTAS diversi da 16 diano risultati migliori, né se uno schema basato sul pipeline parallelism sarebbe stato più veloce. Contiamo di riferire su questi punti man mano che l’addestramento procede.

Il modello e tutte le sue metodologie di addestramento saranno rilasciati in forma aperta con la pubblicazione di EUROPA.

‍

Ringraziamenti

Ringraziamo il Jülich Supercomputing Centre per il supporto, in particolare nell’analisi della memoria host, e i maintainer di PyTorch e TorchTitan per aver revisionato e integrato le nostre pull request. EUROPA è finanziato dall’EuroHPC Joint Undertaking nell’ambito della AI-BOOST Frontier AI Grand Challenge.

‍

Riferimenti bibliografici

  1. Guo, W., Mishra, M., Cheng, X., Stoica, I., & Dao, T. (2026). SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations. https://arxiv.org/abs/2512.14080
  2. Jordan, K., Jin, Y., Boza, V., You, J., Cesista, F., Newhouse, L., & Bernstein, J. (2024). Muon: An optimizer for hidden layers in neural networks. https://kellerjordan.github.io/posts/muon/
  3. Kimi Team. (2025). Kimi K2: Open Agentic Intelligence. arXiv Preprint arXiv:2507.20534.
  4. Kimi Team, Bai, T., Bai, Y., Bao, Y., C., M., Cai, J., Cai, X., Cao, P., Cao, Y., Chai, Z., Charles, Y., Che, H. S., Chen, G., Chen, G., Chen, G., Chen, H., Chen, J., Chen, J., Chen, J., … Zu, X. (2026). Kimi K3: Open Frontier Intelligence. https://arxiv.org/abs/2607.24653
  5. Liu, J., Su, J., Yao, X., Jiang, Z., Lai, G., Du, Y., Qin, Y., Xu, W., Lu, E., Yan, J., & others. (2025). Muon is Scalable for LLM Training. arXiv Preprint arXiv:2502.16982.
Share this post
Scopri le ultime novità di iGenius
Ricevi le ultime notizie dal mondo dell'AI e aggiornamenti sui nostri prodotti.