FSDP frente a DeepSpeed ZeRO: ¿cuál encaja en mi configuración de entrenamiento?
Traducción automática
Este artículo se tradujo automáticamente a partir de la versión original en inglés.
Empieza con FSDP cuando una infraestructura de entrenamiento en PyTorch necesite particionar el estado de forma distribuida. Prueba DeepSpeed si necesitas sus etapas ZeRO seleccionables o sus funciones de descarga a CPU y NVMe. Ningún framework ofrece siempre el mejor rendimiento; compara la configuración que necesitas en tu clúster.
Ambos pueden repartir los parámetros, los gradientes y el estado del optimizador entre GPUs. Sus APIs, políticas y gestión de checkpoints son diferentes.
El particionado completo sigue el mismo proceso básico
Con el particionado completo de parámetros, cada GPU almacena una parte del modelo. Antes de calcular un grupo de capas, el framework reúne los parámetros que necesita ese grupo. Después realiza el cálculo, libera los parámetros reunidos según su política de reparticionado y combina y reparte los gradientes mediante reduce-scatter.
Las decisiones sobre agrupación y precarga controlan la memoria temporal y el solapamiento de la comunicación con el cálculo. Particionar cada operación pequeña por separado puede hacer que la comunicación domine el trabajo; reunir grupos muy grandes puede provocar un pico temporal de memoria.
| Requisito | Candidato que probar |
|---|---|
| APIs nativas de PyTorch y estado de tensores distribuidos | FSDP2 |
| Particionado del estado del optimizador con parámetros aún replicados | DeepSpeed ZeRO-1 o ZeRO-2 |
| Particionado completo del estado | FSDP o ZeRO-3 |
| Descarga del estado a NVMe | DeepSpeed ZeRO-Infinity |
El tutorial de FSDP2 usa fully_shard y parámetros DTensor. También ofrece políticas de reparticionado: conservar los parámetros después de la pasada hacia delante consume más memoria, pero reduce las operaciones de reunión posteriores. Crea el optimizador después de aplicar las transformaciones de particionado necesarias.
Compara un paso de entrenamiento completo
Mantén constantes el modelo, los datos, la precisión, la longitud de secuencia, el tamaño de lote efectivo y la topología de las GPUs. Mide el pico de memoria asignada, la duración de cada paso y la comunicación. Comprueba que un checkpoint guardado puede restaurar el estado del modelo, el estado del optimizador y el progreso de entrenamiento previsto.
Hugging Face Accelerate puede seleccionar FSDP o DeepSpeed mediante configuraciones de lanzamiento para código preparado para esas integraciones. Eso no implica que cualquier script de entrenamiento pueda cambiar sin ajustes, especialmente en la gestión de checkpoints, la descarga de estado y el tratamiento de optimizadores personalizados.
La comparación de FSDP en la LLM Engineering Guide relaciona estas decisiones con la memoria del estado de entrenamiento.