{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "bQbS50fIdHw1" }, "source": [ "(host-offloading)=\n", "# JAX Memories and Host Offloading\n", "\n", "\n", "\n", "This tutorial provides a practical introduction to host offloading techniques in JAX, focusing on:\n", "\n", "- Activation offloading\n", "- Parameter offloading\n", "\n", "By applying offloading strategies, you can better manage memory resources and reduce memory pressure on your devices. To implement these strategies effectively, you'll need to understand JAX's core mechanisms for data placement and movement.\n", "\n", "## Building Blocks for Offloading\n", "\n", "JAX provides several key components for controlling where and how data are stored and moved between the host and the device memory. In the following sections, you'll explore:\n", "\n", "- How to specify data distribution with sharding\n", "- How to control memory placement between host and device\n", "- How to manage data movement in jitted functions\n", "\n", "### NamedSharding and Memory Kinds\n", "\n", "{class}`~jax.sharding.NamedSharding` defines how data are distributed across devices. It includes:\n", "\n", "- Basic data distribution configuration\n", "- `memory_kind` parameter for specifying memory type (`device` or `pinned_host`)\n", "- By default, `memory_kind` is set to `device` memory\n", "- `with_memory_kind` method for creating new sharding with modified memory type" ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "f-6sxUlqrlBn", "outputId": "691a3df2-8341-44a9-a4a0-5521c2d891e3" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "NamedSharding(mesh=Mesh('x': 1, 'y': 1), spec=PartitionSpec('x', 'y'), memory_kind=device)\n", "NamedSharding(mesh=Mesh('x': 1, 'y': 1), spec=PartitionSpec('x', 'y'), memory_kind=pinned_host)\n" ] } ], "source": [ "import jax\n", "import jax.numpy as jnp\n", "from jax.sharding import Mesh, NamedSharding, PartitionSpec as P\n", "import numpy as np\n", "\n", "# Create mesh\n", "# 1x1 mesh represents a single device with two named dimensions (x and y)\n", "mesh = Mesh(np.array(jax.devices()[0]).reshape(1, 1), ('x', 'y'))\n", "\n", "# Device sharding - partitions data along x and y dimensions\n", "s_dev = NamedSharding(mesh, P('x', 'y'), memory_kind=\"device\")\n", "\n", "# Host sharding - same partitioning but in pinned host memory\n", "s_host = s_dev.with_memory_kind('pinned_host')\n", "\n", "print(s_dev) # Shows device memory sharding\n", "print(s_host) # Shows pinned host memory sharding" ] }, { "cell_type": "markdown", "metadata": { "id": "R_pB9465VoMP" }, "source": [ "### Data Placement with device_put\n", "\n", "{func}`jax.device_put` is a function that explicitly transfers arrays to a specified memory location according to a sharding specification." ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "OJFnf7FGp6Lj", "outputId": "c762e1df-2453-4ed9-9d53-0defb6a05ce2" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "pinned_host\n", "device\n" ] } ], "source": [ "# Create a 2x4 array\n", "arr = jnp.arange(8.0).reshape(2, 4)\n", "\n", "# Move arrays to different memory locations based on sharding objects\n", "arr_host = jax.device_put(arr, s_host) # Places in pinned host memory\n", "arr_dev = jax.device_put(arr, s_dev) # Places in device memory\n", "\n", "# Verify memory locations\n", "print(arr_host.sharding.memory_kind) # Output: pinned_host\n", "print(arr_dev.sharding.memory_kind) # Output: device" ] }, { "cell_type": "markdown", "metadata": { "id": "HHXvBpQKTMCR" }, "source": [ "### Output Sharding Controls\n", "\n", "Shardings determine how data is split across devices. JAX provides `out_shardings` to control how output arrays are partitioned when leaving a jitted function.\n", "\n", "Key Features:\n", " - Can differ from input sharding\n", " - Allows different memory kinds for outputs\n", "\n", "Examples:\n", "\n", "#### Device Output Sharding" ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "ZXNj9NUeaIdX", "outputId": "399321ef-082a-4a77-c33a-9de3421f429b" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Result value of H2D: \n", " [[0. 1. 2. 3.]\n", " [4. 5. 6. 7.]]\n" ] } ], "source": [ "f = jax.jit(lambda x:x, out_shardings=s_dev)\n", "out_dev = f(arr_host)\n", "print(\"Result value of H2D: \\n\", out_dev)" ] }, { "cell_type": "markdown", "metadata": { "id": "iYXC5ix384XP" }, "source": [ "Moving data from host to device memory when needed for computation is the essence of host offloading. Use {func}`jax.device_put` to perform this transfer in this example to optimize performance." ] }, { "cell_type": "code", "execution_count": 7, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "cmM6tJTS84XQ", "outputId": "40c353a1-fb55-44bc-bac9-dffc09852f49" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Result value of H2D and add 1 in device memory: \n", " [[1. 2. 3. 4.]\n", " [5. 6. 7. 8.]]\n" ] } ], "source": [ "# Instead of the lambda function, you can define add_func to explicitly\n", "# move data to device before computation\n", "def add_func(x): # Move data to device and add one\n", " x = jax.device_put(x, s_dev)\n", " return x + 1\n", "\n", "f = jax.jit(add_func, out_shardings=s_dev)\n", "out_dev = f(arr_host)\n", "print(\"Result value of H2D and add 1 in device memory: \\n\", out_dev)" ] }, { "cell_type": "markdown", "metadata": { "id": "EbE-eBrJTBuS" }, "source": [ "#### Host Output Sharding" ] }, { "cell_type": "code", "execution_count": 4, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "FjZzkxI8ky4r", "outputId": "2a1b6e7a-1c29-4347-c020-7b47c27a5cc3" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Result value of D2H: \n", " [[0. 1. 2. 3.]\n", " [4. 5. 6. 7.]]\n" ] } ], "source": [ "f = jax.jit(lambda x: x, out_shardings=s_dev)\n", "out_host = f(arr_host) # Input arrays in the device memory while output arrays in the host memory\n", "print(\"Result value of D2H: \\n\", out_host)" ] }, { "cell_type": "markdown", "metadata": { "id": "UhLVvRO2p6Lj" }, "source": [ "## Activation Offloading\n", "\n", "The detailed coverage of activation offloading can be found in the {ref}`gradient-checkpointing` tutorial. Activation offloading helps manage memory by moving intermediate activations to host memory after the forward pass, and bringing them back to device memory during the backward pass when needed for gradient computation.\n", "\n", "To implement activation offloading effectively, you need to understand checkpoint names and policies. Here's how they work in a simple example:\n", "\n", "### Checkpoint Names\n", "\n", "The {func}`checkpoint_name` function allows you to label activations for memory management during computation. Here's a simple example:" ] }, { "cell_type": "code", "execution_count": 5, "metadata": { "id": "sLO9ceS6p6Lj" }, "outputs": [], "source": [ "from jax.ad_checkpoint import checkpoint_name\n", "\n", "def layer(x, w):\n", " w1, w2 = w\n", " x = checkpoint_name(x, \"x\")\n", " y = x @ w1\n", " return y @ w2, None" ] }, { "cell_type": "markdown", "metadata": { "id": "-_T92oCOp6Lk" }, "source": [ "This example shows:\n", "\n", "* A simple neural network layer with two matrix multiplications\n", "* Labeling of input activation x with identifier `\"x\"`\n", "* Sequential operations:\n", " 1. First multiplication: `x @ w1`\n", " 2. Second multiplication: `y @ w2`\n", "\n", "The checkpoint name helps the system decide whether to:\n", "* Keep the activation in device memory or\n", "* Offload it to host memory during computation\n", "\n", "This pattern is common in neural networks, where multiple transformations are applied sequentially to input data.\n", "\n", "\n", "### Checkpoint Policies\n", "\n", "The {func}`jax.remat` transformation manages memory by handling intermediate values through three strategies:\n", "\n", "1. Recomputing during backward pass (default behavior)\n", "2. Storing on device\n", "3. Offloading to host memory after forward pass and loading back during backward pass\n", "\n", "Example of setting an offloading checkpoint policy:" ] }, { "cell_type": "code", "execution_count": 6, "metadata": { "id": "W8Usw_wOp6Lk" }, "outputs": [], "source": [ "from jax import checkpoint_policies as cp\n", "\n", "policy = cp.save_and_offload_only_these_names(\n", " names_which_can_be_saved=[], # No values stored on device\n", " names_which_can_be_offloaded=[\"x\"], # Offload activations labeled \"x\"\n", " offload_src=\"device\", # Move from device memory\n", " offload_dst=\"pinned_host\" # To pinned host memory\n", ")" ] }, { "cell_type": "markdown", "metadata": { "id": "iuDRCXu7ky4r" }, "source": [ "Since {func}`jax.lax.scan` is commonly used in JAX for handling sequential operations (like RNNs or transformers), you need to know how to apply your offloading strategy in this context.\n", "\n", "Key components:\n", "* {func}`jax.remat` applies our checkpoint policy to the layer function\n", "* `prevent_cse=False` enables XLA's common subexpression elimination for better performance\n", "* {func}`jax.lax.scan` iterates the rematerialized layer along an axis" ] }, { "cell_type": "code", "execution_count": 7, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "xCrxjTx_p6Lk", "outputId": "13d46584-9b25-4622-b3c3-f50c1dac02c2" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Sample of results: [3.7363498e-07 3.7363498e-07 3.7363498e-07 3.7363498e-07 3.7363498e-07]\n" ] } ], "source": [ "def scanned(w, x):\n", " remat_layer = jax.remat(layer,\n", " policy=policy, # Use our offloading policy\n", " prevent_cse=False) # Allow CSE optimizations\n", " result = jax.lax.scan(remat_layer, x, w)[0]\n", " return jnp.sum(result)\n", "\n", "# Initialize input and weights with small values (0.0001)\n", "input = jnp.ones((256, 256), dtype=jnp.float32) * 0.001 # Input matrix: 256 x 256\n", "w1 = jnp.ones((10, 256, 1024), dtype=jnp.float32) * 0.001 # 10 layers of 256 x 1024 matrices\n", "w2 = jnp.ones((10, 1024, 256), dtype=jnp.float32) * 0.001 # 10 layers of 1024 x 256 matrices\n", "\n", "# Compile and compute gradients of the scanned function\n", "f = jax.jit(jax.grad(scanned)) # Apply JIT compilation to gradient computation\n", "result_activation = f((w1, w2), input) # Execute the function with weights and input\n", "print(\"Sample of results: \", result_activation[0][0, 0, :5])" ] }, { "cell_type": "markdown", "metadata": { "id": "0tx7aara42pY" }, "source": [ "### Summary of Activation Offloading\n", "\n", "Activation offloading provides a powerful way to manage memory in large computations by:\n", "\n", "* Using checkpoint names to mark specific activations\n", "* Applying policies to control where and how activations are stored\n", "* Supporting common JAX patterns like scan operations\n", "* Moving selected activations to host memory when device memory is under budget\n", "\n", "This approach is particularly useful when working with large models that would otherwise exceed device memory capacity.\n", "\n", "## Parameter Offloading\n", "\n", "Model parameters (also known as weights) can be offloaded to the host memory to optimize device memory usage during initialization. This is achieved by using {func}`jax.jit` with a sharding strategy that specifies host memory kind.\n", "\n", "While parameter offloading and activation offloading are distinct memory optimization techniques, the following example demonstrates parameter offloading built upon the activation offloading implementation shown earlier.\n", "\n", "### Parameter Placement for Computation\n", "\n", "Different from the earlier `layer` function, {func}`jax.device_put` is applied to move parameter `w1` and `w2` to the device before the matrix multiplications. This ensures the parameters are available on the device for both forward and backward passes.\n", "\n", "Note that the activation offloading implementation remains unchanged, using the same:\n", "* Checkpoint name `\"x\"`\n", "* Checkpoint policy\n", "* `scanned` function combining {func}`jax.remat` and {func}`jax.lax.scan`\n", "\n", "### Parameter Initialization with Host Offloading\n", "\n", "During the initialization, parameter `w1` and `w2` are placed on host memory before being passed to the {func}`jax.jit` function `f`, while keeping the `input` variable on the device." ] }, { "cell_type": "code", "execution_count": 8, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "1qGN2hBQdheo", "outputId": "48c09658-f8b6-4be3-ef0e-02e0e2566e10" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Results match within tolerance: True\n" ] } ], "source": [ "# Hybrid version: Both activation and parameter offloading\n", "def hybrid_layer(x, w):\n", " # Move model parameters w1 and w2 to host memory via device_put\n", " w1, w2 = jax.tree.map(lambda x: jax.device_put(x, s_dev), w)\n", " x = checkpoint_name(x, \"x\") # Offload activation x to host memory\n", " y = x @ w1\n", " return y @ w2, None\n", "\n", "def hybrid_scanned(w, x):\n", " remat_layer = jax.remat(hybrid_layer, # Use hybrid_layer instead of layer\n", " policy=policy, # Use offloading policy\n", " prevent_cse=False) # Allow CSE optimizations\n", " result = jax.lax.scan(remat_layer, x, w)[0]\n", " return jnp.sum(result)\n", "\n", "# Move model parameters w1 and w2 to the host via device_put\n", "# Initialize input and weights with small values (0.0001)\n", "wh1 = jax.device_put(w1, s_host)\n", "wh2 = jax.device_put(w2, s_host)\n", "\n", "# Compile and compute gradients of the scanned function\n", "f = jax.jit(jax.grad(hybrid_scanned)) # Apply JIT compilation to gradient computation\n", "result_both = f((wh1, wh2), input) # Execute with both activation and parameter offloading\n", "\n", "# Verify numerical correctness\n", "are_close = jnp.allclose(\n", " result_activation[0], # Result from activation offloading only\n", " result_both[0], # Result from both activation and parameter offloading\n", " rtol=1e-5,\n", " atol=1e-5\n", ")\n", "print(f\"Results match within tolerance: {are_close}\")" ] }, { "cell_type": "markdown", "metadata": { "id": "SVpozzwHflQk" }, "source": [ "The matching results verify that initializing parameters on host memory maintains computational correctness.\n", "\n", "### Limitation of Parameter Offloading\n", "\n", "{func}`jax.lax.scan` is crucial for effective parameter management. Using an explicit for loop would cause parameters to continuously occupy device memory, resulting in the same memory usage as without parameter offloading. While {func}`jax.lax.scan` allows specifying the scan axis, parameter offloading currently works only when scanning over axis 0. Scanning over other axes generates a `transpose` operation during compilation before returning parameters to the device, which is expensive and not supported on all platforms.\n", "\n", "## Tools for Host Offloading\n", "\n", "For device memory analysis, refer to :doc:`device_memory_profiling`. The profiling tools described in {ref}`profiling` can help measure memory savings and performance impact from host offloading." ] } ], "metadata": { "accelerator": "TPU", "colab": { "gpuType": "V28", "provenance": [], "toc_visible": true }, "jupytext": { "formats": "ipynb,md:myst" }, "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.3" } }, "nbformat": 4, "nbformat_minor": 0 }