--- title: Reducing High-Bandwidth Memory Bottlenecks in JAX-Based... description: GPU high-bandwidth memory (HBM) is increasingly the bottleneck in LLM training. JAX's host offloading feature moves selected activations to pinned CPU memory... canonical: https://daily.dev/posts/reducing-high-bandwidth-memory-bottlenecks-in-jax-based-llm-training-with-host-offloading-15fftznzt twitter:card: summary_large_image twitter:site: @dailydotdev og:type: website og:site_name: daily.dev og:title: Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading | daily.dev og:description: GPU high-bandwidth memory (HBM) is increasingly the bottleneck in LLM training. JAX's host offloading feature moves selected activations to pinned CPU memory... og:url: https://daily.dev/posts/reducing-high-bandwidth-memory-bottlenecks-in-jax-based-llm-training-with-host-offloading-15fftznzt og:image: https://api.daily.dev/og/posts/15FFtZNzT.png og:image:alt: Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading og:image:width: 1200 og:image:height: 630 og:locale: en --- > ## Documentation Index > Fetch the complete documentation index at: https://daily.dev/llms.txt > Use this file to discover all available pages before exploring further. # Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading **[NVIDIA Developer](https://daily.dev/sources/nvidiadev)** · 9 min read · 0 upvotes · 0 comments ## Summary GPU high-bandwidth memory (HBM) is increasingly the bottleneck in LLM training. JAX's host offloading feature moves selected activations to pinned CPU memory during the forward pass and streams them back during the backward pass, reducing HBM pressure without recomputing activations. On NVIDIA Grace Blackwell systems, the NVLink-C2C interconnect (900 GB/s bidirectional) makes this practical. Using MaxText with 128 GB200 NVL72 GPUs, DeepSeek-V3 671B achieved 908.2 TFLOPs/s/device with host offloading, LHS, and pipelined transfers — 57% faster than activation rematerialization at the same batch size. It also enabled a 4x larger batch configuration (global batch 1024 vs 256) that would otherwise OOM. Llama 3.1 405B saw a smaller 2.9% throughput gain. Key enablers are XLA's Latency Hiding Scheduler (LHS) and pipelined host offloading flags that overlap data transfers with compute. Practical guidance covers JAX APIs, XLA flags, and profiling with Nsight Systems. ## Full article daily.dev links to this article rather than hosting it. Read it at the original source: ## Similar posts on daily.dev - [Leveraging CPU memory for faster, cost-efficient TPU LLM training](https://daily.dev/posts/leveraging-cpu-memory-for-faster-cost-efficient-tpu-llm-training-qc1ngrujv) · Google Open Source Blog · 0 upvotes · 0 comments - [Accelerating Long-Context Model Training in JAX and XLA](https://daily.dev/posts/accelerating-long-context-model-training-in-jax-and-xla-c4zzbjxps) · NVIDIA Developer · 1 upvotes · 0 comments - [Inside vLLM’s New KV Offloading Connector: Smarter Memory Transfer for Maximizing Inference Throughput](https://daily.dev/posts/inside-vllm-s-new-kv-offloading-connector-smarter-memory-transfer-for-maximizing-inference-throughp-6u1hcazvw) · vLLM · 4 upvotes · 0 comments --- Tags: [#mixture-of-experts](https://daily.dev/tags/mixture-of-experts) [View this post on daily.dev](https://daily.dev/posts/reducing-high-bandwidth-memory-bottlenecks-in-jax-based-llm-training-with-host-offloading-15fftznzt) ```json {"@context":"https://schema.org","@graph":[{"@type":"Organization","@id":"https://daily.dev/#organization","name":"daily.dev","url":"https://daily.dev","logo":{"@type":"ImageObject","url":"https://daily.dev/apple-touch-icon.png","width":180,"height":180},"sameAs":["https://twitter.com/dailydotdev","https://github.com/dailydotdev","https://www.linkedin.com/company/daily-dev-ltd"]},{"@type":"WebSite","@id":"https://daily.dev/#website","url":"https://daily.dev","name":"daily.dev","publisher":{"@id":"https://daily.dev/#organization"},"potentialAction":{"@type":"SearchAction","target":{"@type":"EntryPoint","urlTemplate":"https://daily.dev/search?q={search_term_string}"},"query-input":"required name=search_term_string"}}]} {"@context":"https://schema.org","@type":"TechArticle","headline":"Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading","url":"https://daily.dev/posts/reducing-high-bandwidth-memory-bottlenecks-in-jax-based-llm-training-with-host-offloading-15fftznzt","mainEntityOfPage":{"@type":"WebPage","@id":"https://daily.dev/posts/reducing-high-bandwidth-memory-bottlenecks-in-jax-based-llm-training-with-host-offloading-15fftznzt"},"datePublished":"2026-07-10T18:18:42.478Z","dateModified":"2026-07-10T18:19:06.795Z","description":"GPU high-bandwidth memory (HBM) is increasingly the bottleneck in LLM training. JAX's host offloading feature moves selected activations to pinned CPU memory...","image":"https://media.daily.dev/image/upload/f_auto,q_auto/v1/posts/5798c90f276e1858cf36ccd9883ea3c3?_a=AQAEuop","thumbnailUrl":"https://media.daily.dev/image/upload/f_auto,q_auto/v1/posts/5798c90f276e1858cf36ccd9883ea3c3?_a=AQAEuop","isAccessibleForFree":true,"articleSection":"NVIDIA Developer","inLanguage":"en","publisher":{"@type":"Organization","name":"daily.dev","url":"https://daily.dev","logo":{"@type":"ImageObject","url":"https://daily.dev/apple-touch-icon.png","width":180,"height":180}},"author":{"@type":"Organization","name":"NVIDIA Developer","logo":"https://media.daily.dev/image/upload/t_logo,f_auto/v1/logos/86e45aab42ba48ce83103d01b1119910","url":"https://daily.dev/sources/nvidiadev"},"commentCount":0,"discussionUrl":"https://daily.dev/posts/reducing-high-bandwidth-memory-bottlenecks-in-jax-based-llm-training-with-host-offloading-15fftznzt","interactionStatistic":[{"@type":"InteractionCounter","interactionType":{"@type":"LikeAction"},"userInteractionCount":0},{"@type":"InteractionCounter","interactionType":{"@type":"CommentAction"},"userInteractionCount":0}],"keywords":"mixture-of-experts","timeRequired":"PT9M"} {"@context":"https://schema.org","@type":"BreadcrumbList","itemListElement":[{"@type":"ListItem","position":1,"name":"Home","item":"https://daily.dev"},{"@type":"ListItem","position":2,"name":"NVIDIA Developer","item":"https://daily.dev/sources/nvidiadev"},{"@type":"ListItem","position":3,"name":"Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading"}]} ```