Etapas de DeepSpeed ZeRO: ¿qué estado del entrenamiento se reparte?
Traducción automática
Este artículo se tradujo automáticamente a partir de la versión original en inglés.
ZeRO-1 reparte el estado del optimizador entre las GPUs que trabajan en paralelo por datos; ZeRO-2 reparte también los gradientes, y ZeRO-3, además, los parámetros del modelo. Prueba la etapa más baja que permita ejecutar la carga de trabajo dentro de la memoria disponible con un rendimiento aceptable. Una etapa superior ahorra más memoria del estado del modelo, pero puede requerir más comunicación.
Estas etapas cambian dónde se almacena el estado del entrenamiento. No reducen automáticamente la memoria de las activaciones.
Un ejemplo de memoria con supuestos explícitos
Supongamos que el entrenamiento utiliza pesos y gradientes en FP16, junto con pesos maestros de Adam en FP32 y dos momentos del optimizador. Esto asigna 2 bytes a los pesos, 2 a los gradientes y 12 al estado del optimizador por parámetro: 16 bytes en total.
Para 7.5 mil millones de parámetros en ocho GPUs, el cálculo del estado del modelo es:
| Etapa | Estado repartido entre GPUs | GB del estado del modelo por GPU |
|---|---|---|
| Paralelismo de datos convencional | Ninguno | 120 |
| ZeRO-1 | Estado del optimizador | 41.25 |
| ZeRO-2 | Estado del optimizador y gradientes | 28.125 |
| ZeRO-3 | Los tres componentes | 15 |
Son GB decimales. No incluyen las activaciones, los parámetros reunidos temporalmente, la sobrecarga del asignador de memoria ni los búferes del runtime. Una pila de entrenamiento que almacene los datos con otras precisiones necesita un cálculo distinto.
El artículo de ZeRO explica las particiones del estado y analiza su comunicación. ZeRO-3 reconstruye los parámetros necesarios mediante operaciones all-gather; el análisis del reparto completo del artículo alcanza aproximadamente 1.5 veces el volumen de comunicación del paralelismo de datos convencional. Esa proporción procede de un análisis bajo sus propios supuestos; no es una garantía para cualquier red y configuración.
Elegir según el perfil de memoria
Si los parámetros caben en memoria y los estados del optimizador causan el fallo, prueba la etapa 1 o 2. Si no caben los parámetros completamente replicados, prueba la etapa 3. Si predominan las activaciones, combina la etapa elegida con checkpointing de activaciones o con un lote de entrenamiento más pequeño.
DeepSpeed también permite trasladar el estado a CPU y NVMe. Consulta la configuración actual de ZeRO y mide el coste de transferencia en el equipo y el almacenamiento reales. La capacidad adicional obtenida con este traslado puede ayudar, pero las transferencias lentas pueden reducir el rendimiento.
La sección sobre ZeRO de la LLM Engineering Guide ofrece un cálculo ejecutable para otras cantidades de parámetros y grupos de GPU.