Scaling Transformer Training to 5 Million Tokens

AI Engineer2026-06-08go watch the original →

the gist

To train models with 5 million token context lengths on limited hardware, Together AI combines standard memory-optimization techniques with a novel method called Untied Ulysses that chunks and reuses attention head buffers.

Memory Optimization Stack

To achieve extreme context lengths, the training stack must address linear memory growth and quadratic computation costs. The following techniques are applied sequentially to fit large models onto an 8xH100 node:

Fully Sharded Data Parallelism (FSDP): Distributes model parameters across all available GPUs to resolve initial parameter-loading memory errors.
DeepSpeed Ulysses: Partitions attention heads across GPUs, allowing each device to compute attention over the entire sequence length while reducing activation memory by approximately 8x.
Activation Checkpointing: Recomputes activations during the backward pass instead of storing them, providing an additional 8x reduction in memory usage.
CPU Offloading: Moves transformer block inputs to CPU memory when not in use, prefetching them only when required for backpropagation.
Chunked Sequence Training: Tiles element-wise operations like loss functions and MLPs across the sequence length to prevent the allocation of massive buffers proportional to the 3 million token sequence.

The Untied Ulysses Technique

Even with the standard stack, 5 million tokens exceed available memory. The Untied Ulysses method optimizes the context parallelism step by further subdividing attention heads into smaller chunks. Instead of allocating a buffer for all heads simultaneously, the system iterates through these chunks and reuses the same memory buffers across iterations. This approach reduces activation memory with negligible impact on throughput, allowing for 25% longer sequence lengths compared to standard Ulysses implementations at both 8B and 32B model scales.

Context

Training long-context models is essential for agentic workflows and temporal consistency in tasks like video generation. Standard transformer architectures struggle with memory bottlenecks as sequence lengths grow. By stacking these optimizations, researchers can push context limits significantly further than vanilla implementations, matching the performance of memory-optimized baselines while enabling training at the 5 million token scale.