El entrenamiento de modelos de lenguaje grandes sobre secuencias largas es fundamental para desarrollar sistemas de inteligencia artificial eficientes. A medida que estos modelos se utilizan para tareas como el análisis de documentos, la comprensión de código y el razonamiento complejo, la necesidad de procesar secuencias de cientos de miles e incluso millones de tokens ha aumentado notablemente.
En este contexto, el Ulysses Sequence Parallelism, parte del protocolo Arctic Long Sequence Training (ALST) de Snowflake AI Research, ofrece una solución efectiva al distribuir el cálculo de atención entre múltiples GPUs mediante la paralelización de cabezales de atención. Este artículo explora el funcionamiento de Ulysses y su integración en el ecosistema de Hugging Face, que incluye herramientas como Accelerate y el Transformer Trainer.
El Desafío del Entrenamiento de Secuencias Largas
El mecanismo de atención en los transformadores escala cuadráticamente con la longitud de la secuencia. Esto significa que para una secuencia de longitud n, el cálculo de atención estándar requiere O(n^2) FLOPs y O(n^2) de memoria. Implementaciones optimizadas, como FlashAttention, logran reducir la memoria a O(n) al no materializar la matriz de atención completa, pero el cálculo sigue siendo O(n^2).
El entrenamiento de secuencias muy largas (32k+ tokens) aún supera los límites de memoria de una sola GPU, incluso con FlashAttention. Ejemplos donde el entrenamiento de contexto largo es esencial incluyen:
- Comprensión de documentos: Procesar libros, documentos legales o artículos de investigación.
- Análisis de código: Entender grandes bases de código con múltiples archivos interconectados.
- Tareas de razonamiento: Modelos que generan miles de tokens durante la inferencia.
- Generación aumentada por recuperación: Incorporar múltiples pasajes recuperados en el contexto.
La paralelización de datos tradicional no ayuda en estos casos, ya que cada GPU debe procesar la secuencia completa dentro del bloque de atención. Es necesario un método para dividir la secuencia en múltiples dispositivos.
Cómo Funciona Ulysses
Ulysses Sequence Parallelism (SP), introducido en el artículo de DeepSpeed Ulysses, utiliza un enfoque innovador: además de dividir la dimensión de la secuencia, también particiona los cabezales de atención entre varias GPUs.
El funcionamiento de Ulysses se puede resumir en los siguientes pasos:
- Fragmentación de Secuencia: La secuencia de entrada se divide a lo largo de la dimensión de secuencia entre P GPUs.
- Proyección QKV: Cada GPU calcula las proyecciones de consulta, clave y valor para su fragmento local de secuencia.
- Comunicación All-to-All: Una operación colectiva all-to-all redistribuye los datos para que cada GPU contenga todas las posiciones de secuencia después de las proyecciones, pero solo para un subconjunto de cabezales de atención.
- Atención Local: Cada GPU calcula la atención para sus cabezales asignados utilizando mecanismos de atención estándar.
- Comunicación All-to-All: Una segunda operación all-to-all revierte la redistribución, volviendo al formato de fragmentación de secuencia.
- Proyección de Salida: Cada GPU calcula la proyección de salida para su fragmento local de secuencia.
La clave del enfoque es que los cabezales de atención son independientes, lo que permite que cada cabezal se compute por separado. Al intercambiar la localidad de la secuencia por la localidad del cabezal, Ulysses permite una paralelización eficiente con un bajo costo de comunicación.
