Este artículo presenta el enfoque de Dynamic Context Parallelism (Dynamic-CP), utilizado en el Megatron Core de NVIDIA para el post-entrenamiento de LLM o el pre-entrenamiento de DiT. Este método selecciona dinámicamente el tamaño de CP por micro-batch para manejar eficientemente secuencias de longitud variable, logrando hasta 1.48x de aceleración en conjuntos de datos reales.
En el entrenamiento de modelos a gran escala, una limitación frecuentemente pasada por alto es la variabilidad en la longitud de las secuencias de los conjuntos de datos. Tanto el entrenamiento de LLM como la generación de videos a gran escala presentan distribuciones de longitud de secuencia que son claramente desiguales. Un pequeño porcentaje de muestras extremadamente largas representa una parte desproporcionada de la carga computacional y el consumo de memoria.
Desafíos en el entrenamiento de LLM
En el entrenamiento de LLM, esto se traduce en una amplia variedad de longitudes de secuencias de texto en los lotes. En la generación de video, videos de alta resolución y varios segundos pueden abarcar decenas de miles de tokens. Esto provoca un uso desigual de FLOPs y memoria a nivel de muestra entre los grupos de datos paralelos, modalidades y micro-batches, dificultando la programación y utilización eficiente de los recursos.
Para gestionar entradas de longitud variable, los sistemas de entrenamiento comúnmente emplean empaquetamiento a nivel de muestra, que combina múltiples secuencias cortas en un solo micro-batch cuya longitud total de tokens está limitada por una longitud de secuencia objetivo. En la Figura 1, las secuencias se empaquetan a una longitud igual.
Problemas de carga computacional
A pesar de que las tres muestras empaquetadas son de la misma longitud, sus cargas computacionales no son equivalentes, como se muestra en la Figura 2, debido a la naturaleza cuadrática de la atención por producto punto. Esta variación en la carga computacional entre muestras empaquetadas se conoce como desbalance computacional a nivel de datos paralelos. Esto provoca inactividad en la GPU, ya que algunos grupos de DP esperan a otros con cargas computacionales más altas para realizar la sincronización de gradientes.
En la Figura 3, el perfil de NVIDIA Nsight Systems captura un desequilibrio en el entrenamiento de VLM. Diferentes muestras de imágenes/video tienen longitud de secuencia variable, y se emplea el empaquetamiento. La captura muestra la sobrecarga de sincronización entre diferentes grupos de DP.
Necesidad de un enfoque dinámico
Al utilizar el paralelismo de contexto, el tamaño del sharding de CP se determina por la secuencia más larga en el lote para evitar errores de falta de memoria en las GPUs. Como resultado, las secuencias más cortas que no requieren paralelismo de contexto se dividen. A pesar de que estas secuencias caben en una única GPU, se dividen debido a una secuencia más larga en el mismo lote, lo que genera una sobrecarga innecesaria de comunicación de CP.
Normalmente, el cálculo oculta la comunicación de CP. Sin embargo, cuando los tamaños de CP son grandes, especialmente cuando la comunicación abarca dominios de InfiniBand (IB), la sobrecarga de comunicación puede hacerse evidente cuando las secuencias empaquetadas son más cortas y la carga computacional es menor. Esto representa una ineficiencia computacional de CP.
Los siguientes ejemplos demuestran la necesidad de un enfoque dinámico para el paralelismo de contexto. En lugar de fijar estáticamente el tamaño de CP a la secuencia más larga en un micro-batch, este enfoque adapta el tamaño de CP utilizando la estrategia de empaquetamiento por micro-batch. Trabajos relevantes, como ByteScale y WLB-LLM, abordan problemas similares.
Ventajas de Dynamic-CP
Estas observaciones resaltan la necesidad de un enfoque dinámico al paralelismo de contexto. Este método permite seleccionar el tamaño de CP en función de la estrategia de empaquetamiento por micro-batch. Cambiar el tamaño de CP requiere re-particionar las porciones de secuencia y re-formar los grupos de comunicación de CP utilizados por las operaciones de atención. En comparación con otros esquemas de paralelismo dinámico, como la adaptación de tamaños tensor-paralelos o pipelínicos basados en la longitud de la secuencia, Dynamic-CP agrega una sobrecarga mínima.
El solucionador está diseñado para, dado un conjunto de secuencias de longitud variable, determinar cómo empaquetarlas y seleccionar el tamaño de CP para maximizar la eficiencia computacional sin exceder los límites de memoria de la GPU. Su función es calcular el empaquetamiento óptimo y el tamaño de CP mientras maximiza la eficiencia computacional y respeta estrictamente las limitaciones de memoria de la GPU.
El siguiente ejemplo muestra el beneficio de utilizar Dynamic-CP. Antes de aplicar el balanceo de carga, el desbalance lleva a burbujas de tubería entre diferentes micro-batches, lo que causa un desbalance de DP entre los rangos de DP. Después del balanceo, las burbujas entre micro-batches y rangos de DP se reducen.
Modificaciones en el marco de Megatron Core para soportar Dynamic-CP
Esta sección introduce el proceso de integración de Dynamic CP en Megatron Core.
Construcción de múltiples grupos de contexto paralelo por rango
Con el paralelismo de contexto estándar, cada rango pertenece a un solo grupo (cp_group) con un cp_size fijo que se determina estáticamente durante la inicialización. Sin embargo, el paralelismo de contexto dinámico tiene un cp_size diferente entre iteraciones y micro-batches.
Para soportar esto, un solo rango debe participar en múltiples grupos de CP de diferentes tamaños. Se construyen múltiples grupos de CP durante la inicialización, con cp_size que varía desde 1 hasta dp × cp, restringido a potencias de dos.
Reprogramación dinámica y empaquetado de datos
A diferencia del preentrenamiento, que típicamente utiliza el formato Batch × Sequence × Head × Dim (BSHD), Dynamic-CP opera en un formato THD. En este formato, las secuencias de longitud variable se empaquetan bajo una restricción de longitud, colapsando la dimensión original BS en una dimensión de tokens T.
Como consecuencia, el número de micro-batches ya no es estático. En el formato BSHD, el número de micro-batches se calcula como num_micro_batches = global_batch_size/dp_size/micro_batch_size. Con el empaquetamiento THD, el número de secuencias originales contenidas en cada secuencia empaquetada no es fijo, lo que provoca que num_micro_batches varíe entre iteraciones.
Difusión entre etapas de tubería y extensión de packedSeqParams
Dado que num_micro_batches varía y solo el rango TP 0, así como las primeras y últimas etapas de PP manejan la programación en Megatron Core, el marco transmite num_micro_batches, max_seqlen, y cu_seqlens a todos los rangos de PP relevantes. Esto asegura una ejecución consistente a través de las etapas de la tubería bajo programación dinámica de micro-batch.
Cálculo de pérdida y cálculo de FLOPs
Dadas las secuencias de longitud variable y el formato THD, diferentes secuencias contribuyen a diferentes números de tokens válidos. Por ello, el cálculo de pérdida se realiza a nivel de token: loss = loss_over_valid_tokens / total_number_valid_tokens, evitando sesgos introducidos por tokens de relleno.
Modelado del programador de datos
La carga de trabajo del transformador escala cuadráticamente con la longitud de la secuencia S 
Colaboración del modelo de costos, solucionador y simulador
Un flujo de trabajo completo del programador consiste en tres componentes:
- El modelo de costos estima el tiempo de ejecución para cada muestra en función de su longitud de secuencia, modelando la carga de trabajo por muestra a través de las operaciones del transformador.
- El solucionador utiliza la salida del modelo de costos como entrada y aplica un algoritmo heurístico para determinar una estrategia de empaquetamiento casi óptima para cada muestra.
- El simulador evalúa estos micro-batches bajo el programa de paralelismo de tuberías distribuidas, seleccionando el plan con el tiempo de ejecución mínimo.
Resultados de benchmarks
Con todas las mejoras introducidas, las burbujas de desbalance causadas por distribuciones de secuencias de longitud variable pueden reducirse sustancialmente.
En la Tabla 1, Dynamic CP se evalúa contra una línea base de empaquetamiento puro bajo la configuración: llama-13B, tamaño de lote global 2048, PP=8, CP=8, y recomputación completa. Se ejecutan 10 iteraciones, descartándose la primera como calentamiento, y el tiempo de iteración se promedia sobre las 9 restantes. Dynamic CP logra 1.48x y 1.25x de aceleración en los conjuntos de datos de GitHub y CommonCrawl, respectivamente.
| Tamaño del Modelo | Tipo de Conjunto de Datos | Método | TFLOPS/GPU |
| Llama 13B | GitHub | Solo Empaquetamiento | 195.88 |
| Llama 13B | GitHub | Dynamic CP | 289.32 |
| Llama 13B | Commoncrawl | Solo Empaquetamiento | 139.17 |
| Llama 13B | Commoncrawl | Dynamic CP | 174.39 |
Tabla 1. Comparación de métodos Dynamic CP y empaquetamiento puro en diferentes conjuntos de datos.
Aprende más
Este artículo demostró que Dynamic CP con el backend de Megatron Core mejora el rendimiento de entrenamiento para secuencias de longitud variable en comparación con el método de CP fijo. Con empaquetamiento de secuencias, paralelismo 4D y núcleos optimizados para GPU, Dynamic CP garantiza una alta eficiencia de entrenamiento a través de escalas de modelos.
Comienza con:
- Megatron Core en GitHub para comenzar a entrenar tu modelo con secuencias de longitud variable usando optimizaciones de Megatron Core.
- El scheduler, que también está en GitHub.
Agradecemos a todos los colegas por sus contribuciones a este proyecto.
