17 · Repartir un model entre milers de xips

El capítol anterior parlava de setmanes de càlcul en desenes de milers de xips. Per què tants? I com es posen d'acord tants xips per entrenar un sol model? Aquest capítol és l'enginyeria que fa possible el preentrenament: el problema no és només que el càlcul sigui enorme, sinó que el model no hi cap, i repartir-lo té un preu.

Els xips en qüestió són GPU (graphics processing units): processadors pensats per fer moltes multiplicacions alhora, que és exactament el que fa un model. Una de les més usades per entrenar, la H100, té 80 GB de memòria.

1. Per què no hi cap#

Per fer servir un model només cal guardar-ne els pesos, els milers de milions de números que el defineixen. En el format habitual, cada número ocupa 2 bytes: un model de 70.000 milions de paràmetres ocupa uns 140 GB. Ja no cap en una GPU, però en dues, els pesos sí que hi caben.

Entrenar-lo és una altra cosa. Recorda el capítol 7: l'optimitzador que fa servir tothom, Adam, recorda per a cada pes en quina direcció s'ha anat movent i com de fort sol ser el pendent. A cada pes, doncs, s'hi afegeix el seu gradient (cap on s'ha de moure), una còpia més precisa del mateix pes (perquè els canvis petits no es perdin en arrodonir) i les dues memòries d'Adam. En total, uns 16 bytes per paràmetre: vuit vegades el que cal per fer-lo servir.

Un model de 70.000 milions necessita així uns 1,1 terabytes només per desar el seu estat, abans de fer cap càlcul, i encara hi falten els resultats intermedis de la feina del moment (les activacions). No hi ha cap GPU on càpiga. La conclusió canvia la manera de pensar el problema: repartir el model entre moltes GPU no és una qüestió d'anar més de pressa, sinó de poder fer-ho.

2. Cinc maneres de repartir la feina#

Pensa en una cuina industrial que ha de preparar milions de plats amb un receptari que no cap en cap prestatgeria. Hi ha diverses maneres d'organitzar-la, i cadascuna resol un problema i en crea un altre.

La primera, repartir les dades (data parallelism), és la més senzilla: cada GPU té una còpia sencera del model i processa una part diferent del lot de text; al final de cada pas, totes posen en comú cap on s'han de moure els pesos i en fan la mitjana. Escala molt bé, però no resol res de la memòria, perquè cada GPU continua guardant-ho tot.

La segona la resol: en comptes de copiar l'estat del model a cada GPU, repartir-lo. Les tècniques que ho fan es diuen ZeRO (zero redundancy optimizer: que cap GPU no guardi res repetit) i FSDP (fully sharded data parallel: repartir les dades i, a més, trossejar tot l'estat). Cada GPU en guarda només un tros de cada capa i, quan ha de calcular-ne una, demana a les altres els trossos que li falten, la fa servir sencera i la llença. Amb vuit GPU, cadascuna guarda exactament una vuitena part. El preu és la conversa: les GPU es passen les capes a cada pas, i parlen una vegada i mitja el que parlaven repartint només les dades. Avui és l'opció per defecte per entrenar i ajustar models grans.

La tercera, el paral·lelisme de tensor (tensor parallelism), parteix cada capa entre diverses GPU: cadascuna fa una part de les multiplicacions de la mateixa capa. El problema és que, abans de passar a la capa següent, han de reunir els resultats, i això passa diverses vegades a cada capa. Només surt a compte si les GPU estan unides per un cable molt ràpid, i per això gairebé sempre es fa només entre les vuit GPU d'una mateixa màquina (un node), unides per una connexió interna (NVLink). Entre màquines, la connexió és unes nou vegades més lenta per GPU, i la conversa trigaria més que el càlcul.

La quarta és una cadena de muntatge (pipeline parallelism): cada grup de GPU s'encarrega d'unes quantes capes i passa el resultat al grup següent. Parlen poc, però té un problema propi. Al principi, les últimes etapes no tenen res a fer fins que els arriba la primera feina; al final, les primeres ja han acabat i esperen. Aquest temps mort és la bombolla, i la manera de reduir-la és partir el lot en molts trossos petits, perquè la cadena estigui plena més estona.

Encara n'hi ha dues més. Amb textos molt llargs, es pot repartir el text mateix entre GPU (el paral·lelisme de context). I en un model de mescla d'experts (capítol 13) es reparteixen els experts, i cada paraula viatja a la GPU on viu el seu expert.

Els entrenaments de debò ho combinen tot: capes partides dins de cada màquina, una cadena de muntatge entre màquines, i el conjunt copiat moltes vegades per repartir-se les dades. Trobar la millor combinació per a un model i un clúster és un problema d'optimització en si mateix, i hi ha eines que el resolen automàticament.

3. Estalviar memòria#

Repartir no és l'única manera de fer-hi cabre el model. N'hi ha unes quantes més, i cada una canvia memòria per alguna altra cosa.

La més important canvia memòria per temps. Per corregir els pesos, el model ha de recórrer les capes al revés, i per fer-ho necessita els resultats intermedis de l'anada. Normalment els desa tots, com un excursionista que fotografia cada cruïlla per saber tornar. L'alternativa és desar-ne només uns quants i, a la tornada, refer el tros que falta just quan cal. Se'n diu activation checkpointing.

Les xifres habituals: fins a deu vegades menys memòria per als resultats intermedis, a canvi d'uns 30 % més de temps. El 30 % no és casual: tornar enrere costa aproximadament el doble que anar endavant, de manera que refer una anada afegeix un terç a la feina. També es pot fer de manera selectiva: desar el que és car de refer i refer el que és barat.

Les altres maneres: moure part de l'estat a la memòria normal de l'ordinador o al disc (offloading), que salva entrenaments impossibles però és lent; fer els càlculs amb números de menys precisió, per exemple d'1 byte en comptes de 2; i fer servir optimitzadors que recorden menys coses per pes. El cas extrem d'aquest últim és no entrenar tots els pesos, sinó només una peça petita afegida al model, que és el que fa LoRA i que veurem al capítol 20: s'estalvia el 99 % de les memòries de l'optimitzador.

4. Quina part del temps es fa feina útil#

Una cuina amb mil cuiners no cuina mil vegades més de pressa que amb un: es perd temps esperant ingredients, passant-se plats, esperant que el del costat acabi. Amb les GPU passa el mateix, i la mesura que ho diu té nom: l'MFU (model FLOPs utilization, l'aprofitament de les operacions de càlcul del model), la fracció del temps en què les GPU fan càlcul útil al seu ritme màxim.

Un MFU d'entre el 35 i el 55 % es considera bo en un entrenament gran; per sota del 25 %, alguna cosa va malament. Amb models de mescla d'experts i textos molt llargs, és normal que sigui més baix. El temps es perd en converses entre GPU que no s'han pogut fer mentre es calculava, en bombolles, en càlculs mal aprofitats, en esperar que arribi el lot següent de dades i, a aquesta escala, en fallades: quan milers de GPU treballen setmanes seguides, que se n'espatlli alguna és cosa de cada dia. Per això cal desar l'estat sovint, detectar automàticament les màquines defectuoses i saber tornar a arrencar de pressa.

La lliçó pràctica és més de mètode que de llista: quan l'MFU és baix, primer cal mesurar on se'n va el temps, i arreglar el que pesa més. Cada pèrdua que es treu fa que les altres pesin més, i el sospitós següent canvia.

5. El maquinari en una pàgina#

Una GPU té dues velocitats que importen: la velocitat a la qual calcula i la velocitat a la qual pot portar dades de la seva memòria. Quina de les dues mana depèn de quanta feina es fa amb cada dada que arriba.

D'aquí ve la frase que resumeix el maquinari: en fer servir un model, mana la memòria; en entrenar-lo, mana el càlcul. Escriure una paraula per a un sol usuari fa servir cada pes per a dues operacions i prou, i la GPU es passa el temps esperant que arribin. Per això els servidors agrupen molts usuaris alhora (el capítol 15): cada pes que arriba serveix per a més paraules.

Les GPU actuals tenen entre 80 i més de 190 GB de memòria, que es llegeix a uns quants terabytes per segon. Dins d'una màquina, les GPU es parlen per una connexió de centenars de GB per segon; entre màquines, per una xarxa més lenta. Aquesta distància decideix quines maneres de repartir el model són viables on, i per això el paral·lelisme de tensor es queda dins de la màquina. Hi ha alternatives a les GPU: els xips de Google (TPU), pensats per treballar en grups grans amb la seva pròpia xarxa, i xips especialitzats a respondre molt de pressa. I hi ha un límit cada cop més real: l'energia que tot això consumeix.

Què t'has d'endur#

• Entrenar un model demana uns 16 bytes per paràmetre, vuit vegades el que cal per fer-lo servir. Cap model gran no cap en una GPU: repartir-lo és condició per entrenar-lo.
• Repartir les dades no estalvia memòria; repartir l'estat (ZeRO, FSDP) sí, a canvi de més conversa entre GPU.
• Partir cada capa demana una connexió molt ràpida, i per això es fa dins d'una màquina; la cadena de muntatge parla poc però té bombolles, que es redueixen partint el lot en molts trossos.
• Refer càlculs en comptes de desar-los estalvia molta memòria a canvi d'un terç més de temps.
• L'MFU diu quina part del temps es fa feina útil. Entre el 35 i el 55 % és bo, i per millorar-lo primer cal mesurar on es perd.
• Fent servir un model, mana la velocitat de la memòria; entrenant-lo, la velocitat de càlcul.

Comprova que ho has entès#

Per anar més lluny#