Entrenar y servir modelos de transformadores grandes a escala es fundamentalmente un problema de gestión de memoria. Cada GPU en un clúster tiene una cantidad fija de VRAM y, a medida que crecen los tamaños de los modelos y la duración del contexto, los ingenieros tienen que hacer concesiones constantemente sobre cómo distribuir el trabajo entre el hardware. Una nueva técnica de Zyphra, llamada Tensor and Sequence Parallelism (TSP), ofrece una manera de repensar esa compensación, y en pruebas comparativas en hasta 1,024 GPU AMD MI300X, ofrece consistentemente una memoria máxima por GPU más baja que cualquiera de los esquemas de paralelismo estándar utilizados hoy en día, tanto para cargas de trabajo de entrenamiento como de inferencia.
El problema que TSP está resolviendo
Para comprender por qué TSP es importante, primero debe comprender las dos estrategias de paralelismo que reúne.
El paralelismo tensorial (TP) divide los pesos del modelo entre las GPU. Si tiene una matriz de peso en una capa de atención o MLP, cada GPU en el grupo TP contiene solo una fracción de esa matriz. Esto reduce directamente la memoria por GPU ocupada por parámetros, gradientes y estados del optimizador: la memoria de “estado del modelo”. La desventaja es que TP requiere operaciones de comunicación colectiva (normalmente pares de reducción total o reducción-dispersión/reunión total) cada vez que se calcula una capa. Esta comunicación es proporcional al tamaño de la activación, por lo que se vuelve cada vez más costosa a medida que crece la longitud de la secuencia.
El paralelismo de secuencia (SP) adopta un enfoque diferente. En lugar de dividir pesos, divide la secuencia de tokens de entrada entre las GPU. Cada GPU procesa solo una fracción de los tokens, lo que reduce la memoria de activación y el costo cuadrático del cálculo de la atención. Sin embargo, SP deja los pesos del modelo completamente replicados en cada GPU, lo que significa que la memoria de estado del modelo permanece exactamente igual independientemente de cuántas GPU agregue al grupo de SP.
En el paralelismo multidimensional estándar, los ingenieros combinan TP y SP colocándolos en ejes ortogonales de una malla de dispositivo. Si desea un grado TP de T y un grado SP de Σ, la réplica de su modelo consume T.Σ GPU. Esto es caro en dos sentidos. En primer lugar, utiliza más GPU para el grupo de modelos paralelos, lo que deja menos disponibles para las réplicas de datos paralelos. En segundo lugar, si T.Σ es lo suficientemente grande como para abarcar múltiples nodos, parte de la comunicación colectiva tiene que viajar a través de interconexiones entre nodos más lentas como InfiniBand o Ethernet en lugar de la estructura intranodo de alto ancho de banda, como AMD Infinity Fabric o NVIDIA NVLink. El paralelismo de datos (DP), la otra línea de base común, evita por completo estos costos de modelo paralelo, pero replica todo el estado del modelo en cada dispositivo, lo que lo hace poco práctico para modelos grandes o contextos prolongados por sí solo.
Lo que realmente significa plegar
La idea central de TSP es el plegado en paralelismo: en lugar de colocar TP y SP en dimensiones de malla ortogonales separadas, colapsa ambos en un único eje de malla de dispositivo de tamaño D. Cada GPU en el grupo TSP contiene simultáneamente 1/D de los pesos del modelo y 1/D de la secuencia de tokens. Debido a que ambas están fragmentadas en las mismas GPU D, la huella de memoria por dispositivo disminuye en 1/D tanto para la memoria de parámetros como para la memoria de activación, algo que ningún esquema de paralelismo estándar logra por sí solo. Por lo tanto, TSP es el único esquema que reduce simultáneamente la memoria proporcional al peso (parámetros, gradientes, estados del optimizador) y la memoria de activación por el mismo factor 1/D en un solo eje sin requerir un diseño de dispositivo T.Σ bidimensional.
El desafío es que si cada GPU solo tiene parte de los pesos y parte de la secuencia, necesita coordinarse con otras GPU para completar el avance de cada capa. TSP utiliza dos programas de comunicación diferentes para manejar esto, uno para atención y otro para el MLP cerrado.
Para llamar la atención, TSP itera sobre fragmentos de peso. En cada paso, una GPU transmite sus fragmentos de peso de atención empaquetados (WQ, WK, WV y WO) a todas las demás GPU del grupo. Luego, cada GPU aplica esos pesos a sus tokens de secuencia local para calcular las proyecciones locales Q, K y V. Dado que la atención causal requiere acceso al contexto clave/valor completo, los tensores K y V locales se reúnen en todo el grupo TSP y se reordenan utilizando un esquema de partición en zigzag antes de aplicar FlashAttention. La partición en zigzag garantiza que la carga de trabajo de atención causal esté equilibrada entre los rangos, ya que los tokens posteriores atienden a prefijos más grandes y, de lo contrario, provocarían un desequilibrio de carga.
Para el MLP cerrado, TSP utiliza un programa de timbre. Cada GPU comienza con fragmentos locales de la puerta, proyecciones hacia arriba y hacia abajo. Estos fragmentos de peso circulan por el grupo TSP mediante operaciones de envío/recepción punto a punto, y cada GPU acumula salidas parciales localmente a medida que llegan los fragmentos. Fundamentalmente, esto elimina la reducción total que requiere el TP estándar para la salida MLP: la secuencia permanece local y solo se mueven los pesos. El anillo está diseñado para superponer las transferencias de peso con el cálculo GEMM, por lo que la comunicación se produce en segundo plano mientras la GPU está computando.
Resultados de memoria y rendimiento
Probado en un único nodo MI300X de 8 GPU en longitudes de secuencia de tokens de 16K a 128K, TSP logra el pico de memoria más bajo en cada punto. En tokens de 16K, TSP y TP son casi equivalentes, 31,0 GB frente a 31,5 GB por GPU, porque la memoria de estado del modelo domina en el contexto corto. Con 128.000 tokens, el panorama cambia drásticamente: TSP usa 38,8 GB por GPU, en comparación con 70,0 GB para TP y 85,0 GB y 140,0 GB para dos factorizaciones diferentes de TP+SP en el mismo nodo. Las cifras teóricas a lo largo de esta investigación se basan en un modelo de referencia de transformador denso solo decodificador 7B (dimensión oculta h=4096, 32 capas, 32 cabezales de consulta, cabezales de 32 KV, factor de expansión FFN F=4, precisión bf16), lo que proporciona una línea de base reproducible para comparar los esquemas.
Los resultados de rendimiento en 128 nodos completos (1024 GPU MI300X) muestran que TSP supera consistentemente las líneas base de TP+SP coincidentes. Con un grado de plegado de D=8 y una longitud de secuencia de 128 000 tokens, TSP logra 173 millones de tokens por segundo en comparación con 66,30 millones de tokens por segundo para la línea base TP+SP coincidente (aproximadamente una aceleración de 2,6 veces). La ventaja aumenta con un mayor grado de paralelismo y una mayor longitud de secuencia.
Compensaciones prácticas que hay que entender
TSP aumenta el volumen total de comunicación en comparación con TP solo. Agrega un término de movimiento de peso por capa además del mismo conjunto K/V proporcional a la activación que utiliza SP. Sin embargo, el equipo de investigación muestra que cuando el tamaño de lote B y la longitud de secuencia S satisfacen BS > 8h (donde h es la dimensión de incrustación del modelo), el volumen de comunicación directa de TSP es competitivo con el de TP. Esta condición se cumple en la mayoría de los escenarios de inferencia y entrenamiento de contexto largo.
La idea clave que enfatiza el equipo de Zyphra es que el volumen de comunicación y el costo de la comunicación no son lo mismo. Que un volumen de comunicación adicional se traduzca en una desaceleración del reloj de pared depende de si los colectivos están limitados por la latencia o por el ancho de banda, y cuánto de ese tráfico puede superponerse con la multiplicación de matrices. Su implementación canaliza transferencias de peso detrás de las operaciones GEMM dominantes, de modo que la comunicación de peso consume ancho de banda sin aumentar el tiempo de ruta crítica.
TSP no está diseñado para reemplazar TP, SP o TP+SP en todas las configuraciones. Está pensado como un eje adicional en el espacio de diseño de paralelismo multidimensional. Se compone ortogonalmente con paralelismo de canalización, paralelismo experto y paralelismo de datos. Esto significa que los equipos pueden ubicar TSP en una configuración de paralelismo existente donde el diseño estándar obligaría a grupos de modelos paralelos a través de enlaces entre nodos más lentos.
Conclusiones clave
El tensor y el paralelismo de secuencia (TSP) de Zyphra pliega el paralelismo del tensor y el paralelismo de secuencia en un único eje de malla de dispositivo, de modo que cada GPU contiene simultáneamente 1/D de los pesos del modelo y 1/D de la secuencia de tokens, lo que reduce la sobrecarga de memoria tanto para el entrenamiento como para la inferencia. TSP es el único esquema de paralelismo que reduce tanto la memoria proporcional al peso (parámetros, gradientes, estados del optimizador) como la memoria de activación por el mismo factor 1/D en un solo eje, sin requerir una malla de dispositivo T.Σ bidimensional. Los resultados empíricos en un único nodo MI300X de 8 GPU muestran que TSP utiliza 38,8 GB por GPU con una longitud de secuencia de 128 K, en comparación con 70,0 GB para TP y 85,0–140,0 GB para configuraciones TP+SP. A gran escala (1024 GPU MI300X, contexto de 128K, D=8), TSP logra 173 millones de tokens por segundo frente a 66,30 millones de tokens por segundo para una línea base TP+SP coincidente (aproximadamente una ventaja de rendimiento de 2,6 veces). TSP se compone ortogonalmente con paralelismo de canalización, expertos y datos y es más adecuado para cargas de trabajo de inferencia y entrenamiento de contexto largo y con memoria limitada, donde la eliminación del peso y la replicación de activación superan el volumen de comunicación agregado.
Consulte el documento y los detalles técnicos. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 130.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.
¿Necesita asociarse con nosotros para promocionar su repositorio de GitHub O su página principal de Hugging O su lanzamiento de producto O seminario web, etc.? Conéctate con nosotros