Accelerate ND‑Parallel: Guía para Entrenamiento Multi‑GPU Eficiente

TL;DR

• Accelerate añade herramientas ND‑Parallel para componer múltiples estrategias de paralelismo (DP, FSDP, TP, CP) en un único script de entrenamiento.
• La integración con Axolotl ofrece configs de ejemplo y campos para activar combinaciones sin reescribir el código base.
• Parámetros clave: ParallelismConfig (dp_replicate_size, dp_shard_size, tp_size) y tensor_parallel_size en Axolotl; sharding de contexto para secuencias muy largas.

Contexto y antecedentes

Entrenar modelos grandes en múltiples GPUs requiere elegir y combinar estrategias de paralelismo para equilibrar memoria, computación y comunicación. El paralelismo por datos (DP) replica el modelo y el estado del optimizador en cada dispositivo y reparte los lotes entre réplicas; los gradientes se sincronizan antes de la actualización. DP mejora el rendimiento por paso pero exige que el modelo quepa en un solo dispositivo. Si el modelo no cabe en una GPU, Fully Sharded Data Parallel (FSDP) fragmenta (shards) pesos, gradientes y estados del optimizador entre GPUs, reuniendo parámetros de una capa cuando es necesario. Tensor Parallelism (TP) hace sharding permanente de parámetros y divide la computación de capas lineales grandes. Context Parallelism (CP) shardea la secuencia de entrada para manejar contextos extremadamente largos. Accelerate, junto con Axolotl, ofrece un enfoque integrado y configurable para componer estas estrategias y escalar desde entrenamientos en una sola GPU hasta configuraciones multi‑node, controlando compensaciones entre memoria y comunicación. El repo de Accelerate contiene un script end‑to‑end y Axolotl tiene configs de ejemplo para empezar.

Qué hay de nuevo

• Accelerate expone una manera rápida de usar cualquier combinación de estrategias de paralelismo desde un único script, mediante ParallelismConfig y plugins.
• Axolotl integra técnicas ND‑Parallel para simplificar fine‑tuning a escala y componer paralelismos con distintas técnicas; basta añadir campos ND‑Parallel en el archivo de config de Axolotl.
• Se proporcionan configs de ejemplo y scripts end‑to‑end que demuestran cómo configurar dataloaders, optimizadores y el bucle de entrenamiento mientras se combinan estrategias.

Por qué importa (impacto para desarrolladores/empresas)

• Iteración más rápida: los desarrolladores pueden probar combinaciones de DP, FSDP, TP y CP sin reescribir la infraestructura de entrenamiento.
• Escalar a modelos más grandes: componer sharding y model parallel permite entrenar modelos que no caben en una sola GPU o un solo nodo.
• Control de costes y recursos: ajustando réplicas, tamaños de shard y tamaños de grupo TP, los equipos pueden intercambiar uso de memoria por coste de comunicación y adaptar configuraciones a la topología del clúster (NVLink intra‑nodo vs Infiniband inter‑nodo).
• Reutilización: los campos de config de Axolotl y ParallelismConfig de Accelerate facilitan reutilizar el mismo script en distintos entornos.

Detalles técnicos o implementación

Estrategias de paralelismo clave

| Estrategia | Qué hace | Cuándo usarla | Principales compensaciones |---|---|---|---| | Paralelismo de Datos (DP) | Replica modelo y estado del optimizador; divide lotes y sincroniza gradientes | Cuando el modelo cabe en un dispositivo y se busca mayor throughput | Baja comunicación por paso (all‑reduce) pero requiere que el modelo quepa en el dispositivo; dp_replicate_size controla réplicas |Fully Sharded Data Parallel (FSDP) | Sharde pesos, gradientes y estado del optimizador; reúne parámetros por capa | Cuando el modelo es demasiado grande para una GPU | Cambia uso máximo de memoria por comunicación (gather por capa); la granularidad de gather afecta memoria vs comunicación; dp_shard_size configura el grado |Tensor Parallelism (TP) | Sharde permanentemente grandes capas lineales entre dispositivos | Cuando capas individuales son demasiado grandes para gather en FSDP | Reducción de memoria proporcional al tamaño del grupo TP, pero requiere sincronización frecuente de activaciones; efectivo dentro de un nodo; tp_size (tensor_parallel_size en Axolotl) define el tamaño |Context Parallelism (CP) | Sharde la entrada por la dimensión de secuencia para procesar contextos muy largos | Cuando se entrena con secuencias muy largas (cientos de miles a millones de tokens) | Reduce memoria de atención por dispositivo al dividir la matriz de atención; cada query aún necesita K y V completos para el softmax |

Composición de estrategias y consideraciones de clúster

• DP es una estrategia de nivel superior: dp_replicate_size=2 crea dos réplicas completas, cada una puede a su vez ser shardeada por FSDP o TP.
• FSDP puede aplicarse a varios nodos tratando todos los dispositivos como un único dominio de sharding (por ejemplo, 4 nodos × 8 GPUs = 32 dispositivos). La comunicación entre nodos es más costosa, por lo que a menudo se evita FSDP más allá de un nodo completo.
• TP es más eficaz dentro de un único nodo debido a la necesidad de comunicación de alta velocidad (NVLink). TP entre nodos puede ser ineficaz.
• Para contextos muy largos, la atención crece cuadráticamente con la longitud de la secuencia; el post muestra que secuencias enormes (por ejemplo, 128k tokens) pueden hacer que las activaciones de atención sean masivas, justificando CP.

Parámetros en Accelerate y Axolotl

• ParallelismConfig en Accelerate expone dp_replicate_size, dp_shard_size y tp_size para configurar DP, FSDP y TP y su composición.
• En Axolotl, existen campos similares (tensor_parallel_size, dp_shard_size) por lo que añadir técnicas ND‑Parallel a una config suele ser sólo sumar campos.
• El repo de Accelerate incluye un FullyShardedDataParallelPlugin preparado y ejemplos end‑to‑end.

Puntos clave

• Accelerate + Axolotl permiten componer DP, FSDP, TP y CP en un único flujo de entrenamiento.
• Usar DP cuando el modelo quepa en un dispositivo; aumentar dp_replicate_size para más throughput.
• Usar FSDP para sharding de estado cuando el modelo no quepa en un dispositivo; ajustar la granularidad de gather.
• Usar TP para shardear grandes capas dentro de un nodo cuando el gather de FSDP sea inviable.
• Usar CP para entrenar con secuencias muy largas.

Referencias

• Blog post: https://huggingface.co/blog/accelerate-nd-parallel