← back
Road to 5 Million Tokens: Breaking Barriers in Long Context Training — Max Ryabinin, Together AI
Takeaway
Reusing and offloading activation memory can extend context-parallel training beyond limits reached by standard techniques.
Summary
- Long-context training faces quadratic attention computation and growing activation memory, motivating careful memory accounting.
- Together AI combines sharded training, Ulysses context parallelism, activation checkpointing, CPU offload, and sequence-wise tiling of MLP and loss computations.
- These techniques make multi-million-token training feasible, but attention-head activation buffers remain a constraint.
- The proposed extension processes attention heads in smaller groups and reuses buffers, reducing memory without materially hurting throughput at small scales.
long-contextcontext-parallelismtraining
Original description
Training a standard LLaMA 3B model with a 3 million token context on a single 8xH100 node fails before you even start: the model parameters alone exhaust GPU memory. Max Ryabinin from Together AI walks through the full stack of techniques needed to get there: fully sharded data parallelism, DeepSpeed Ulysses context parallelism for an 8x activation reduction, activation checkpointing for another 8x, CPU offloading for transformer block inputs, and chunked sequence training to avoid allocating buffers 3 million tokens wide. Even that stack falls short at 5 million tokens. The novel contribution, Untied Ulysses, goes deeper into the context parallelism step: instead of allocating one large buffer per attention head group, it chunks the heads further and reuses those buffers across iterations, cutting activation memory with negligible throughput impact. At both 8B and 32B scale the results match the most memory optimized transformer training baselines while pushing sequence length 25% further than prior Ulysses implementations. Speaker info: / max-ryabinin https://x.com/m_ryabinin