JAX Pathways
Cloud TPU Mini-Course · Step 3 of 3Cloud TPU VM + CPU verifierLesson 1 of 2 in Phase 21

Phase 21 / TPU Systems: Generations, Memory, Precision & XProf

TPU generations and HBM memory budgeting

Revision v0.7.0 ·

Lesson · about 25 minutes · CPU practice

Compare TPU v5e, v5p, and v6e specifications and roofline ridge points, budget FP32 vs mixed-BF16 training state and KV-cache HBM, and size the minimum chip count for a model.

Phase study guide: preparation, milestones and project evidence

Before you start

Before this: TPU Workflows: Launching Jobs & Checkpointed Experiments
Hardware: CPU capacity & roofline planner; Cloud TPU VM for target HBM verification

Download notebook
Other downloadsPython worked examplesLesson text
Setup help: Python, notebook, or Cloud TPU VM

Use the tested CPU packages (the recorded environment uses Python 3.14.3). Other Python versions and accelerator environments require separate checks. The optional Cloud TPU specialization explains the target-specific execution boundaries. Logical CPU devices practice partitioning; they do not simulate TPU performance.

Download the repository ZIP or use your existing checkout. Extract it and open a terminal in the folder containing requirements-cpu.txt. The beginner ZIP is separate; these lessons use the full repository folder.

macOS / Linux terminal · repository folder
python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-cpu.txt

PowerShell · repository folder
py -3 -m venv .venv
.\.venv\Scripts\python.exe -m pip install -r requirements-cpu.txt

After activation, run the complete worked example from that same folder:

macOS / Linux terminal · activated environment
python phases/21-tpu-systems/01-tpu-generations-and-hbm-memory-budgeting/code/main.py

PowerShell command · repository folder
.\.venv\Scripts\python.exe phases/21-tpu-systems/01-tpu-generations-and-hbm-memory-budgeting/code/main.py

For an incremental build, create main.py in your own practice folder, append each build block in order, and run python main.py with this environment active. Keep that file separate from the downloaded reference solution.

Check: Prints the ridge-point FLOP/byte ratios for v5e (240.5), v5p (166.0), and v6e (573.8) along with the 7B and 1.5B memory budgets in GiB.

Revisit the guided CPU setup · Open 19: TPU Setup. Package instructions are CPU instructions; Windows and accelerator environments need their own validation.

Phase 21.01 Stage 1 of 5 · Understand
1 Understand 2 Build & Run 3 Experiment 4 Exercise & Check 5 Evidence
Stage 1 of 5 · Understand the MechanismCore mental model, equations & visual walkthrough

What you will be able to do

  • Compare Google Cloud TPU generations (v5e, v5p, v6e Trillium, and Ironwood) by HBM capacity, HBM bandwidth, and MXU BF16 throughput.
  • Calculate the Roofline ridge point (bf16_tflops / bw_gbs) in FLOPs per byte for each TPU generation.
  • Compute training-state and KV-cache HBM from explicit tensor byte counts, including master weights, compute weights, gradients, and Adam moments.
  • Select the minimum single-host or Pod TPU slice that fits a given model's training state with activation headroom.

The problem

Provisioning and launching a TPU VM is only half the engineering job: if you pick the wrong TPU generation or miscalculate parameter + optimizer + KV-cache HBM bytes, your TPU experiment will either fail with RESOURCE_EXHAUSTED: Out of memory or under-utilize the Matrix Multiply Unit (MXU).

The idea

Budget capacity with a byte ledger before selecting a TPU slice, and treat the roofline ridge point as a theoretical screening estimate. GiB capacity and FLOP/byte ridge points are different quantities and must be read in separate panels. The worked training-state estimates state their tensor dtypes and exclusions; they are not observed runtime memory.

Connect MXU arithmetic intensity and HBM capacity to TPU slice sizing

The ridge point is peak BF16 FLOP/s divided by HBM bandwidth. It is a theoretical balance point for a roofline model, not a measured bottleneck; use a target workload profile to investigate utilization. HBM capacity is a separate quantity.

For this lesson, the full-FP32 training-state ledger is 4 bytes/parameter for weights + 4 for gradients + 4 each for Adam first and second moments = 16 bytes/parameter. The mixed-BF16 ledger assumes an FP32 master weight (4), a BF16 compute copy (2), an FP32 gradient buffer (4), and two FP32 moments (8), totaling 18 bytes/parameter. This deliberately conservative case includes the compute copy and does not include activations, communication buffers, allocator overhead, or temporary copies. If gradients are kept in BF16, subtract 2 bytes/parameter and state that assumption.

TPU HBM-to-MXU dataflow and mixed-precision state budget

Predict: Which tensors stay in float32 in HBM and which operands stream as bfloat16 into the MXU?

Read top to bottom: HBM holds fp32 master weights and Adam moments plus bf16/int8 caches, streams bf16 operands into the MXU, accumulates dot products in fp32, and updates sharded HBM state.
Architecture and dataflow mechanism diagram.

Read top to bottom: HBM stores master fp32 parameters and Adam moments (m, v) alongside bf16 activations and KV cache; operands are cast to bf16 (or int8) to stream across the HBM bandwidth bottleneck into the MXU, which accumulates dot products in fp32 before writing updated state back.

Open full-size diagram
Throughputtok/s=Bglobal×LseqTstep\text{Throughput}_{\text{tok/s}} = \frac{B_{\text{global}} \times L_{\text{seq}}}{T_{\text{step}}}

1. Compare TPU generations (v5e, v5p, v6e, Ironwood) and calculate HBM budgets

Published per-chip TPU specifications determine which slice fits your model:

  • TPU v5e (v5litepod): 1616 GiB HBM per chip, 819819 GB/s HBM bandwidth, 197197 BF16 TFLOP/s (240.5 FLOP/byte ridge point). Single-host slices: 1, 4 (64 GiB), or 8 (128 GiB) chips.
  • TPU v5p: 9595 GiB HBM per chip, 27652765 GB/s HBM bandwidth, 459459 BF16 TFLOP/s (166.0 FLOP/byte ridge point).
  • TPU v6e (Trillium): 3232 GiB HBM per chip, 16001600 GB/s HBM bandwidth, 918918 BF16 TFLOP/s (573.8 FLOP/byte ridge point). Single-host slices: 1, 4 (128 GiB), or 8 (256 GiB) chips.
  • TPU Ironwood (7th-gen scale): 192192 GiB HBM per chip, 73707370 GB/s HBM bandwidth.

Run the TPU generation & HBM budgeting verifier on your TPU VM or workstation

bash
# Run the TPU generation & HBM budgeting verifier on your TPU VM or workstation
gcloud compute tpus tpu-vm scp public/exercises/tpu-03.py "$TPU_NAME":~/jax-tpu-lab/ --zone="$ZONE"
gcloud compute tpus tpu-vm ssh "$TPU_NAME" --zone="$ZONE" \
  --command="JAX_PLATFORMS=tpu ~/jax-tpu-lab/.venv/bin/python ~/jax-tpu-lab/tpu-03.py"

Expected:

Prints the TPU ridge points (FLOP/byte) and the 7B and 1.5B training + KV-cache HBM budgets in GiB.

Calculated TPU memory capacity and theoretical ridge point

Predict before looking: How do per-chip HBM capacity (GiB) and the compute-to-memory ridge point (FLOP/byte) change across TPU v5e, v5p, and v6e?

Two-panel chart: HBM capacity per chip in GiB and theoretical BF16 ridge point in FLOP/byte, shown on separate axes for v5e, v5p, and v6e.
Calculated scenario; no device execution was performed

Read the figure

Two panels show per-chip HBM capacity in GiB and theoretical BF16 ridge point in FLOP/byte for the same TPU generations. Read the panels independently; their heights are not comparable.

Connect it to the computation

These bars are calculations from published specifications and stated tensor byte assumptions. They are not device measurements or a workload profile.

Open full-size figureDownload data and provenance

This figure is a calculated expectation. It does not report device execution.

Stage 2 of 5 · Build & Run the ReferenceStep-by-step code construction & verified execution output

Where to put & run this code: Work in phases/21-tpu-systems/01-tpu-generations-and-hbm-memory-budgeting/code/main.py (or downloaded exercises/tpu-03.py / exercises/tpu-03.ipynb) and execute from the repository root with python3 phases/21-tpu-systems/01-tpu-generations-and-hbm-memory-budgeting/code/main.py.

1 Define TPU generation specs and compute Roofline ridge points

Create main.py with TPU_SPECS and calculate ridge_flops_per_byte for v5e, v5p, and v6e.

main.py · append build step 1
# Step 1 — Define TPU generation specs and compute Roofline ridge points: Dividing peak BF16 FLOP/s by HBM bandwidth (B/s) gives the minimum...
# Import jax for this computation.
import jax
import jax.numpy as jnp
import numpy as np

# Initialize list `TPU_SPECS` for the stage values.
TPU_SPECS = [
    {"gen": "v5e", "hbm_gib": 16.0, "bw_gbs": 819.0, "bf16_tflops": 197.0},
    {"gen": "v5p", "hbm_gib": 95.0, "bw_gbs": 2765.0, "bf16_tflops": 459.0},
    {"gen": "v6e", "hbm_gib": 32.0, "bw_gbs": 1600.0, "bf16_tflops": 918.0},
]
# Iterate over `s` to step through the computation:
for s in TPU_SPECS:
    # Run `round` to compute `s['ridge_flops_per_byte']`.
    s["ridge_flops_per_byte"] = round((s["bf16_tflops"] * 1e12) / (s["bw_gbs"] * 1e9), 1)

Code building blocks in this step

  • bandwidth(...) — Call bandwidth with your updated parameters or inputs from this lesson's workspace.
  • round(...) — Call round with your updated parameters or inputs from this lesson's workspace.

Dividing peak BF16 FLOP/s by HBM bandwidth (B/s) gives the minimum FLOPs per byte needed to saturate the MXU.

2 Estimate training-state and KV-cache HBM budgets across precisions

Append estimate_memory_gib and compare 7B vs 1.5B model memory footprints.

main.py · append build step 2
# Step 2 — Estimate training-state and KV-cache HBM budgets across precisions: Computing both training state and KV-cache GiB upfront prevents...
def estimate_memory_gib(params_billions, kv_tokens, layers=32, kv_heads=8, head_dim=128):
    # Compute `n_params` from `params_billions * 1e9`
    n_params = params_billions * 1e9
    # Compute `kv_elements` from `2.0 * layers * kv_tokens * kv_heads * head_dim`
    kv_elements = 2.0 * layers * kv_tokens * kv_heads * head_dim
    # Evaluate `1024 ** 3` and convert the result into Python scalar/collection `gib`.
    gib = float(1024 ** 3)
    # Return `{'fp32_train_state_gib': round(n_params * 16.0 / gib, 3), 'mixed_bf16_train_state_gib': round(n_params * 18.0 / gib, 3), 'kv_fp32_gib': round(kv_elements * 4.0 / gib, 3), 'kv_bf16_gib': round(kv_elements * 2.0 / gib, 3), 'kv_int8_gib': round(kv_elements * 1.0 / gib, 3)}` to the caller.
    return {
        "fp32_train_state_gib": round((n_params * 16.0) / gib, 3),
        "mixed_bf16_train_state_gib": round((n_params * 18.0) / gib, 3),
        "kv_fp32_gib": round((kv_elements * 4.0) / gib, 3),
        "kv_bf16_gib": round((kv_elements * 2.0) / gib, 3),
        "kv_int8_gib": round((kv_elements * 1.0) / gib, 3),
    }



# Run `estimate_memory_gib` to compute `budget_7b`.
budget_7b = estimate_memory_gib(params_billions=7.0, kv_tokens=8192)
# Run `estimate_memory_gib` to compute `budget_1p5b`.
budget_1p5b = estimate_memory_gib(params_billions=1.5, kv_tokens=4096)
# Assert that `budget_7b["mixed_bf16_train_state_gib"] > 64.0 and budget_7b["mixed_bf16_train_state_gib"] <`.
assert budget_7b["mixed_bf16_train_state_gib"] > 64.0 and budget_7b["mixed_bf16_train_state_gib"] < 128.0
# Assert invariant `budget_1p5b["mixed_bf16_train_state_gib"] < 64.0` holds
assert budget_1p5b["mixed_bf16_train_state_gib"] < 64.0
# Print the observed values to compare against the expected result.
print("TPU ridge points (FLOP/byte):", {s["gen"]: s["ridge_flops_per_byte"] for s in TPU_SPECS})
# Print diagnostic summary of the computed outputs.
print("7B memory budget (GiB):", budget_7b)
# Print diagnostic summary of the computed outputs.
print("1.5B memory budget (GiB):", budget_1p5b)

Code building blocks in this step

  • estimate_memory_gib(...) — Call estimate_memory_gib with your updated parameters or inputs from this lesson's workspace.
  • round(...) — Call round with your updated parameters or inputs from this lesson's workspace.
  • assert condition — Verify that the observed output shape, status, or numerical value satisfies the contract.

Computing the named-tensor ledger identifies whether a topology has enough total HBM, but usable capacity must also reserve room for activations, communication buffers, temporary copies and allocator overhead. The 7B estimate exceeds the 80%-usable budget of a 4-chip v6e slice; it is an estimate, not a fit guarantee.

3 Step 3: Verify invariants on the completed state

Run the final shape and numerical assertions to confirm the state built in Steps 1 and 2.

main.py · append build step 3
assert budget_7b["mixed_bf16_train_state_gib"] > 64.0 and budget_7b["mixed_bf16_train_state_gib"] < 128.0
assert budget_1p5b["mixed_bf16_train_state_gib"] < 64.0

Code building blocks in this step

  • assert condition — Verify that the observed output shape, status, or numerical value satisfies the contract.

Checking these invariants confirms the computation is ready for the full worked experiment.

View the cumulative build in main.py

These blocks append in order to the same file.

main.py · cumulative Python build
# Step 1 — Define TPU generation specs and compute Roofline ridge points: Dividing peak BF16 FLOP/s by HBM bandwidth (B/s) gives the minimum...
# Import jax for this computation.
import jax
import jax.numpy as jnp
import numpy as np

# Initialize list `TPU_SPECS` for the stage values.
TPU_SPECS = [
    {"gen": "v5e", "hbm_gib": 16.0, "bw_gbs": 819.0, "bf16_tflops": 197.0},
    {"gen": "v5p", "hbm_gib": 95.0, "bw_gbs": 2765.0, "bf16_tflops": 459.0},
    {"gen": "v6e", "hbm_gib": 32.0, "bw_gbs": 1600.0, "bf16_tflops": 918.0},
]
# Iterate over `s` to step through the computation:
for s in TPU_SPECS:
    # Run `round` to compute `s['ridge_flops_per_byte']`.
    s["ridge_flops_per_byte"] = round((s["bf16_tflops"] * 1e12) / (s["bw_gbs"] * 1e9), 1)

# Step 2 — Estimate training-state and KV-cache HBM budgets across precisions: Computing both training state and KV-cache GiB upfront prevents...
def estimate_memory_gib(params_billions, kv_tokens, layers=32, kv_heads=8, head_dim=128):
    # Compute `n_params` from `params_billions * 1e9`
    n_params = params_billions * 1e9
    # Compute `kv_elements` from `2.0 * layers * kv_tokens * kv_heads * head_dim`
    kv_elements = 2.0 * layers * kv_tokens * kv_heads * head_dim
    # Evaluate `1024 ** 3` and convert the result into Python scalar/collection `gib`.
    gib = float(1024 ** 3)
    # Return `{'fp32_train_state_gib': round(n_params * 16.0 / gib, 3), 'mixed_bf16_train_state_gib': round(n_params * 18.0 / gib, 3), 'kv_fp32_gib': round(kv_elements * 4.0 / gib, 3), 'kv_bf16_gib': round(kv_elements * 2.0 / gib, 3), 'kv_int8_gib': round(kv_elements * 1.0 / gib, 3)}` to the caller.
    return {
        "fp32_train_state_gib": round((n_params * 16.0) / gib, 3),
        "mixed_bf16_train_state_gib": round((n_params * 18.0) / gib, 3),
        "kv_fp32_gib": round((kv_elements * 4.0) / gib, 3),
        "kv_bf16_gib": round((kv_elements * 2.0) / gib, 3),
        "kv_int8_gib": round((kv_elements * 1.0) / gib, 3),
    }



# Run `estimate_memory_gib` to compute `budget_7b`.
budget_7b = estimate_memory_gib(params_billions=7.0, kv_tokens=8192)
# Run `estimate_memory_gib` to compute `budget_1p5b`.
budget_1p5b = estimate_memory_gib(params_billions=1.5, kv_tokens=4096)
# Assert that `budget_7b["mixed_bf16_train_state_gib"] > 64.0 and budget_7b["mixed_bf16_train_state_gib"] <`.
assert budget_7b["mixed_bf16_train_state_gib"] > 64.0 and budget_7b["mixed_bf16_train_state_gib"] < 128.0
# Assert invariant `budget_1p5b["mixed_bf16_train_state_gib"] < 64.0` holds
assert budget_1p5b["mixed_bf16_train_state_gib"] < 64.0
# Print the observed values to compare against the expected result.
print("TPU ridge points (FLOP/byte):", {s["gen"]: s["ridge_flops_per_byte"] for s in TPU_SPECS})
# Print diagnostic summary of the computed outputs.
print("7B memory budget (GiB):", budget_7b)
# Print diagnostic summary of the computed outputs.
print("1.5B memory budget (GiB):", budget_1p5b)

assert budget_7b["mixed_bf16_train_state_gib"] > 64.0 and budget_7b["mixed_bf16_train_state_gib"] < 128.0
assert budget_1p5b["mixed_bf16_train_state_gib"] < 64.0

Run the example

python3 phases/21-tpu-systems/01-tpu-generations-and-hbm-memory-budgeting/code/main.py
About this command

This shortcut uses an activated macOS/Linux course environment in the repository folder. For Windows, use the explicit interpreter command in “Setup help” above. The script includes worked examples and reference solutions.

Inspect the complete assembled example
main.py · complete Python example
# Step 1 — Define TPU generation specs and compute Roofline ridge points: Dividing peak BF16 FLOP/s by HBM bandwidth (B/s) gives the minimum...
# Import jax for this computation.
import jax
import jax.numpy as jnp
import numpy as np

# Initialize list `TPU_SPECS` for the stage values.
TPU_SPECS = [
    {"gen": "v5e", "hbm_gib": 16.0, "bw_gbs": 819.0, "bf16_tflops": 197.0},
    {"gen": "v5p", "hbm_gib": 95.0, "bw_gbs": 2765.0, "bf16_tflops": 459.0},
    {"gen": "v6e", "hbm_gib": 32.0, "bw_gbs": 1600.0, "bf16_tflops": 918.0},
]
# Iterate over `s` to step through the computation:
for s in TPU_SPECS:
    # Run `round` to compute `s['ridge_flops_per_byte']`.
    s["ridge_flops_per_byte"] = round((s["bf16_tflops"] * 1e12) / (s["bw_gbs"] * 1e9), 1)


# Step 2 — Estimate training-state and KV-cache HBM budgets across precisions: Computing both training state and KV-cache GiB upfront prevents...
def estimate_memory_gib(params_billions, kv_tokens, layers=32, kv_heads=8, head_dim=128):
    # Compute `n_params` from `params_billions * 1e9`
    n_params = params_billions * 1e9
    # Compute `kv_elements` from `2.0 * layers * kv_tokens * kv_heads * head_dim`
    kv_elements = 2.0 * layers * kv_tokens * kv_heads * head_dim
    # Evaluate `1024 ** 3` and convert the result into Python scalar/collection `gib`.
    gib = float(1024 ** 3)
    # Return `{'fp32_train_state_gib': round(n_params * 16.0 / gib, 3), 'mixed_bf16_train_state_gib': round(n_params * 18.0 / gib, 3), 'kv_fp32_gib': round(kv_elements * 4.0 / gib, 3), 'kv_bf16_gib': round(kv_elements * 2.0 / gib, 3), 'kv_int8_gib': round(kv_elements * 1.0 / gib, 3)}` to the caller.
    return {
        "fp32_train_state_gib": round((n_params * 16.0) / gib, 3),
        "mixed_bf16_train_state_gib": round((n_params * 18.0) / gib, 3),
        "kv_fp32_gib": round((kv_elements * 4.0) / gib, 3),
        "kv_bf16_gib": round((kv_elements * 2.0) / gib, 3),
        "kv_int8_gib": round((kv_elements * 1.0) / gib, 3),
    }



# Run `estimate_memory_gib` to compute `budget_7b`.
budget_7b = estimate_memory_gib(params_billions=7.0, kv_tokens=8192)
# Run `estimate_memory_gib` to compute `budget_1p5b`.
budget_1p5b = estimate_memory_gib(params_billions=1.5, kv_tokens=4096)
# Assert that `budget_7b["mixed_bf16_train_state_gib"] > 64.0 and budget_7b["mixed_bf16_train_state_gib"] <`.
assert budget_7b["mixed_bf16_train_state_gib"] > 64.0 and budget_7b["mixed_bf16_train_state_gib"] < 128.0
# Assert invariant `budget_1p5b["mixed_bf16_train_state_gib"] < 64.0` holds
assert budget_1p5b["mixed_bf16_train_state_gib"] < 64.0
# Print the observed values to compare against the expected result.
print("TPU ridge points (FLOP/byte):", {s["gen"]: s["ridge_flops_per_byte"] for s in TPU_SPECS})
# Print diagnostic summary of the computed outputs.
print("7B memory budget (GiB):", budget_7b)
# Print diagnostic summary of the computed outputs.
print("1.5B memory budget (GiB):", budget_1p5b)

Expected result: Prints the ridge-point FLOP/byte ratios for v5e (240.5), v5p (166.0), and v6e (573.8) along with the 7B and 1.5B memory budgets in GiB.

Recorded reference execution

Reference output captured from executing the complete lesson script (code/main.py) and its assertions in the course CPU environment.

2026-10-20T10:58:55.101094+00:00 · Python 3.14.3 · JAX 0.9.2 · 1 CPU device

TPU ridge points (FLOP/byte): {'v5e': 240.5, 'v5p': 166.0, 'v6e': 573.8}
7B memory budget (GiB): {'fp32_train_state_gib': 104.308, 'mixed_bf16_train_state_gib': 117.347, 'kv_fp32_gib': 2.0, 'kv_bf16_gib': 1.0, 'kv_int8_gib': 0.5}
1.5B memory budget (GiB): {'fp32_train_state_gib': 22.352, 'mixed_bf16_train_state_gib': 25.146, 'kv_fp32_gib': 1.0, 'kv_bf16_gib': 0.5, 'kv_int8_gib': 0.25}
TPU ridge points (FLOP/byte): {'v5e': 240.5, 'v5p': 166.0, 'v6e': 573.8}
7B memory budget (GiB): {'fp32_train_state_gib': 104.308, 'mixed_bf16_train_state_gib': 117.347, 'kv_fp32_gib': 2.0, 'kv_bf16_gib': 1.0, 'kv_int8_gib': 0.5}
1.5B memory budget (GiB): {'fp32_train_state_gib': 22.352, 'mixed_bf16_train_state_gib': 25.146, 'kv_fp32_gib': 1.0, 'kv_bf16_gib': 0.5, 'kv_int8_gib': 0.25}
KV cache GiB (FP32 / BF16 / INT8): 2.0 1.0 0.5
7B mixed train state GiB: 117.347 Minimum chips at 80% usable HBM: 10
1.5B state plus BF16 KV cache (GiB): 25.646
Arithmetic intensity: 341.3333333333333 FLOPs/byte
v6e-4 80% usable GiB: 102.4 Fits 7B state estimate: False
PASS: tpu-03

Download execution record · The notebook includes executed cells and figure outputs. Compare your run with these values; timings depend on the machine.

Stage 3 of 5 · Controlled ExperimentsPredict before running & test boundary conditions

Compare KV-cache memory across FP32, BF16, and INT8

Compare KV-cache memory across FP32, BF16, and INT8 · Python experiment
# Experiment — Compare KV-cache memory across FP32, BF16, and INT8: Halving bytes per element halves both the HBM capacity consumed...
# Assert that `np.isclose(budget_7b["kv_fp32_gib"], 2.0 * budget_7b["kv_bf16_gib"])`.
assert np.isclose(budget_7b["kv_fp32_gib"], 2.0 * budget_7b["kv_bf16_gib"])
# Assert that `np.isclose(budget_7b["kv_bf16_gib"], 2.0 * budget_7b["kv_int8_gib"])`.
assert np.isclose(budget_7b["kv_bf16_gib"], 2.0 * budget_7b["kv_int8_gib"])
# Print the observed values to compare against the expected result.
print("KV cache GiB (FP32 / BF16 / INT8):", budget_7b["kv_fp32_gib"], budget_7b["kv_bf16_gib"], budget_7b["kv_int8_gib"])

Expected: For 8192 tokens, KV cache drops from 2.0 GiB (FP32) to 1.0 GiB (BF16) and 0.5 GiB (INT8).

Halving bytes per element halves both the HBM capacity consumed by the KV cache and the HBM bytes read on every decode step.

Verify how many v5e chips are needed to hold a 7B mixed-precision training state

Verify how many v5e chips are needed to hold a 7B mixed-precision training state · Python experiment
v5e_chip_gib = 16.0
usable_fraction = 0.8
min_v5e_chips = int(np.ceil(budget_7b["mixed_bf16_train_state_gib"] / (v5e_chip_gib * usable_fraction)))
print("7B mixed train state GiB:", budget_7b["mixed_bf16_train_state_gib"], "Minimum chips at 80% usable HBM:", min_v5e_chips)
assert min_v5e_chips == 10
assert min_v5e_chips > 8 and min_v5e_chips <= 16

Expected: The 18-byte-per-parameter ledger estimates 117.347 GiB. At 16 GiB per chip and 80% usable HBM, at least 10 chips are needed; the listed 16-chip v5litepod-16 is the first listed topology large enough.

The estimate includes 2-byte BF16 weights, 4-byte FP32 master weights, 4-byte gradients and two 4-byte FP32 Adam moments, for 18 bytes per parameter. This is a named-tensor ledger, not a complete fit guarantee: activations, communication buffers, temporary copies and allocator overhead need additional headroom.

Stage 4 of 5 · Hands-On Exercise & CheckpointModify the code yourself, solve transfer tasks & test your understanding
Hands-On Exercise · Modify & Verify Run in the declared lesson environment

Make it yours

Where to put your code: Append your solution at the bottom of phases/21-tpu-systems/01-tpu-generations-and-hbm-memory-budgeting/code/main.py (or inside the Make it yours cell of exercises/tpu-03.ipynb) and run python3 phases/21-tpu-systems/01-tpu-generations-and-hbm-memory-budgeting/code/main.py to verify your prediction before opening the reference solution.

Compute estimate_memory_gib for a 1.5B parameter model with 4096 KV tokens and verify that its mixed-precision training-state estimate is below 30 GiB and that state plus BF16 KV cache is below 80% of a single-host v5litepod-4 (64 GiB total HBM).

Show a writing guide and starter scaffold

How to write this exercise code: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

  • estimate_memory_gib(...) — Call estimate_memory_gib with your updated parameters or inputs from this lesson's workspace.
  • cache(...) — Call cache with your updated parameters or inputs from this lesson's workspace.
  • assert condition — Verify that the observed output shape, status, or numerical value satisfies the contract.

2. Step-by-step implementation plan:

  1. Set up the changed input or configuration parameters requested in the prompt.
  2. Run the target function or transformation and assert that the output matches your prediction.
Starter code scaffold · copy into your file & fill in the TODOs 
budget_check = estimate_memory_gib(...)  # TODO: compute budget_check
usable_hbm_gib = ...  # TODO: compute usable_hbm_gib
state_and_kv_gib = ...  # TODO: compute state_and_kv_gib
print("1.5B state plus BF16 KV cache (GiB):", state_and_kv_gib)
assert budget_check["mixed_bf16_train_state_gib"]  # TODO: complete assertion check
assert state_and_kv_gib  # TODO: complete assertion check
Reveal the reference solution
Reference solution
budget_check = estimate_memory_gib(params_billions=1.5, kv_tokens=4096)
usable_hbm_gib = 4 * 16.0 * 0.8
state_and_kv_gib = budget_check["mixed_bf16_train_state_gib"] + budget_check["kv_bf16_gib"]
print("1.5B state plus BF16 KV cache (GiB):", state_and_kv_gib)
assert budget_check["mixed_bf16_train_state_gib"] < 30.0
assert state_and_kv_gib < usable_hbm_gib

Compare arithmetic intensity against the TPU v5e and v6e ridge points

Foundations

Compute the arithmetic intensity (FLOPs/byte) of a bfloat16 square matrix multiply of size N=1024N = 1024 (2 * N^3 FLOPs divided by 3 * 2 * N^2 bytes) and check whether it exceeds the v5e ridge point.

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

  • assert condition — Verify that the observed output shape, status, or numerical value satisfies the contract.

2. Step-by-step implementation plan:

  1. Set up the changed input or configuration parameters requested in the prompt.
  2. Run the target function or transformation and assert that the output matches your prediction.
Starter code scaffold · copy into your file & fill in the TODOs 
n = ...  # TODO: compute n
intensity = ...  # TODO: compute intensity
print("Arithmetic intensity:", intensity, "FLOPs/byte")
assert intensity  # TODO: complete assertion check
assert intensity  # TODO: complete assertion check
Show a hint

For two N x N BF16 inputs and one BF16 output, arithmetic intensity simplifies to N / 3.0 FLOPs/byte.

Reveal reference solution and reasoning
Compare arithmetic intensity against the TPU v5e and v6e ridge points · Python solution
n = 1024
intensity = (2 * n**3) / (3 * 2 * n**2)
print("Arithmetic intensity:", intensity, "FLOPs/byte")
assert intensity == n / 3.0
assert intensity > TPU_SPECS[0]["ridge_flops_per_byte"]

The matrix multiply performs about 341.3 FLOPs per byte. That exceeds the v5e ridge point of 240.5 FLOPs/byte, so this idealized operation is compute-bound by the Roofline comparison. Actual kernels may have different data reuse and overhead.

Size a Trillium v6e-4 slice for a 7B training state

Transfer / diagnosis

Compute the total HBM (in GiB) of a 4-chip v6e-4 slice (32 GiB/chip) and check whether budget_7b['mixed_bf16_train_state_gib'] fits inside 80% of that slice's HBM.

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

  • assert condition — Verify that the observed output shape, status, or numerical value satisfies the contract.

2. Step-by-step implementation plan:

  1. Set up the changed input or configuration parameters requested in the prompt.
  2. Run the target function or transformation and assert that the output matches your prediction.
Starter code scaffold · copy into your file & fill in the TODOs 
v6e4_usable_gib = ...  # TODO: compute v6e4_usable_gib
fits_v6e4 = ...  # TODO: compute fits_v6e4
print("v6e-4 80% usable GiB:", v6e4_usable_gib, "Fits 7B state estimate:", fits_v6e4)
assert fits_v6e4 is False  # TODO: complete assertion check
Show a hint

Compare budget_7b['mixed_bf16_train_state_gib'] with 4 * 32.0 * 0.8.

Reveal reference solution and reasoning
Size a Trillium v6e-4 slice for a 7B training state · Python solution
v6e4_usable_gib = 4 * 32.0 * 0.8
fits_v6e4 = budget_7b["mixed_bf16_train_state_gib"] <= v6e4_usable_gib
print("v6e-4 80% usable GiB:", v6e4_usable_gib, "Fits 7B state estimate:", fits_v6e4)
assert fits_v6e4 is False

A 4-chip v6e slice has 128 GiB total HBM and 102.4 GiB at an 80% budget. The 117.347 GiB named-tensor estimate exceeds that budget, before activations and other excluded allocations.

Check your understanding

What does a TPU generation's Roofline ridge point (bf16_tflops / bw_gbs in FLOPs/byte) tell you about a kernel?

Choose one answer

Diagnose the result

If your TPU job fails with RESOURCE_EXHAUSTED: Out of memory while trying to allocate..., compare your parameter + optimizer + activation + KV-cache byte estimate against per-chip HBM (16 GiB on v5e, 32 GiB on v6e) and shard the state across chips (NamedSharding / FSDP) or reduce microbatch size.

Stage 5 of 5 · Keep Your EvidenceRecord your verification receipt & carry takeaways forward

Carry forward

  • Size your TPU slice (v5e, v5p, v6e) from explicit HBM byte budgets for training state and KV caches before launching.
  • Compare your workload's FLOPs/byte against the chip's ridge point to know whether fusion/quantization (bandwidth) or MXU tiling (compute) will speed it up.

Keep your evidence

Download evidence template

Keep the TPU ridge-point table, 7B and 1.5B training/KV-cache HBM budgets, and minimum chip count calculation at an 80% HBM ceiling.

Save your modified code, its output, and an explanation of your prediction. Checkpoint answers and exercise evidence are tracked separately.

Saved in this browser. Export or back up your portfolio anytime in My learning.

Save work for this lesson

Verified on local CPU. Stage to a Cloud TPU VM via Cloud TPU setup to record accelerator execution receipts (TODO: add automated Cloud TPU CI execution receipts).

Primary references