Speed-up JAX LLM Training on Intel® Xeon® 6 CPU: Activation Offloading on Heterogeneous Systems

JAX-based Activation Offloading on Intel® Xeon® 6 with P-cores systems offers an effective alternative to activation recomputation, repurposing the CPU’s large DDR5 host memory as a live activation store and leveraging XLA’s asynchronous compute–communication overlap to avoid throughput loss.