Ulysses SP: entrenar modelos con millones de tokens | Keryc
Entrenar modelos con contextos de libro completo ya no es una curiosidad de laboratorio: es una necesidad práctica para tareas como análisis documental, razonamiento extendido, revisión de código y sistemas RAG. ¿El problema? La atención escala cuadráticamente y una secuencia de cientos de miles de tokens no cabe en una sola GPU.
Qué problema resuelve Ulysses
La atención en transformers requiere calcular puntuaciones entre pares de tokens, lo que hace que memoria y FLOPs crezcan como el cuadrado de la longitud de la secuencia. FlashAttention y optimizaciones similares alivian memoria evitando materializar toda la matriz, pero el cómputo sigue ahí. ¿Qué haces cuando una novela son ~250k tokens y necesitas entrenar con varios documentos a la vez?
La solución tradicional de data parallel no ayuda: cada GPU aún tendría que ver la secuencia completa dentro del bloque de atención. Ulysses Sequence Parallelism (parte del protocolo ALST de Snowflake AI Research) ofrece una forma elegante de repartir la atención entre GPUs mediante paralelismo por cabezas de atención.
Cómo funciona Ulysses Sequence Parallelism (SP)
La idea clave es cambiar la localización: en vez de mantener todas las cabezas y partir la secuencia, Ulysses parte la secuencia entre GPUs y reparte las cabezas entre ellas. Pasos simplificados:
Sequence sharding: la secuencia se divide a lo largo de la dimensión temporal; cada GPU tiene un chunk de tokens.
QKV projection: cada GPU proyecta q, k y v para su chunk local.
All-to-all: una operación colectiva redistribuye datos para que cada GPU tenga todas las posiciones de secuencia pero solo un subconjunto de cabezas.
Atención local: cada GPU calcula la atención para sus cabezas asignadas (FlashAttention o SDPA).
All-to-all: se invierte la redistribución para volver al formato shard de secuencia.
Output projection: cada GPU proyecta la salida para su chunk local.
Esto requiere dos all-to-all por capa de atención. El beneficio: las cabezas son independientes, así que se puede paralelizar con bajo overhead de comunicación si la interconexión lo permite.
Comunicación: Ulysses comunica O(total_seq * hidden / sp_size) por GPU, mientras Ring Attention comunica un factor mayor. All-to-all aprovecha mejor la bisección de banda disponible; Ring serializa por saltos.
Integración con el ecosistema Hugging Face
Accelerate actúa como la base para habilitar Ulysses vía ParallelismConfig y la integración con DeepSpeed. Con accelerator.prepare() el modelo se registra con UlyssesSPAttentionHF y el dataloader se envuelve para manejar el sharding de secuencia y shift_labels.
Transformers Trainer acepta TrainingArguments.parallelism_config y se encarga automáticamente de:
envolver el dataloader con UlyssesSPDataLoaderAdapter
agregar shift_labels y detectar pre-shifted labels
agregar pérdidas por tokens válidos entre ranks y computar la pérdida ponderada
ajustar cálculos de tamaño de batch y longitud del dataloader
TRL SFTTrainer extiende esto con optimizaciones para fine-tuning supervisado de secuencias largas (pre-shifted labels, packing, pad_to_multiple_of igual a sp_size).
Si usas un loop propio con accelerate, debes agregar la agregación ponderada de pérdidas entre SP ranks. En Trainer y SFTTrainer esto ya está automatizado.
Comparación práctica: Ulysses vs Ring Attention
Método de paralelismo: Ulysses usa particionado por cabezas; Ring Attention usa intercambio en anillo de KV.
Backend: Ulysses requiere DeepSpeed ZeRO; Ring funciona con PyTorch FSDP2.
Soporte de atención: Ulysses soporta FlashAttention 2/3 y SDPA; Ring suele usar SDPA.
Comunicación: Ulysses hace dos all-to-all por capa (mejor uso de bisección de banda), Ring hace comunicación punto a punto serializada.
Restricciones: Ulysses exige num_heads >= sp_size; Ring no tiene esa limitación.
¿Conclusión? Cambiar de uno a otro suele ser tan sencillo como ajustar la config de Accelerate. Prueba ambos en tu hardware.
Buenas prácticas y recomendaciones
Asegura divisibilidad: la longitud global max_length debe ser divisible por sp_size. Usa pad_to_multiple_of = sp_size.
Usa FlashAttention 2 para Ampere y FlashAttention 3 para Hopper. Evita FA2 en Blackwell; espera FA4 cuando esté disponible.
Para modelos grandes combina Ulysses con ZeRO Stage 3 y offload de optimizer/params si hace falta.
Activa PYTORCH_ALLOC_CONF=expandable_segments:True para permitir mayores longitudes.
Empareja sp_size y dp_shard_size para tu recuento de GPUs: por ejemplo, en 4 GPUs puedes usar SP=4/DP=1 para secuencias máximas o SP=2/DP=2 para equilibrio entre longitud y rendimiento.
Si está disponible, habilita use_liger_kernel=True y optimizaciones como FusedLinearCrossEntropy y TiledMLP para ahorrar memoria en logits y MLPs.
Benchmarks esenciales y verificación de pérdida
Experimentos con Qwen3-4B en H100 80GB + ZeRO-3 muestran:
SP=4 en 4 GPUs permite saltar de 8K a 96K tokens en la misma configuración de memoria (pico ~66 GB por GPU a 96K). A 128K hubo OOM en esa configuración.
Throughput: en 64K tokens SP=4 obtuvo ~13.4K tokens/s, ~3.7x sobre baseline de 8K en 1 GPU. A medida que la secuencia crece, la computación cuadrática domina y la paralelización por secuencia es más eficiente.
Equivalencia de pérdida: con el presupuesto de tokens igualado (ajustando GAS = SP cuando corresponde), el entrenamiento con SP y DP coincide en pérdida token-normalizada; las diferencias residuales aparecen en logs reportados pero no en la función objetivo canónica.
Estos resultados implican que Ulysses es una herramienta práctica para entrenar con contextos muy largos sin sacrificar la calidad del entrenamiento, siempre que ajustes correctamente batching y acumulación de gradientes.
Requisitos y versiones recomendadas
DeepSpeed >= 0.18.1
accelerate >= 1.12
transformers >= 5.0 para Trainer integration
trl >= 0.18.0 para SFTTrainer
FlashAttention 2/3 según GPU
Reflexión final
Ulysses Sequence Parallelism convierte un cuello de botella fundamental de la arquitectura transformer en una palanca práctica: repartir cabezas en lugar de duplicar la secuencia. ¿Qué significa esto en términos prácticos? Que ahora es viable entrenar sobre libros, colecciones de documentos y sesiones de razonamiento largas sin necesitar clusters imposibles. Si tu proyecto exige contexto extendido, vale la pena probar Ulysses y compararlo con Ring Attention en tu hardware. La elección correcta depende de tu topología de red, número de GPUs, y la arquitectura del modelo.
