{"id":112140,"date":"2026-02-03T09:30:00","date_gmt":"2026-02-03T17:30:00","guid":{"rendered":"https:\/\/developer.nvidia.com\/blog\/?p=112140"},"modified":"2026-03-05T11:20:03","modified_gmt":"2026-03-05T19:20:03","slug":"accelerating-long-context-model-training-in-jax-and-xla","status":"publish","type":"post","link":"https:\/\/developer.nvidia.com\/blog\/accelerating-long-context-model-training-in-jax-and-xla\/","title":{"rendered":"Accelerating Long-Context Model Training in JAX and XLA"},"content":{"rendered":"\n

Large language models (LLMs)<\/a> are rapidly expanding their context windows, with recent models supporting sequences of 128K tokens, 256K tokens, and beyond. However, training these models with extended context lengths presents significant computational and communication challenges. As context lengths grow, the memory and communication overhead of attention mechanisms scale quadratically, creating bottlenecks that traditional parallelism strategies struggle to address efficiently.<\/p>\n\n\n\n

This post demonstrates that integrating the NVSHMEM<\/a> communication library into Accelerated Linear Algebra (XLA) compiler optimizes context parallelism. This integration enables the efficient training of Llama 3 8B model in JAX framework with sequences up to 256K tokens. Our results show that NVSHMEM provides up to 36% speedup over NVIDIA Collective Communications Library (NCCL)<\/a> for long-context training workloads, particularly when combined with tensor parallelism across multiple nodes.<\/p>\n\n\n\n

The long-context training challenge<\/i><\/a><\/h2>\n\n\n\n

To understand why NVSHMEM provides significant speedups for long-context training, it\u2019s necessary to first understand how context parallelism works and the unique communication patterns it creates. This section explains why the fine-grained, latency-sensitive communication of ring attention makes it an ideal candidate for optimization.<\/p>\n\n\n\n

Context parallelism and ring attention<\/i><\/a><\/h3>\n\n\n\n

Context parallelism <\/em>(CP) is a parallelization strategy designed specifically for handling long sequences in transformer models. Unlike data parallelism, which splits the batch, or tensor parallelism, which splits the model, context parallelism splits the sequence dimension across multiple devices.<\/p>\n\n\n\n

Ring attention<\/em> is a memory-efficient implementation of context parallelism that uses a ring-based communication pattern. During attention computation, each device:<\/p>\n\n\n\n