# Intel Xeon

Published articles for Intel Xeon.

This is one page of public article previews, not the complete archive. Follow Next page to continue. Summaries are not the original full articles.

## Leveraging CPU memory for faster, cost-efficient TPU LLM training

DevFeed: [Leveraging CPU memory for faster, cost-efficient TPU LLM training](<https://devfeed.tech/articles/leveraging-cpu-memory-for-faster-cost-efficient-tpu-llm-training-34299.md>)

Original publisher: [Read original article](<http://opensource.googleblog.com/2026/04/leveraging-cpu-memory-for-faster-cost-efficient-tpu-llm-training.html>)

Author: Google Open Source (noreply@blogger.com)

Published: 2026-04-10T18:30:00Z

Content type: tutorial

Language: en

Sources: [Google Open Source Blog](<https://devfeed.tech/sources/google-open-source-blog.md>)

Topics: [cpu](<https://devfeed.tech/topics/cpu.md>), [Large Language Model](<https://devfeed.tech/topics/llm.md>), [intel](<https://devfeed.tech/topics/intel.md>), [Google Cloud Platform (GCP)](<https://devfeed.tech/topics/google-cloud.md>), [Cloud](<https://devfeed.tech/topics/cloud.md>), [implementation](<https://devfeed.tech/topics/implementation.md>)

Tags: [accelerators](<https://devfeed.tech/tags/accelerators.md>), [cloud-tpu](<https://devfeed.tech/tags/cloud-tpu.md>), [cost](<https://devfeed.tech/tags/cost.md>), [cpu](<https://devfeed.tech/tags/cpu.md>), [host-offloading](<https://devfeed.tech/tags/host-offloading.md>), [intel-xeon](<https://devfeed.tech/tags/intel-xeon.md>), [jax](<https://devfeed.tech/tags/jax.md>), [llm-training](<https://devfeed.tech/tags/llm-training.md>), [memory](<https://devfeed.tech/tags/memory.md>), [performance](<https://devfeed.tech/tags/performance.md>), [tpu](<https://devfeed.tech/tags/tpu.md>)

### AI overview

This practical guide explains how to use host activation offloading with JAX on TPU platforms. It describes moving selected activations from TPU device memory to Intel Xeon CPU memory to reduce accelerator memory pressure and support larger models or batch sizes, while discussing potential throughput and cost benefits.

### Source excerpt

by Keyur Ruganathbhai Ranipa, Qinglan Xiang, Vrushabh Sanghavi, Ramesh AG & Weilin Wang, Intel and Penporn Koanantakool, Google Host offloading with JAX on Intel® Xeon® processors As Large Language Models (LLMs) continue to scale into the hundreds of billions of parameters, device memory capacity has become a big limiting factor in training, as intermediate activations from every layer in the forward pass are needed in the backward pass. To reduce device memory pressure, these activations can be rematerialized during the backward pass, trading memory for recomputation. While rematerialization enables larger models to fit within limited device memory, it significantly increases training time and cost. Intel® Xeon® processors (5th and 6th Gen) with Advanced Matrix Extensions (AMX) enable practical host offloading of selected memory- and compute-intensive components in JAX training workflows. This approach can help teams train larger models, relieve accelerator memory pressure, improve end-to-end throughput, and reduce total cost of ownership--particularly on TPU-based Google Cloud instances. By publishing these results and implementation details, Google and Intel aim to promote transparency and share practical guidance with the community. This post describes how to enable activation offloading for JAX on TPU platforms and outlines considerations for building scalable, cost-aware hybrid CPU-accelerator training workflows. Figure 1. Google Cloud TPU Pod commonly used in LLM training. Host offloading Traditional LLM training is usually done on device accelerators alone. However, modern host machines have much larger memory size than accelerators (512GB or more) and can offer extra compute power, e.g., TFLOPS in case of Intel® Xeon® Scalable Processor with AMX capability. Leveraging host resources can be a great alternative to rematerialization. Host offloading selectively moves computation or data between host and device to optimize performance and memory usage. Host memo