ParalelismoRuta Optimización
ZeRO y FSDP
Leerlo en el Capítulo 4 · checkpoint +50 XP →
Reparte el estado del optimizador, los gradientes y los parámetros entre los rangos de paralelismo de datos en lugar de replicarlos.
El paralelismo de datos simple desperdicia memoria: cada rango guarda copias idénticas de todo. ZeRO (y FSDP de PyTorch) lo reparten:
| Etapa | Repartido | Estado por GPU (Adam en precisión mixta) |
|---|---|---|
| 0 | nada | 16 B/parámetro |
| 1 | estados del optimizador | 4 + 12/N B/parámetro |
| 2 | + gradientes | 2 + 14/N B/parámetro |
| 3 / FSDP | + parámetros | 16/N B/parámetro |
El problema de la etapa 3: los parámetros deben reunirse con all-gather antes del forward y del backward de cada capa, en cada micro-batch. Para un modelo de 70B son ~140 GB movidos dos veces por micro-batch. Sobre NVLink va bien. Sobre una NIC de 400 Gb/s puede dominar el paso.
Guía práctica
- ZeRO-1 es casi gratis y casi siempre vale la pena.
- ZeRO-3/FSDP es excelente dentro de un nodo o en redes rápidas, y para hacer caber un modelo que de otro modo no cabría.
- A gran escala, combínalo: TP dentro del nodo, ZeRO-1 o sharding híbrido (repartir dentro del nodo, replicar entre nodos).
Úsalo en una misión
Relacionado
Paralelismo de datos (DDP)HBM y la jerarquía de memoriaParalelismo tensorialCheckpointing de activaciones
Fuentes y lecturas adicionales
- ZeRO (Rajbhandari et al., 2020)
- Distributed Training (T. Tsugriy) — parallelism strategies, collectives, memory
Escrito para LazyLab con nuestras propias palabras; las fuentes se citan para profundizar y verificar.