Google ha incorporado paralelismo de contexto a MaxText, su biblioteca de código abierto para entrenar grandes modelos de lenguaje en TPU, y asegura que con ella se pueden entrenar secuencias de hasta 14 millones de tokens en TPU 5p y TPU 7 (Ironwood). La versión 0.2.4 ofrece cuatro estrategias para repartir el contexto entre chips y la presentó el blog de desarrolladores de Google el 9 de noviembre de 2026, firmada por Michael Tran, Nuojin Cheng y Deepak Patil.
Las claves del paralelismo de contexto en MaxText en 20 segundos
- MaxText v0.2.4 añade cuatro estrategias: All-gather, Ulysses, Ring y USP.
- Google entrenó con 14 millones de tokens en 64 chips TPU 7.
- En esa prueba logró un 48,7 % de MFU, frente al 37,1 % de TPU 5p.
Entrenar con contextos largos choca con un límite físico: la memoria y el cálculo de la atención crecen con la longitud de la secuencia, y a partir de cierto punto una secuencia no cabe en un solo chip. El paralelismo de contexto la divide entre varios aceleradores, cada uno procesa su trozo y se intercambian la información necesaria para que la atención vea la secuencia completa. La dificultad está en que esa comunicación no se coma el rendimiento.
Cuatro formas de repartir la secuencia
MaxText ofrece cuatro estrategias, todas apoyadas en el kernel Splash Attention de la biblioteca Tokamax. La de All-gather reúne en cada chip las claves y valores de toda la secuencia, una opción sencilla que consume más memoria. Ulysses usa comunicación de todos a todos (A2A) para que cada chip calcule la atención completa sobre un subconjunto de cabezas. Ring pasa los bloques de claves y valores de chip en chip (P2P) mientras se calcula, solapando comunicación y cálculo. USP combina Ulysses y Ring.
La elección se hace en la configuración mediante el parámetro context_parallel_strategy, con los valores ring, ulysses o usp. Para estirar la memoria disponible, los autores describen un orden de descarga a la memoria del anfitrión: primero las salidas de la atención, después el estado del optimizador, las entradas y los parámetros, y solo al final la rematerialización, que recalcula activaciones a costa de más cómputo.
Las cifras de rendimiento
Google mide la eficiencia con el MFU (utilización de FLOPs del modelo), el porcentaje de la capacidad teórica de cálculo que se aprovecha de verdad. Con Llama 3.1 8B, MaxText mantiene un 70 % de MFU hasta 4 millones de tokens en TPU 5p. En TPU 7 ronda el 50 % por debajo de 6 millones. En la prueba más exigente, 14 millones de tokens en 64 chips TPU 7 con la estrategia Ring, el MFU fue del 48,7 %, frente al 37,1 % en TPU 5p.
Son resultados de Google sobre su propio hardware y con un modelo de 8.000 millones de parámetros, así que no se trasladan directamente a modelos mayores ni a otras configuraciones. Los propios autores recogen limitaciones en la publicación, por lo que conviene revisarlas antes de planificar un entrenamiento real. También hay que distinguir entrenar con contexto largo de que el modelo resultante use bien ese contexto, algo que depende de los datos y de la evaluación, no solo de la infraestructura.
El código está en el repositorio de MaxText en GitHub, dentro de la organización AI-Hypercomputer, y Google publica recetas de configuración en su repositorio tpu-recipes para reproducir las pruebas.
Preguntas frecuentes
¿Qué es MaxText?
Una biblioteca de código abierto de Google para entrenar grandes modelos de lenguaje en TPU.
¿Qué es el paralelismo de contexto?
Una técnica que reparte una secuencia muy larga entre varios chips para que el entrenamiento no quede limitado por la memoria de uno solo.
¿Cuál es el contexto máximo probado?
Google informa de entrenamientos con hasta 14 millones de tokens en TPU 5p y TPU 7 Ironwood.
¿Cómo se activa?
Con el parámetro context_parallel_strategy de la configuración de MaxText v0.2.4, eligiendo ring, ulysses o usp.
Fuentes:
