What compute do you need for RL training of LLM agents?
RL training for an LLM agent needs two kinds of compute. GPUs generate the agent's actions and update its weights, and a fleet of isolated CPU sandboxes runs those actions so they can be scored. On Nodus, the GPUs come from a TrainingJob or a Job, and the sandboxes are Nodus Sandboxes that start with no network access and stop at a dollar cap. In our worked example below, the sandboxes are about a tenth of the bill; the expensive part is GPUs waiting while tool calls finish.
That last point decides most design choices, so this post walks through the parts of an agent RL setup, what each one costs, and when the managed path is enough.
What are the moving parts of an agent RL run?
An agent RL run is a loop of three jobs: generate, act, and learn. The policy model writes a reply or a tool call. An environment executes it, returns the result, and the episode continues for several turns. At the end, a grader turns the episode into a reward, and the trainer updates the model.
GRPO is a common choice for agent RL. It samples a group of episodes per task and rewards the ones that beat their group's average. It needs no critic model, which keeps GPU memory down. The DeepSeekMath paper introduced it, and our post on what GRPO is and how to run it covers the method itself.
The three jobs want different hardware:
| Part | What it does | Hardware it needs |
|---|---|---|
| Policy and trainer | Samples actions, computes the loss, updates weights | GPUs sized to the model and method |
| Environment | Runs code, shell commands or tool calls for each episode | Isolated CPU containers, one per live episode |
| Grader | Scores the finished episode, holds the answers | A separate process the model cannot read |
The environment is the part that is new compared with ordinary RL on text. A coding agent needs a real filesystem, a shell and a test runner for every episode in flight. That is why agent RL is as much a sandbox problem as a GPU problem.
How much GPU does the trainer need?
The trainer's GPU follows from the model and method, the same as any fine-tune. With a LoRA adapter on a small model, one GPU is enough. The nodus/grpo-lora runtime's presets put Qwen3-8B with QLoRA on one L40S, RTX A6000 or A100 80GB, and the L40S is listed at $0.793/hr on the Nodus pricing page.
Full-parameter training is another matter. In our PPO vs DPO vs GRPO comparison we worked out that full GRPO on an 8B model needs about 131 GB at 16 bytes per parameter. That means several 80 GB GPUs before you count activations.
The group size also multiplies generation work. The Nodus rl.train recipe samples 8 replies for each of 4 tasks per step, capped at 256 tokens each. That is up to 8,192 generated tokens per step for single-turn tasks. A multi-turn agent episode repeats that generation on every turn, so plan generation capacity per turn, not per task.
How many sandboxes does agent RL need, and what do they cost?
Plan on one sandbox per episode that is running at the same moment. With 16 tasks per step and a group size of 8, that is 128 episodes in flight, so 128 sandboxes. Fewer sandboxes means episodes queue, and the GPUs wait for them.
On Nodus, a 1 vCPU, 2 GiB Sandbox is listed at $0.0179/hr on the pricing page. Keeping 128 of them running costs 128 x $0.0179 = $2.29 per hour. Sandboxes run on CPU machines Nodus operates, not on the GPU machines used for your jobs, so they never take GPU memory from the trainer.
Each one is isolated in ways that matter for code a model wrote. The Sandbox isolation page lists the layers: a user-space kernel that filters system calls, a default-deny firewall with cloud metadata addresses always blocked, a 256-process limit and a non-root user. The Ephemeral continuity mode starts every run empty, which fits a fresh sandbox per task.
Why does GPU wait time matter more than sandbox price?
Because the GPUs bill for every minute a step takes, including the minutes they spend waiting on tool calls. The worked example below shows the size of the effect.
Take an illustrative run on one 8 x H100 node with 128 sandboxes at 1 vCPU and 2 GiB. It trains for 300 steps at 3 minutes of wall-clock time per step, most of it spent waiting on tests and shell commands. These inputs are our assumptions for illustration, not a benchmark. The rates are from the Nodus pricing page, where the H100 SXM is listed at $2.600/hr. Prices checked November 4, 2026.
The run lasts 300 x 3 minutes = 15 hours. The GPUs cost 8 x $2.60 x 15 = $312.00. The sandboxes cost 128 x $0.0179 x 15 = $34.37. The total is $346.37, and the sandboxes are 9.9% of it.
!Cost breakdown of a 15-hour agent RL run: $312.00 for 8 x H100 GPUs and $34.37 for 128 sandboxes.
In this example the sandboxes are under a tenth of the bill; the GPUs are the rest.
Now cut the step time in half, to 1.5 minutes, by running tool calls in parallel or overlapping generation with execution. The run drops to 7.5 hours. GPUs cost $156.00 and sandboxes $17.18, for $173.18 in total.
!Total cost of the same 300-step run at 3 minutes per step ($346.37) and at 1.5 minutes per step ($173.18).
Halving the step time saves $156.00 of GPU time, more than four times the entire $34.37 sandbox bill.
From this arithmetic we get a decision rule. Never starve rollouts of sandboxes to save money, because a sandbox costs cents per hour while every idle minute costs the whole GPU node. Spend your engineering effort on step time: parallel episodes, fast environment resets and generation that does not sit idle during tool calls.
When is the managed Nodus runtime enough?
The managed path fits when your environment's tools are plain Python that runs offline. Nodus Environments are versioned task sets with a grader, and the Environments guide shows nodus/grpo-lora training on multi-turn games and tool-calling environments. The tools run inside the grader, and the trainer never sees the expected answers.
A tool-calling run is a short manifest, copied from that guide:
spec:
runtime: nodus/grpo-lora
environment: {name: tool-test@0.1.1, trainTasks: 21, heldOutTasks: 43, seed: 42}
parameters:
maxTurns: 3
maxCompletionLength: 256
Each episode can call tools for up to maxTurns rounds, and the reward covers the whole episode. Program graders run in grading Sandboxes with no network access, billed to the TrainingJob and counted against its maxCostUSD. Before you launch, the recipe's preview() returns a cost estimate from the runtime's published per-step times.
Grading results are also built to keep training honest. A grader that crashes returns InfrastructureFailure, which carries no reward and does not count toward the pass rate. An outage never trains the model to believe a correct action was wrong.
When do you need your own trainer and sandboxes?
You need your own trainer when tools need the network, a browser, an MCP server or another model as judge. The managed grader has no network, so the Environments guide says those environments fail the build. For them, run an open-source trainer such as verl or SkyRL on GPUs, and give each episode a Nodus Sandbox through the Python SDK.
This sketch shows the sandbox side of one episode, using calls from the Python Sandboxes guide. The policy and score functions are placeholders for your own generation and reward code:
import nodus
def run_episode(task, policy, score, max_turns=6):
sb = nodus.Sandbox.create(cpu=1, memory="2Gi", idle_timeout="5m", max_cost=1)
try:
with sb.open("/workspace/task.json", "w") as f:
f.write(task.setup_json)
history = [task.prompt]
for _ in range(max_turns):
command = policy(history) # your model proposes a shell command
if command is None:
break
p = sb.exec("bash", "-c", command)
history.append(p.stdout.read())
p.wait()
return score(sb, task) # e.g. run the hidden tests
finally:
sb.terminate()
The egress default is deny, so the agent cannot reach the internet unless you open it with network=nodus.Egress.open(). Keep it closed unless a tool needs it, because code in a sandbox can read every secret you put there.
Where does this setup fall short?
Be clear about the limits before you plan around them. Nodus has no managed PPO runtime, so a learned reward model with a value head means running your own trainer. Allow-list egress, where only named hosts are reachable, is available in Sandbox manifests but not yet in the Python SDK. Our worked example uses assumed step times; your environment's real reset and tool latency decide the numbers, so time a few episodes before you size a run.
Training recipes, TrainingJobs and GPU Sandbox state snapshots are Beta; Environments are generally available.
FAQ
Can agent RL run on a single GPU?
Yes, for a LoRA adapter on a model up to about 8B parameters. The smallest nodus/grpo-lora presets, Qwen3-0.6B and Qwen3-1.7B with LoRA, fit on one RTX 4090, listed at $0.290/hr on the pricing page. That is a cheap place to debug an environment before you scale up.
Should rollouts run on GPU sandboxes?
Only if the agent's tools need a GPU, such as an agent that writes and benchmarks CUDA kernels. Shell and test-running tools work fine on CPU sandboxes, which cost cents per hour. Nodus GPU Sandboxes run on dedicated on-demand machines and are never preempted; our GPU sandboxes post compares the options.
How do you stop a reward hack from reading the answers?
Keep the expected answers outside the model's reach. In Nodus grading, code from a completion runs as a separate user that cannot read the answers, and the trainer only sees prompts and metadata. With your own trainer, run hidden tests from a location the agent's sandbox never mounts.
Can I train on environments built for verifiers, OpenEnv or verl?
Many of them, yes. The Environments guide covers verifiers packages, OpenEnv environments and verl datasets with their reward functions unchanged, as long as grading needs no network.
Originally published on the Nodus blog.
Running Captum on Cloud TPUs with TorchTPU: a work log
This is the log of a few days spent running Captum's attribution methods on Cloud TPUs through TorchTPU, in the order the work happened, with what I learned at each step. Unless stated otherwise, measurements are from a Cloud TPU v6e-1 in us-east5-b with torch 2.15.0, torch_tpu 0.2.0, Captum 0.9.0 and transformers 5.22.1. Code, scripts and raw findings are in joinalahmed/captum-on-torchtpu; the last section lists them.
---
What I ended up with
Before the log, the short version:
- Captum works on TorchTPU without modification. Every public captum.attr method, including the LLMAttribution wrappers, runs on torch.device("tpu"). The 27 deterministic methods match CPU to about 1e-7 relative error; the 11 stochastic ones run and differ only through random sampling.
- The one setting that matters for correctness is matmul precision. At the default, TPU attributions differ from CPU by about 1e-3 (1e-2 for LayerGradCam). With torch.set_float32_matmul_precision("highest") they agree to 1e-7 on small models and 1e-5 on Qwen3-0.6B, at no measurable cost for this workload.
- Attribution turned out to be a good lens on TorchTPU's execution modes. On Qwen3-0.6B, Layer Integrated Gradients with 64 steps takes 0.80 s in strict eager mode, 0.36 s with DeferAndFuse and 0.16 s under torch.compile(backend="tpu"), against 4 to 17 s on CPU.
- Along the way I hit, and fixed or filed, a buffer leak for Python scalars, a missing max_unpool2d lowering, a compiled-mode bug that read CPU view arguments past their buffer, a Captum bug that placed a CPU tensor in the compiled graph, eight unregistered ops and 11 stale deny-list entries. Five pull requests are open upstream; five more branches wait on a PR-count limit.
- On a four-chip v6e-4 everything held, every fix passed the tests upstream CI runs, and a device-side trace explained the eager-mode numbers: 363 XLA programs per attribution call and a 2% TensorCore duty cycle.
The rest is how I got there.
---
1. Starting point: why attribution, and what already existed
I wanted a workload that would exercise TorchTPU's execution modes, its numerical defaults and its profiler, and that was small enough to iterate on quickly. Attribution fit. Integrated Gradients interpolates between a baseline and the input over n_steps points, runs a forward and backward pass at each point and integrates the gradients, which means:
- repeated forward and backward passes with identical shapes, which a compiler can exploit;
- Python scalars (step sizes, scaling factors) flowing into device operations, which exercises eager dispatch and caching;
- forward hooks and torch.autograd.grad on intermediate layers, which exercises the autograd integration;
- outputs that are differences of gradients, which make arithmetic error visible.
The first thing I did was check whether anyone had done this. A survey of Captum, SHAP, Inseq, TransformerLens, nnsight and pyvene (the report is in reports/) found no TPU support statements, no TPU CI and no TPU-related issues or pull requests in any of them. Captum has around 355k downloads a month and runs CI on CPU and GPU only.
Learned: there was no prior art to lean on, which meant every result below needed its own CPU comparison rather than a reference I could cite.
---
2. Does anything run at all
I started with a smoke test (scripts/tpu_smoke.py) covering the autograd features attribution methods depend on, on deliberately small models so that any failure would come from the runtime and not from memory.
| Primitive | Result on TPU |
|---|---|
| loss.backward() | pass |
| torch.autograd.grad with respect to inputs | pass |
| forward hooks | pass |
| full backward hooks | pass |
| double backward (create_graph=True) | pass |
| custom autograd.Function | pass |
| torch.func.grad | pass |
| torch.compile(backend="tpu") with backward on a view input | pass |
| Captum IntegratedGradients | pass |
| Captum LayerIntegratedGradients | pass |
All twelve passed. torch.accelerator.current_accelerator() returned tpu, is_available() true, device_count() 1.
Encouraged, I wrote scripts/captum_tpu_matrix.py to run every captum.attr class on an MLP and a small CNN, on CPU and on TPU, and compare the outputs.
- 27 methods matched CPU within 1e-7 to 1e-8: IntegratedGradients, Saliency, InputXGradient, DeepLift, GuidedBackprop, Deconvolution, GuidedGradCam, FeatureAblation, Occlusion, ShapleyValues, LRP and all the Layer* and Neuron* variants.
- 11 ran but did not match bit for bit: GradientShap, ShapleyValueSampling, NoiseTunnel and the *Shap family. They sample random noise and the streams differ between devices, so that is expected.
- Nothing failed for a TPU reason. Lime and KernelShap failed on the first run because sklearn was not installed.
Then a real model. scripts/demo_captum_qwen_tpu.py loads Qwen3-0.6B with AutoModelForCausalLM.from_pretrained(..., device_map="tpu"), which puts the weights on the device directly (2.9 s). LayerIntegratedGradients over the embedding layer, LLMAttribution(FeatureAblation) and LLMGradientAttribution(LayerIntegratedGradients) all ran and the HTML heatmap rendered.
For completeness I tried a few other libraries (scripts/other_xai_libs_tpu.py). SHAP's DeepExplainer and GradientExplainer, pytorch-grad-cam's GradCAM and ScoreCAM, and nnsight's trace with device_map="tpu" all worked. Inseq failed with ValueError: Unknown device tpu from a hard-coded device list, a small upstream fix. TransformerLens failed on a transformers v5 incompatibility unrelated to TPUs.
Learned: TorchTPU registers through the standard accelerator interface, so a library that is device-agnostic in the PyTorch sense needs no TPU-specific code. The failures I did see were allow-lists, not kernels.
---
3. The numbers were off by 1e-3, and why
The matrix above was run with matmul precision set to "highest". My first pass had used the default, and the agreement was much worse. That turned into the single most useful finding of the project.
By default TorchTPU runs float32 matmuls with bfloat16 inputs on the MXU. Model outputs tolerate this. Attributions do not, because they are differences of gradients and a difference magnifies the relative error of its terms.
| Method | Default precision | torch.set_float32_matmul_precision("highest") |
|---|---|---|
| Saliency | 2.2e-3 | 1.7e-7 |
| InputXGradient | 1.7e-3 | 7.0e-8 |
| IntegratedGradients (50 steps) | 2.1e-3 | 4.4e-8 |
| LayerIntegratedGradients | 3.1e-3 | 1.0e-7 |
| LayerGradCam | 1.0e-2 | 2.2e-7 |
| GradientShap | 3.6e-2 | 2.5e-2 (stochastic) |
!Max relative error vs CPU, default vs highest matmul precision
On Qwen3-0.6B in fp32 at "highest", LayerIntegratedGradients on "The capital of France is" with target " Paris" came within 8.5e-6 of CPU and gave the same three top-attributed tokens (" France", " of", "The").
I expected "highest" to cost something. It did not, measurably: on Qwen3 at 32 and 128 steps the warm eager times were 0.86 s and 1.17 s at "highest" versus 0.88 s and 1.08 s at "high". The profiles in section 5 explain why: at this model size the workload is bound by dispatch and host-device traffic, not the MXU.
Learned: set torch.set_float32_matmul_precision("highest") before any attribution work on TPU. It is the difference between 1e-3 and 1e-7, and for this workload it is free.
---
4. A detour: my own batched IG, and a reference that had not converged
Early on I assumed that Captum would be slow on TPU because it loops, and wrote a batched Integrated Gradients (torchtpu_xai) that runs all steps as a single fixed-shape call. Two things came out of that.
First, Captum already does this. With internal_batch_size=None, IntegratedGradients expands the input to n_steps × batch and runs one forward and backward. My timings matched Captum's, and I dropped the separate implementation.
Second, my implementation disagreed with Captum by 0.71 relative error on Qwen3, which took a while to understand. It was not a bug. Captum defaults to Gauss-Legendre quadrature and I had used a right Riemann sum; with method="riemann_right" the two agreed to the printed digit. Checking that, though, showed the reference itself had not converged: Captum on CPU with 256 Gauss-Legendre steps reported a convergence delta of 27.3 against a total attribution of 44.4, and the total moved from -26.6 to -8.2 to +44.4 across 32, 128 and 256 steps. A raw-logit target with a pad-token baseline is simply not a well-conditioned IG setup for a language model.
!IG total attribution vs steps on Qwen3-0.6B
Learned: before comparing two attribution implementations, check that the thing being compared has converged. return_convergence_delta=True is cheap, and I put that advice into the docs I later contributed. Also: read the library before reimplementing it.
---
5. Execution modes, with attribution as the probe
TorchTPU has three ways to run a program: strict eager, where each op is dispatched as called; DeferAndFuse (TORCH_TPU_DEFER_AND_FUSE=1), where ops are deferred and fused into larger programs; and torch.compile(backend="tpu"). Attribution exposes the differences more than a single forward pass does.
Qwen3-0.6B, 53-token prompt, Captum LayerIntegratedGradients, precision "high"
| Steps | Captum, CPU | Captum, TPU strict eager | Captum, TPU DeferAndFuse | Batched IG, torch.compile |
|---|---|---|---|---|
| 32 | 4.48 s | 0.880 s | 0.238 s (first call 26 s) | 0.176 s (first call 68 to 112 s) |
| 128 | 17.04 s | 1.078 s | 0.450 s (first call 29 s) | out of memory (see section 6) |
!LayerIntegratedGradients on Qwen3-0.6B across execution modes
With DeferAndFuse, unmodified Captum ran 3.7 times faster than strict eager and 38 times faster than CPU at 128 steps, for a one-time fusion compile of 26 to 29 s.
The shipped example: Qwen3-0.6B, 42-token prompt, 64 steps, precision "highest"
These are the numbers behind the example I later added to torch_tpu (examples/explainability/captum_attribution.py).
| Mode | Warm median | First call | Max rel. error vs CPU | Peak HBM |
|---|---|---|---|---|
| strict eager | 0.80 s | 5 s | 9.2e-05 | 10.6 GiB |
| DeferAndFuse | 0.36 s | 43 s | 9.3e-05 | 11.1 GiB |
| torch.compile(backend="tpu") | 0.16 s | 130 s | 7.0e-05 | 15.3 GiB |
!The shipped example in three execution modes
On a small MLP the three modes take 0.065, 0.039 and 0.062 s with 2e-7 error each, so compiled mode only pays off once there is enough work per call to amortise the launch.
Three things I had to learn the hard way here:
- First-call time is real. TorchTPU's tier-2 compilation cache is off for single-device programs, so each new process recompiles. The tier-3 cache (TORCH_TPU_TIER3_COMPILATION_CACHE_ROOT=gs://...) persists compiled programs across processes.
- The compiled column did not work out of the box. Captum's LayerIntegratedGradients around a compiled Hugging Face model failed with "expected tensor to be on tpu, got cpu". Section 7 covers the two causes, one in Captum and one in torch_tpu.
- At small sizes TPU time does not scale with steps. On a 5-token prompt it was about 0.55 s whether I ran 8, 32 or 64 steps. That is fixed overhead, and the differences between modes only show at tens of tokens and tens of steps.
Learned: for an unmodified Captum workload, DeferAndFuse is the setting that gives most of the speed-up with no code changes. Compile is faster still but needs patches and a long first call.
---
6. Profiling: where the time actually goes
I profiled LayerIntegratedGradients at 32 steps in each mode with torch.profiler, ProfilerActivity.TPU and TpuProfilerConfig (scripts/profile_ig_modes.py), and opened the traces in XProf.
| Mode | Self CPU total | aten::_to_copy | _to_copy calls | Average per copy | Self TPU total |
|---|---|---|---|---|---|
| strict eager | 1.93 s | 1.35 s (70%) | 522 | 2.6 ms | 5.0 s |
| DeferAndFuse | 0.78 s | 0.29 s | 522 | 0.57 ms | 1.23 s |
!XProf: aten::_to_copy share of host time
In strict eager mode 70% of host time is aten::_to_copy: host-to-device and device-to-host copies that block the Python thread. Captum crosses the boundary constantly, for step sizes, for the scaling of inputs along the path, for the cumsum offsets it uses to slice layer outputs. Each is a synchronous round trip in eager mode. DeferAndFuse does the same 522 copies but each costs a fifth as much, because a copy no longer waits for every preceding op individually.
The same picture from the device side
Later, on the v6e-4, I captured a trace of a single warm 50-step IntegratedGradients call on the small MLP with torch_tpu._internal.profiler and read it with the XProf tools.
| Metric | Value |
|---|---|
| Profiled window | 23.3 ms |
| XLA programs launched for one call | 363 |
| Gaps between programs (median 57 us) | 22.3 ms, 95.8% of the window |
| TensorCore duty cycle | 2.0% |
| Device op time | 476 us |
| X64SplitLow / X64SplitHigh custom calls (302) | 231 us, 48.6% of op time |
| The MLP's four matmul fusions | 5.2 us, 0.02% of the window |
| MXU utilization | 0.0% |
Eager TorchTPU dispatches one XLA program per PyTorch op, so 50 interpolation points become 363 launches with a median 57 us gap between them, and the TensorCore sits idle 98% of the time. Half the device op time is X64SplitLow/X64SplitHigh, the custom calls that split 64-bit Python scalars into two 32-bit halves on arrival, which is the same scalar path as the leak in section 7.1. The matmuls that are the "model" take 5 microseconds. The trace and the full report are in tpu_agent_project/results/.
Memory
Peak HBM is set by the logits tensor. A language-model attribution materialises [steps, T, V] logits before selecting the last position and the target token. For Qwen3 at 128 steps and 53 tokens that is [128, 53, 151936] in fp32, 3.84 GB, alongside the interpolation path and its gradients. That is what ran out of memory for the batched implementation at 128 steps and for everything at 256.
Two user-side changes fix it:
- model(input_ids, logits_to_keep=1) (transformers v5) keeps only the last position. With it, 128 steps in one batch took 0.857 s warm on TPU at 23.9 GB peak, against 11.2 s on CPU.
- Captum's internal_batch_size chunks the steps. 128 steps with internal_batch_size=64 took 1.42 s at 10.6 GiB peak instead of 23.9 GB.
!Peak HBM by step count and internal_batch_size
While measuring this I found the torch.accelerator memory APIs (memory_stats, max_memory_allocated, empty_cache and friends) undocumented for TorchTPU, and wrote the section that became the docs-memory-apis branch.
Learned: eager attribution on TorchTPU is launch-bound and copy-bound, from both the host and the device side, and the fix is a different execution mode, not a faster chip. Memory is a logits problem; logits_to_keep and internal_batch_size are the two knobs.
---
7. The bugs, in the order I hit them
Every item here started as something Captum does that a training loop does not.
7.1 Device buffers for Python scalars were never released (issue #6036, PR #6102)
IG scales inputs by a different Python float at every step. An open issue reported that eager-mode device buffers for distinct Python scalars were never evicted, and I wanted to know whether Captum triggered it. scripts/scalar_leak_check.py:
| Loop | Live allocations added after 5000 iterations |
|---|---|
| x.add_(0.5) (same scalar each time) | +2 |
| x.add_(i / 1024) (new scalar each time) | +4999 |
| after empty_cache() | none released |
It did. Twenty attribution calls with the same n_steps added nothing, but sweeping n_steps from 10 to 109 added 6361 allocations, because each n_steps yields a new set of step-size floats.
!Scalar buffer leak before and after the fix
The fix is a bounded per-thread cache in MakeBuffer in csrc/eager/op_dispatcher.cc, with a regression test. With the patch, the 5000-scalar loop adds 903 allocations (the 1024-entry cache clears four times) and the Captum sweep adds 280. On current main the strict-eager case no longer leaks after a transfer refactor, but DeferAndFuse still does (+5001), so the fix stays relevant.
7.2 aten::max_unpool2d was not implemented (issue #6103, PR #6105)
To get beyond my own two toy models I ported Captum's 78 attribution test configurations into a device-generic suite that runs them on CPU and on CAPTUM_TEST_DEVICE=tpu and compares. Of 266 generated tests, 262 passed and 4 failed, in 5 min 25 s.
One failure was a tolerance miss at a near-zero attribution (+6.1e-5 vs -6.1e-5), precision noise. The other three were one gap: DeepLift's MaxPool rule calls F.max_unpool2d, TorchTPU had no lowering, and every DeepLift variant on a CNN failed with RuntimeError: operator 'aten::max_unpool2d' is not implemented for TPU.
I filed the issue and implemented max_unpool2d and max_unpool3d as a StableHLO batched scatter. 25 unit tests and 4 error-message tests pass on v6e, and the three DeepLift configurations now pass.
7.3 Compiled partitions read CPU view arguments past their buffer (PR #6106)
Chasing the "expected tensor to be on tpu, got cpu" failure under torch.compile, I reduced it and found a different bug on the way: x * cpu_tensor[1].to(tpu) under torch.compile(backend="tpu") returned zero. A StableHLO dump showed the TPU partition compiled against the view's strides and storage offset, while at run time the argument is copied to the device contiguous. The compiled program read past the buffer it was given.
The fix in split_compiler.py compiles TPU partitions against the contiguous layout the argument is actually copied with. Test included, verified on v6e, CI green.
7.4 Captum placed a CPU tensor in the compiled graph (Captum PR #1952)
The other half of the compile failure is in Captum. LayerIntegratedGradients computes output offsets with torch.cumsum on a CPU tensor and slices layer outputs with it; under torch.compile that CPU tensor becomes a graph input on the wrong device. Replacing it with Python integers fixes it, and 12 LayerIntegratedGradients tests pass on TPU.
With #1952 and #6106 applied, the compiled row in section 5 runs: 0.16 s warm, 7.0e-05 error.
7.5 Eight unregistered ops (branch missing-ops-batch)
scripts/probe_unregistered_ops.py calls every ATen op TorchTPU does not register with a plausible input. From the list I implemented eight with tests: smooth_l1_loss and huber_loss with backward, logaddexp, logaddexp2, nansum, std_mean, trace and addbmm. 53 tests pass on v6e against a wheel built from the branch; 23 files, 1733 lines.
The thing I did not know going in: TorchTPU does not accept new composite kernels. If a TT_KERNEL calls another ATen op, OpNameStack aborts at run time with CompositeOpCheck: Op trace is decomposed into as_strided ... implement the trace op by lowering it to SHLO instead, and the permitted-composite list is capped at 150 entries with a comment not to add more. nansum, trace, std_mean and addbmm compiled fine as composites and failed at run time, and I rewrote each as a direct StableHLO lowering: a diagonal mask and reduction for trace; Select(IsNan(x), 0, x) then a sum for nansum; Sqrt(Var) and Mean on the shared reduction code for std_mean; a batched matmul then a sum over the batch axis for addbmm.
Smaller things for the next contributor: IgnoreInCacheKey rejects PromotedScalar because scalars are already excluded from cache keys; UnaryOpOptions has no default constructor, so .op_param_cache_keys must be set; TPU matmul tests need set_float32_matmul_precision("highest") to match CPU.
Still unregistered: triu_indices, median, nanmedian, mode, kthvalue, frexp, renorm, and the .out variants of about a dozen more.
7.6 The bounded-dynamism deny list (branch dynamism-denylist-triage)
TorchTPU's bounded-dynamism tests carry a deny list of ops marked "untriaged: invalid result". I removed the skips for 12 of them and ran the eager and compiled dynamic targets on v6e. On a single-chip host this needs --local_test_jobs=1, otherwise parallel shards fight over /dev/vfio/0.
Eleven pass in both modes: eq, maximum, ne, neg, pow, remainder, round, sign, sinh, sub, trunc. Their entries go. One is a real defect: clamp.Tensor (tensor-valued min and max) mismatches 80% of elements whenever any dimension is bounded, while the scalar variants pass. Its entry now carries the diagnosis and reproduction conditions.
7.7 Smaller items
- Loader error message (branch loader-import-error-context): a wheel built from source only imports against the torch pinned in the repo lockfile, and the failure was a bare undefined symbol c10::report_overflow. The loader now names the glue library it chose and the PyTorch version it found. Test included.
- Memory API docs (branch docs-memory-apis, addresses #4413): the torch.accelerator memory APIs documented with values measured on v6e.
- Triage on #4408: the reported NaN in compiled SDPA backward does not reproduce on 0.2.0 with torch 2.15 (Qwen3; compiled vs CPU 1.85e-5).
- #4409 (compile with backward on a view input) did not reproduce in my smoke test. Noted, not yet verified numerically.
Learned: the bugs clustered around two things Captum does unusually often, Python scalars into device ops and CPU-side bookkeeping tensors next to device tensors. Neither shows up in a training loop. Also: read a project's contributor constraints (the composite-op cap, the three-open-PR limit for outside contributors) before writing the code, not after.
---
8. Four chips: what changes and what does not (2026-11-01)
Everything so far ran on one v6e chip. I queued a v6e-4 (four chips on one host) to answer two questions: do the results hold on a multi-chip host, and does splitting an attribution across chips help. Capacity took several hours to arrive.
Process model. The first thing I learned is that libtpu is exclusive per process. A second process touching the TPU while one is running fails with "The TPU is already in use by process with pid N". Multi-chip work on TorchTPU is one worker process per chip, launched with torchrun after exporting the slice topology:
eval $(python -m torch.tpu.distributed | sed 's/^/export /')
torchrun --nproc_per_node=$WORLD_SIZE captum_multichip_ig.py --steps 256
Inside a worker, torch.accelerator.device_count() is 1 and torch.device("tpu") is that worker's chip. The tpu_dist backend provides the usual torch.distributed collectives.
Correctness. The 38-method matrix on chip 0 gave the same result as the v6e-1: 27 match CPU, 11 stochastic run. I then sharded Integrated Gradients across the four chips, each rank evaluating a quarter of the Riemann grid in one batched forward and backward, with an all_reduce assembling the attribution. It reproduces Captum's single-chip LayerIntegratedGradients to 5.6e-8 when both sides use the same batch size and 1.2e-4 when they do not (reduction order, as in section 3).
Speed. Qwen3-0.6B, fp32, 53 tokens, strict eager, released wheel:
| n_steps | 4 chips, sharded | 1 chip, Captum | speedup |
|---|---|---|---|
| 64 | 0.355 s | 0.219 s (all 64 in one batch) | 0.62x |
| 256 | 0.327 s | 0.547 s (internal_batch_size=64) | 1.67x |
| 256 | 0.357 s | out of memory in a single batch | - |
At 64 steps one chip finishes in a fifth of a second and splitting four ways costs more in dispatch and all-reduce than it saves. At 256 steps one chip cannot hold the batch, Captum has to chunk it into four sequential pieces, and four chips run those in parallel. 1.67x rather than 4x because per-chunk fixed costs and the all-reduce are not amortised at this size.
The fixes on four chips. I rebuilt each branch from source on the host and ran the bazel targets upstream CI uses (_tpu-v6e variant) plus the matching pytest files against the installed wheel. The bounded scalar cache test passes and a four-worker torchrun of the leak check gives every rank the same counts as the single chip; the max_unpool, eight-op and eleven-op dynamic tests pass; the split-compile test passes; the loader test passes; the memory APIs return per-process values; the example runs in all three modes at 2.0e-7. Logs are under results/v6e-4/.
What went wrong. Partway through, the host-side TPU session manager on the first v6e-4 started refusing every process (tpunetd START_SESSION failed). A guest reboot did not clear it, nor did restarting the runtime service. Queued-resource nodes cannot be stopped and started, so the only remedy was to delete the node and queue another with the same persistent build disk attached. The bazel cache on that disk took the replacement from bare image to its first branch test in under 20 minutes, which is the strongest argument I have for keeping build state on a separate disk.
Not covered: the v5e-8 that upstream CI runs on, and tests/distributed/. v5e capacity was unavailable all day and a v5litepod-8 exceeds the project's 4-chip per-zone serving quota.
Learned: multi-chip attribution is worth it only when a single chip runs out of memory. Build from source on the host, with the bazel cache on a detachable disk. Expect a TPU host to need replacing at some point and plan the setup so that costs 20 minutes rather than a day.
---
9. What went upstream
To torch_tpu, on the explainability-example branch:
- examples/explainability/captum_attribution.py: one script with flags for the model (tiny or any Hugging Face causal LM), steps, runs, precision, --compile, --internal-batch-size and --profile-dir. It reports CPU reference time, first-call time, warm median, max relative error against CPU and peak HBM.
- docs/source/case_studies_explainability.py and its Markdown export: a seven-section case study covering running Captum on a TPU, setting matmul precision, comparing execution modes, reading the XProf trace, bounding memory, checking convergence and running the example. Linked from the case-study index and the examples README; mkdocs build --strict passes.
To Captum, tests/attr/test_accelerator.py (meta-pytorch/captum#1951), which re-runs Captum's 78 attribution configurations on CPU and on any accelerator and compares them, so TPU compatibility is covered by tests rather than by this post.
---
10. If you want to try this
- Strict eager and DeferAndFuse work on torch_tpu 0.2.0 as released. Compiled mode needs two unmerged patches (Captum #1952, torch_tpu #6106). Until they land, DeferAndFuse is the no-changes option.
- Set matmul precision to "highest" and check the convergence delta.
- Five branches (missing-ops-batch, explainability-example, dynamism-denylist-triage, docs-memory-apis, loader-import-error-context) are pushed to the fork without upstream PRs yet. The repository allows users without write access three open PRs and the first three fill the quota. Each is tracked by an issue (#6111 to #6114, and a comment on #4413) so review can start before the PR exists.
---
11. Where everything lives
!Contributions map
Open upstream
- google-pytorch/torch_tpu#6102: bounded scalar buffer cache (fixes #6036)
- google-pytorch/torch_tpu#6103: issue, aten::max_unpool2d missing
- google-pytorch/torch_tpu#6105: max_unpool2d / max_unpool3d lowerings
- google-pytorch/torch_tpu#6106: fix compiled TPU partitions reading CPU view arguments past their buffer
- meta-pytorch/captum#1951: device-generic accelerator comparison tests
- meta-pytorch/captum#1952: Python ints for LayerIntegratedGradients output offsets
Pushed, PR pending (fork joinalahmed/torch_tpu)
- missing-ops-batch: google-pytorch/torch_tpu#6111
- explainability-example: google-pytorch/torch_tpu#6114
- dynamism-denylist-triage: google-pytorch/torch_tpu#6112
- docs-memory-apis: comment on google-pytorch/torch_tpu#4413
- loader-import-error-context: google-pytorch/torch_tpu#6113
In this repository (joinalahmed/captum-on-torchtpu)
- images/: the figures in this post, generated by scripts/make_blog_images.py
- FINDINGS.md: the measurements, in the order they were taken
- scripts/: every benchmark, probe and reproduction above, plus tpu_vm.sh for creating, setting up and deleting the VM
- captum_on_torchtpu.ipynb: notebook version of sections 2 to 6
- reports/Explainability tools on TPUs.md: the prior-art survey
- results/v6e-4/: the four-chip matrix, the sharded IG numbers and every branch's multi-chip test log
- tpu_agent_project/: the same workload dispatched through Google's TPU agent suite, with the XProf report quoted in section 6
Cost. The single-chip work ran on a v6e-1 created and deleted per work session; the four-chip work on one v6e-4 from a queued resource after several hours of waiting. VMs and disks bill while they exist; everything was deleted when the work finished.
---
Written with Claude Code as the pair programmer. Commits in torch_tpu carry an Assisted-by trailer, as that project's contributing guide requests; AI is not listed as a co-author there.
Mentions
What compute do you need for RL training of LLM agents? RL training for an LLM agent needs two kinds of compute. GPUs generate the agent's actions and update its weights, and a fleet of isolated CPU sandboxes runs those actions so they can be scored. On Nodus, the GPUs come from a TrainingJob or a Job, and the sandboxes are Nodus Sandboxes that start with no network access and stop at a dollar cap. In our worked example below, the sandboxes are about a tenth of the bill; the expensive part is GPUs waiting while tool calls finish. That last point decides most design choices, so this post walks through the parts of an agent RL setup, what each one costs, and when the managed path is enough. What are the moving parts of an agent RL run? An agent RL run is a loop of three jobs: generate, act, and learn. The policy model writes a reply or a tool call. An environment executes it, returns the result, and the episode continues for several turns. At the end, a grader turns the episode into a reward, and the trainer updates the model. GRPO is a common choice for agent RL. It samples a group of episodes per task and rewards the ones that beat their group's average. It needs no critic model, which keeps GPU memory down. The DeepSeekMath paper introduced it, and our post on what GRPO is and how to run it covers the method itself. The three jobs want different hardware: | Part | What it does | Hardware it needs | |---|---|---| | Policy and trainer | Samples actions, computes the loss, updates weights | GPUs sized to the model and method | | Environment | Runs code, shell commands or tool calls for each episode | Isolated CPU containers, one per live episode | | Grader | Scores the finished episode, holds the answers | A separate process the model cannot read | The environment is the part that is new compared with ordinary RL on text. A coding agent needs a real filesystem, a shell and a test runner for every episode in flight. That is why agent RL is as much a sandbox problem as a GPU problem. How much GPU does the trainer need? The trainer's GPU follows from the model and method, the same as any fine-tune. With a LoRA adapter on a small model, one GPU is enough. The nodus/grpo-lora runtime's presets put Qwen3-8B with QLoRA on one L40S, RTX A6000 or A100 80GB, and the L40S is listed at $0.793/hr on the Nodus pricing page. Full-parameter training is another matter. In our PPO vs DPO vs GRPO comparison we worked out that full GRPO on an 8B model needs about 131 GB at 16 bytes per parameter. That means several 80 GB GPUs before you count activations. The group size also multiplies generation work. The Nodus rl.train recipe samples 8 replies for each of 4 tasks per step, capped at 256 tokens each. That is up to 8,192 generated tokens per step for single-turn tasks. A multi-turn agent episode repeats that generation on every turn, so plan generation capacity per turn, not per task. How many sandboxes does agent RL need, and what do they cost? Plan on one sandbox per episode that is running at the same moment. With 16 tasks per step and a group size of 8, that is 128 episodes in flight, so 128 sandboxes. Fewer sandboxes means episodes queue, and the GPUs wait for them. On Nodus, a 1 vCPU, 2 GiB Sandbox is listed at $0.0179/hr on the pricing page. Keeping 128 of them running costs 128 x $0.0179 = $2.29 per hour. Sandboxes run on CPU machines Nodus operates, not on the GPU machines used for your jobs, so they never take GPU memory from the trainer. Each one is isolated in ways that matter for code a model wrote. The Sandbox isolation page lists the layers: a user-space kernel that filters system calls, a default-deny firewall with cloud metadata addresses always blocked, a 256-process limit and a non-root user. The Ephemeral continuity mode starts every run empty, which fits a fresh sandbox per task. Why does GPU wait time matter more than sandbox price? Because the GPUs bill for every minute a step takes, including the minutes they spend waiting on tool calls. The worked example below shows the size of the effect. Take an illustrative run on one 8 x H100 node with 128 sandboxes at 1 vCPU and 2 GiB. It trains for 300 steps at 3 minutes of wall-clock time per step, most of it spent waiting on tests and shell commands. These inputs are our assumptions for illustration, not a benchmark. The rates are from the Nodus pricing page, where the H100 SXM is listed at $2.600/hr. Prices checked November 4, 2026. The run lasts 300 x 3 minutes = 15 hours. The GPUs cost 8 x $2.60 x 15 = $312.00. The sandboxes cost 128 x $0.0179 x 15 = $34.37. The total is $346.37, and the sandboxes are 9.9% of it. !Cost breakdown of a 15-hour agent RL run: $312.00 for 8 x H100 GPUs and $34.37 for 128 sandboxes. In this example the sandboxes are under a tenth of the bill; the GPUs are the rest. Now cut the step time in half, to 1.5 minutes, by running tool calls in parallel or overlapping generation with execution. The run drops to 7.5 hours. GPUs cost $156.00 and sandboxes $17.18, for $173.18 in total. !Total cost of the same 300-step run at 3 minutes per step ($346.37) and at 1.5 minutes per step ($173.18). Halving the step time saves $156.00 of GPU time, more than four times the entire $34.37 sandbox bill. From this arithmetic we get a decision rule. Never starve rollouts of sandboxes to save money, because a sandbox costs cents per hour while every idle minute costs the whole GPU node. Spend your engineering effort on step time: parallel episodes, fast environment resets and generation that does not sit idle during tool calls. When is the managed Nodus runtime enough? The managed path fits when your environment's tools are plain Python that runs offline. Nodus Environments are versioned task sets with a grader, and the Environments guide shows nodus/grpo-lora training on multi-turn games and tool-calling environments. The tools run inside the grader, and the trainer never sees the expected answers. A tool-calling run is a short manifest, copied from that guide: spec: runtime: nodus/grpo-lora environment: {name: tool-test@0.1.1, trainTasks: 21, heldOutTasks: 43, seed: 42} parameters: maxTurns: 3 maxCompletionLength: 256 Each episode can call tools for up to maxTurns rounds, and the reward covers the whole episode. Program graders run in grading Sandboxes with no network access, billed to the TrainingJob and counted against its maxCostUSD. Before you launch, the recipe's preview() returns a cost estimate from the runtime's published per-step times. Grading results are also built to keep training honest. A grader that crashes returns InfrastructureFailure, which carries no reward and does not count toward the pass rate. An outage never trains the model to believe a correct action was wrong. When do you need your own trainer and sandboxes? You need your own trainer when tools need the network, a browser, an MCP server or another model as judge. The managed grader has no network, so the Environments guide says those environments fail the build. For them, run an open-source trainer such as verl or SkyRL on GPUs, and give each episode a Nodus Sandbox through the Python SDK. This sketch shows the sandbox side of one episode, using calls from the Python Sandboxes guide. The policy and score functions are placeholders for your own generation and reward code: import nodus def run_episode(task, policy, score, max_turns=6): sb = nodus.Sandbox.create(cpu=1, memory="2Gi", idle_timeout="5m", max_cost=1) try: with sb.open("/workspace/task.json", "w") as f: f.write(task.setup_json) history = [task.prompt] for _ in range(max_turns): command = policy(history) # your model proposes a shell command if command is None: break p = sb.exec("bash", "-c", command) history.append(p.stdout.read()) p.wait() return score(sb, task) # e.g. run the hidden tests finally: sb.terminate() The egress default is deny, so the agent cannot reach the internet unless you open it with network=nodus.Egress.open(). Keep it closed unless a tool needs it, because code in a sandbox can read every secret you put there. Where does this setup fall short? Be clear about the limits before you plan around them. Nodus has no managed PPO runtime, so a learned reward model with a value head means running your own trainer. Allow-list egress, where only named hosts are reachable, is available in Sandbox manifests but not yet in the Python SDK. Our worked example uses assumed step times; your environment's real reset and tool latency decide the numbers, so time a few episodes before you size a run. Training recipes, TrainingJobs and GPU Sandbox state snapshots are Beta; Environments are generally available. FAQ Can agent RL run on a single GPU? Yes, for a LoRA adapter on a model up to about 8B parameters. The smallest nodus/grpo-lora presets, Qwen3-0.6B and Qwen3-1.7B with LoRA, fit on one RTX 4090, listed at $0.290/hr on the pricing page. That is a cheap place to debug an environment before you scale up. Should rollouts run on GPU sandboxes? Only if the agent's tools need a GPU, such as an agent that writes and benchmarks CUDA kernels. Shell and test-running tools work fine on CPU sandboxes, which cost cents per hour. Nodus GPU Sandboxes run on dedicated on-demand machines and are never preempted; our GPU sandboxes post compares the options. How do you stop a reward hack from reading the answers? Keep the expected answers outside the model's reach. In Nodus grading, code from a completion runs as a separate user that cannot read the answers, and the trainer only sees prompts and metadata. With your own trainer, run hidden tests from a location the agent's sandbox never mounts. Can I train on environments built for verifiers, OpenEnv or verl? Many of them, yes. The Environments guide covers verifiers packages, OpenEnv environments and verl datasets with their reward functions unchanged, as long as grading needs no network. Originally published on the Nodus blog.
Running Captum on Cloud TPUs with TorchTPU: a work log This is the log of a few days spent running Captum's attribution methods on Cloud TPUs through TorchTPU, in the order the work happened, with what I learned at each step. Unless stated otherwise, measurements are from a Cloud TPU v6e-1 in us-east5-b with torch 2.15.0, torch_tpu 0.2.0, Captum 0.9.0 and transformers 5.22.1. Code, scripts and raw findings are in joinalahmed/captum-on-torchtpu; the last section lists them. --- What I ended up with Before the log, the short version: - Captum works on TorchTPU without modification. Every public captum.attr method, including the LLMAttribution wrappers, runs on torch.device("tpu"). The 27 deterministic methods match CPU to about 1e-7 relative error; the 11 stochastic ones run and differ only through random sampling. - The one setting that matters for correctness is matmul precision. At the default, TPU attributions differ from CPU by about 1e-3 (1e-2 for LayerGradCam). With torch.set_float32_matmul_precision("highest") they agree to 1e-7 on small models and 1e-5 on Qwen3-0.6B, at no measurable cost for this workload. - Attribution turned out to be a good lens on TorchTPU's execution modes. On Qwen3-0.6B, Layer Integrated Gradients with 64 steps takes 0.80 s in strict eager mode, 0.36 s with DeferAndFuse and 0.16 s under torch.compile(backend="tpu"), against 4 to 17 s on CPU. - Along the way I hit, and fixed or filed, a buffer leak for Python scalars, a missing max_unpool2d lowering, a compiled-mode bug that read CPU view arguments past their buffer, a Captum bug that placed a CPU tensor in the compiled graph, eight unregistered ops and 11 stale deny-list entries. Five pull requests are open upstream; five more branches wait on a PR-count limit. - On a four-chip v6e-4 everything held, every fix passed the tests upstream CI runs, and a device-side trace explained the eager-mode numbers: 363 XLA programs per attribution call and a 2% TensorCore duty cycle. The rest is how I got there. --- 1. Starting point: why attribution, and what already existed I wanted a workload that would exercise TorchTPU's execution modes, its numerical defaults and its profiler, and that was small enough to iterate on quickly. Attribution fit. Integrated Gradients interpolates between a baseline and the input over n_steps points, runs a forward and backward pass at each point and integrates the gradients, which means: - repeated forward and backward passes with identical shapes, which a compiler can exploit; - Python scalars (step sizes, scaling factors) flowing into device operations, which exercises eager dispatch and caching; - forward hooks and torch.autograd.grad on intermediate layers, which exercises the autograd integration; - outputs that are differences of gradients, which make arithmetic error visible. The first thing I did was check whether anyone had done this. A survey of Captum, SHAP, Inseq, TransformerLens, nnsight and pyvene (the report is in reports/) found no TPU support statements, no TPU CI and no TPU-related issues or pull requests in any of them. Captum has around 355k downloads a month and runs CI on CPU and GPU only. Learned: there was no prior art to lean on, which meant every result below needed its own CPU comparison rather than a reference I could cite. --- 2. Does anything run at all I started with a smoke test (scripts/tpu_smoke.py) covering the autograd features attribution methods depend on, on deliberately small models so that any failure would come from the runtime and not from memory. | Primitive | Result on TPU | |---|---| | loss.backward() | pass | | torch.autograd.grad with respect to inputs | pass | | forward hooks | pass | | full backward hooks | pass | | double backward (create_graph=True) | pass | | custom autograd.Function | pass | | torch.func.grad | pass | | torch.compile(backend="tpu") with backward on a view input | pass | | Captum IntegratedGradients | pass | | Captum LayerIntegratedGradients | pass | All twelve passed. torch.accelerator.current_accelerator() returned tpu, is_available() true, device_count() 1. Encouraged, I wrote scripts/captum_tpu_matrix.py to run every captum.attr class on an MLP and a small CNN, on CPU and on TPU, and compare the outputs. - 27 methods matched CPU within 1e-7 to 1e-8: IntegratedGradients, Saliency, InputXGradient, DeepLift, GuidedBackprop, Deconvolution, GuidedGradCam, FeatureAblation, Occlusion, ShapleyValues, LRP and all the Layer* and Neuron* variants. - 11 ran but did not match bit for bit: GradientShap, ShapleyValueSampling, NoiseTunnel and the *Shap family. They sample random noise and the streams differ between devices, so that is expected. - Nothing failed for a TPU reason. Lime and KernelShap failed on the first run because sklearn was not installed. Then a real model. scripts/demo_captum_qwen_tpu.py loads Qwen3-0.6B with AutoModelForCausalLM.from_pretrained(..., device_map="tpu"), which puts the weights on the device directly (2.9 s). LayerIntegratedGradients over the embedding layer, LLMAttribution(FeatureAblation) and LLMGradientAttribution(LayerIntegratedGradients) all ran and the HTML heatmap rendered. For completeness I tried a few other libraries (scripts/other_xai_libs_tpu.py). SHAP's DeepExplainer and GradientExplainer, pytorch-grad-cam's GradCAM and ScoreCAM, and nnsight's trace with device_map="tpu" all worked. Inseq failed with ValueError: Unknown device tpu from a hard-coded device list, a small upstream fix. TransformerLens failed on a transformers v5 incompatibility unrelated to TPUs. Learned: TorchTPU registers through the standard accelerator interface, so a library that is device-agnostic in the PyTorch sense needs no TPU-specific code. The failures I did see were allow-lists, not kernels. --- 3. The numbers were off by 1e-3, and why The matrix above was run with matmul precision set to "highest". My first pass had used the default, and the agreement was much worse. That turned into the single most useful finding of the project. By default TorchTPU runs float32 matmuls with bfloat16 inputs on the MXU. Model outputs tolerate this. Attributions do not, because they are differences of gradients and a difference magnifies the relative error of its terms. | Method | Default precision | torch.set_float32_matmul_precision("highest") | |---|---|---| | Saliency | 2.2e-3 | 1.7e-7 | | InputXGradient | 1.7e-3 | 7.0e-8 | | IntegratedGradients (50 steps) | 2.1e-3 | 4.4e-8 | | LayerIntegratedGradients | 3.1e-3 | 1.0e-7 | | LayerGradCam | 1.0e-2 | 2.2e-7 | | GradientShap | 3.6e-2 | 2.5e-2 (stochastic) | !Max relative error vs CPU, default vs highest matmul precision On Qwen3-0.6B in fp32 at "highest", LayerIntegratedGradients on "The capital of France is" with target " Paris" came within 8.5e-6 of CPU and gave the same three top-attributed tokens (" France", " of", "The"). I expected "highest" to cost something. It did not, measurably: on Qwen3 at 32 and 128 steps the warm eager times were 0.86 s and 1.17 s at "highest" versus 0.88 s and 1.08 s at "high". The profiles in section 5 explain why: at this model size the workload is bound by dispatch and host-device traffic, not the MXU. Learned: set torch.set_float32_matmul_precision("highest") before any attribution work on TPU. It is the difference between 1e-3 and 1e-7, and for this workload it is free. --- 4. A detour: my own batched IG, and a reference that had not converged Early on I assumed that Captum would be slow on TPU because it loops, and wrote a batched Integrated Gradients (torchtpu_xai) that runs all steps as a single fixed-shape call. Two things came out of that. First, Captum already does this. With internal_batch_size=None, IntegratedGradients expands the input to n_steps × batch and runs one forward and backward. My timings matched Captum's, and I dropped the separate implementation. Second, my implementation disagreed with Captum by 0.71 relative error on Qwen3, which took a while to understand. It was not a bug. Captum defaults to Gauss-Legendre quadrature and I had used a right Riemann sum; with method="riemann_right" the two agreed to the printed digit. Checking that, though, showed the reference itself had not converged: Captum on CPU with 256 Gauss-Legendre steps reported a convergence delta of 27.3 against a total attribution of 44.4, and the total moved from -26.6 to -8.2 to +44.4 across 32, 128 and 256 steps. A raw-logit target with a pad-token baseline is simply not a well-conditioned IG setup for a language model. !IG total attribution vs steps on Qwen3-0.6B Learned: before comparing two attribution implementations, check that the thing being compared has converged. return_convergence_delta=True is cheap, and I put that advice into the docs I later contributed. Also: read the library before reimplementing it. --- 5. Execution modes, with attribution as the probe TorchTPU has three ways to run a program: strict eager, where each op is dispatched as called; DeferAndFuse (TORCH_TPU_DEFER_AND_FUSE=1), where ops are deferred and fused into larger programs; and torch.compile(backend="tpu"). Attribution exposes the differences more than a single forward pass does. Qwen3-0.6B, 53-token prompt, Captum LayerIntegratedGradients, precision "high" | Steps | Captum, CPU | Captum, TPU strict eager | Captum, TPU DeferAndFuse | Batched IG, torch.compile | |---|---|---|---|---| | 32 | 4.48 s | 0.880 s | 0.238 s (first call 26 s) | 0.176 s (first call 68 to 112 s) | | 128 | 17.04 s | 1.078 s | 0.450 s (first call 29 s) | out of memory (see section 6) | !LayerIntegratedGradients on Qwen3-0.6B across execution modes With DeferAndFuse, unmodified Captum ran 3.7 times faster than strict eager and 38 times faster than CPU at 128 steps, for a one-time fusion compile of 26 to 29 s. The shipped example: Qwen3-0.6B, 42-token prompt, 64 steps, precision "highest" These are the numbers behind the example I later added to torch_tpu (examples/explainability/captum_attribution.py). | Mode | Warm median | First call | Max rel. error vs CPU | Peak HBM | |---|---|---|---|---| | strict eager | 0.80 s | 5 s | 9.2e-05 | 10.6 GiB | | DeferAndFuse | 0.36 s | 43 s | 9.3e-05 | 11.1 GiB | | torch.compile(backend="tpu") | 0.16 s | 130 s | 7.0e-05 | 15.3 GiB | !The shipped example in three execution modes On a small MLP the three modes take 0.065, 0.039 and 0.062 s with 2e-7 error each, so compiled mode only pays off once there is enough work per call to amortise the launch. Three things I had to learn the hard way here: - First-call time is real. TorchTPU's tier-2 compilation cache is off for single-device programs, so each new process recompiles. The tier-3 cache (TORCH_TPU_TIER3_COMPILATION_CACHE_ROOT=gs://...) persists compiled programs across processes. - The compiled column did not work out of the box. Captum's LayerIntegratedGradients around a compiled Hugging Face model failed with "expected tensor to be on tpu, got cpu". Section 7 covers the two causes, one in Captum and one in torch_tpu. - At small sizes TPU time does not scale with steps. On a 5-token prompt it was about 0.55 s whether I ran 8, 32 or 64 steps. That is fixed overhead, and the differences between modes only show at tens of tokens and tens of steps. Learned: for an unmodified Captum workload, DeferAndFuse is the setting that gives most of the speed-up with no code changes. Compile is faster still but needs patches and a long first call. --- 6. Profiling: where the time actually goes I profiled LayerIntegratedGradients at 32 steps in each mode with torch.profiler, ProfilerActivity.TPU and TpuProfilerConfig (scripts/profile_ig_modes.py), and opened the traces in XProf. | Mode | Self CPU total | aten::_to_copy | _to_copy calls | Average per copy | Self TPU total | |---|---|---|---|---|---| | strict eager | 1.93 s | 1.35 s (70%) | 522 | 2.6 ms | 5.0 s | | DeferAndFuse | 0.78 s | 0.29 s | 522 | 0.57 ms | 1.23 s | !XProf: aten::_to_copy share of host time In strict eager mode 70% of host time is aten::_to_copy: host-to-device and device-to-host copies that block the Python thread. Captum crosses the boundary constantly, for step sizes, for the scaling of inputs along the path, for the cumsum offsets it uses to slice layer outputs. Each is a synchronous round trip in eager mode. DeferAndFuse does the same 522 copies but each costs a fifth as much, because a copy no longer waits for every preceding op individually. The same picture from the device side Later, on the v6e-4, I captured a trace of a single warm 50-step IntegratedGradients call on the small MLP with torch_tpu._internal.profiler and read it with the XProf tools. | Metric | Value | |---|---| | Profiled window | 23.3 ms | | XLA programs launched for one call | 363 | | Gaps between programs (median 57 us) | 22.3 ms, 95.8% of the window | | TensorCore duty cycle | 2.0% | | Device op time | 476 us | | X64SplitLow / X64SplitHigh custom calls (302) | 231 us, 48.6% of op time | | The MLP's four matmul fusions | 5.2 us, 0.02% of the window | | MXU utilization | 0.0% | Eager TorchTPU dispatches one XLA program per PyTorch op, so 50 interpolation points become 363 launches with a median 57 us gap between them, and the TensorCore sits idle 98% of the time. Half the device op time is X64SplitLow/X64SplitHigh, the custom calls that split 64-bit Python scalars into two 32-bit halves on arrival, which is the same scalar path as the leak in section 7.1. The matmuls that are the "model" take 5 microseconds. The trace and the full report are in tpu_agent_project/results/. Memory Peak HBM is set by the logits tensor. A language-model attribution materialises [steps, T, V] logits before selecting the last position and the target token. For Qwen3 at 128 steps and 53 tokens that is [128, 53, 151936] in fp32, 3.84 GB, alongside the interpolation path and its gradients. That is what ran out of memory for the batched implementation at 128 steps and for everything at 256. Two user-side changes fix it: - model(input_ids, logits_to_keep=1) (transformers v5) keeps only the last position. With it, 128 steps in one batch took 0.857 s warm on TPU at 23.9 GB peak, against 11.2 s on CPU. - Captum's internal_batch_size chunks the steps. 128 steps with internal_batch_size=64 took 1.42 s at 10.6 GiB peak instead of 23.9 GB. !Peak HBM by step count and internal_batch_size While measuring this I found the torch.accelerator memory APIs (memory_stats, max_memory_allocated, empty_cache and friends) undocumented for TorchTPU, and wrote the section that became the docs-memory-apis branch. Learned: eager attribution on TorchTPU is launch-bound and copy-bound, from both the host and the device side, and the fix is a different execution mode, not a faster chip. Memory is a logits problem; logits_to_keep and internal_batch_size are the two knobs. --- 7. The bugs, in the order I hit them Every item here started as something Captum does that a training loop does not. 7.1 Device buffers for Python scalars were never released (issue #6036, PR #6102) IG scales inputs by a different Python float at every step. An open issue reported that eager-mode device buffers for distinct Python scalars were never evicted, and I wanted to know whether Captum triggered it. scripts/scalar_leak_check.py: | Loop | Live allocations added after 5000 iterations | |---|---| | x.add_(0.5) (same scalar each time) | +2 | | x.add_(i / 1024) (new scalar each time) | +4999 | | after empty_cache() | none released | It did. Twenty attribution calls with the same n_steps added nothing, but sweeping n_steps from 10 to 109 added 6361 allocations, because each n_steps yields a new set of step-size floats. !Scalar buffer leak before and after the fix The fix is a bounded per-thread cache in MakeBuffer in csrc/eager/op_dispatcher.cc, with a regression test. With the patch, the 5000-scalar loop adds 903 allocations (the 1024-entry cache clears four times) and the Captum sweep adds 280. On current main the strict-eager case no longer leaks after a transfer refactor, but DeferAndFuse still does (+5001), so the fix stays relevant. 7.2 aten::max_unpool2d was not implemented (issue #6103, PR #6105) To get beyond my own two toy models I ported Captum's 78 attribution test configurations into a device-generic suite that runs them on CPU and on CAPTUM_TEST_DEVICE=tpu and compares. Of 266 generated tests, 262 passed and 4 failed, in 5 min 25 s. One failure was a tolerance miss at a near-zero attribution (+6.1e-5 vs -6.1e-5), precision noise. The other three were one gap: DeepLift's MaxPool rule calls F.max_unpool2d, TorchTPU had no lowering, and every DeepLift variant on a CNN failed with RuntimeError: operator 'aten::max_unpool2d' is not implemented for TPU. I filed the issue and implemented max_unpool2d and max_unpool3d as a StableHLO batched scatter. 25 unit tests and 4 error-message tests pass on v6e, and the three DeepLift configurations now pass. 7.3 Compiled partitions read CPU view arguments past their buffer (PR #6106) Chasing the "expected tensor to be on tpu, got cpu" failure under torch.compile, I reduced it and found a different bug on the way: x * cpu_tensor[1].to(tpu) under torch.compile(backend="tpu") returned zero. A StableHLO dump showed the TPU partition compiled against the view's strides and storage offset, while at run time the argument is copied to the device contiguous. The compiled program read past the buffer it was given. The fix in split_compiler.py compiles TPU partitions against the contiguous layout the argument is actually copied with. Test included, verified on v6e, CI green. 7.4 Captum placed a CPU tensor in the compiled graph (Captum PR #1952) The other half of the compile failure is in Captum. LayerIntegratedGradients computes output offsets with torch.cumsum on a CPU tensor and slices layer outputs with it; under torch.compile that CPU tensor becomes a graph input on the wrong device. Replacing it with Python integers fixes it, and 12 LayerIntegratedGradients tests pass on TPU. With #1952 and #6106 applied, the compiled row in section 5 runs: 0.16 s warm, 7.0e-05 error. 7.5 Eight unregistered ops (branch missing-ops-batch) scripts/probe_unregistered_ops.py calls every ATen op TorchTPU does not register with a plausible input. From the list I implemented eight with tests: smooth_l1_loss and huber_loss with backward, logaddexp, logaddexp2, nansum, std_mean, trace and addbmm. 53 tests pass on v6e against a wheel built from the branch; 23 files, 1733 lines. The thing I did not know going in: TorchTPU does not accept new composite kernels. If a TT_KERNEL calls another ATen op, OpNameStack aborts at run time with CompositeOpCheck: Op trace is decomposed into as_strided ... implement the trace op by lowering it to SHLO instead, and the permitted-composite list is capped at 150 entries with a comment not to add more. nansum, trace, std_mean and addbmm compiled fine as composites and failed at run time, and I rewrote each as a direct StableHLO lowering: a diagonal mask and reduction for trace; Select(IsNan(x), 0, x) then a sum for nansum; Sqrt(Var) and Mean on the shared reduction code for std_mean; a batched matmul then a sum over the batch axis for addbmm. Smaller things for the next contributor: IgnoreInCacheKey rejects PromotedScalar because scalars are already excluded from cache keys; UnaryOpOptions has no default constructor, so .op_param_cache_keys must be set; TPU matmul tests need set_float32_matmul_precision("highest") to match CPU. Still unregistered: triu_indices, median, nanmedian, mode, kthvalue, frexp, renorm, and the .out variants of about a dozen more. 7.6 The bounded-dynamism deny list (branch dynamism-denylist-triage) TorchTPU's bounded-dynamism tests carry a deny list of ops marked "untriaged: invalid result". I removed the skips for 12 of them and ran the eager and compiled dynamic targets on v6e. On a single-chip host this needs --local_test_jobs=1, otherwise parallel shards fight over /dev/vfio/0. Eleven pass in both modes: eq, maximum, ne, neg, pow, remainder, round, sign, sinh, sub, trunc. Their entries go. One is a real defect: clamp.Tensor (tensor-valued min and max) mismatches 80% of elements whenever any dimension is bounded, while the scalar variants pass. Its entry now carries the diagnosis and reproduction conditions. 7.7 Smaller items - Loader error message (branch loader-import-error-context): a wheel built from source only imports against the torch pinned in the repo lockfile, and the failure was a bare undefined symbol c10::report_overflow. The loader now names the glue library it chose and the PyTorch version it found. Test included. - Memory API docs (branch docs-memory-apis, addresses #4413): the torch.accelerator memory APIs documented with values measured on v6e. - Triage on #4408: the reported NaN in compiled SDPA backward does not reproduce on 0.2.0 with torch 2.15 (Qwen3; compiled vs CPU 1.85e-5). - #4409 (compile with backward on a view input) did not reproduce in my smoke test. Noted, not yet verified numerically. Learned: the bugs clustered around two things Captum does unusually often, Python scalars into device ops and CPU-side bookkeeping tensors next to device tensors. Neither shows up in a training loop. Also: read a project's contributor constraints (the composite-op cap, the three-open-PR limit for outside contributors) before writing the code, not after. --- 8. Four chips: what changes and what does not (2026-11-01) Everything so far ran on one v6e chip. I queued a v6e-4 (four chips on one host) to answer two questions: do the results hold on a multi-chip host, and does splitting an attribution across chips help. Capacity took several hours to arrive. Process model. The first thing I learned is that libtpu is exclusive per process. A second process touching the TPU while one is running fails with "The TPU is already in use by process with pid N". Multi-chip work on TorchTPU is one worker process per chip, launched with torchrun after exporting the slice topology: eval $(python -m torch.tpu.distributed | sed 's/^/export /') torchrun --nproc_per_node=$WORLD_SIZE captum_multichip_ig.py --steps 256 Inside a worker, torch.accelerator.device_count() is 1 and torch.device("tpu") is that worker's chip. The tpu_dist backend provides the usual torch.distributed collectives. Correctness. The 38-method matrix on chip 0 gave the same result as the v6e-1: 27 match CPU, 11 stochastic run. I then sharded Integrated Gradients across the four chips, each rank evaluating a quarter of the Riemann grid in one batched forward and backward, with an all_reduce assembling the attribution. It reproduces Captum's single-chip LayerIntegratedGradients to 5.6e-8 when both sides use the same batch size and 1.2e-4 when they do not (reduction order, as in section 3). Speed. Qwen3-0.6B, fp32, 53 tokens, strict eager, released wheel: | n_steps | 4 chips, sharded | 1 chip, Captum | speedup | |---|---|---|---| | 64 | 0.355 s | 0.219 s (all 64 in one batch) | 0.62x | | 256 | 0.327 s | 0.547 s (internal_batch_size=64) | 1.67x | | 256 | 0.357 s | out of memory in a single batch | - | At 64 steps one chip finishes in a fifth of a second and splitting four ways costs more in dispatch and all-reduce than it saves. At 256 steps one chip cannot hold the batch, Captum has to chunk it into four sequential pieces, and four chips run those in parallel. 1.67x rather than 4x because per-chunk fixed costs and the all-reduce are not amortised at this size. The fixes on four chips. I rebuilt each branch from source on the host and ran the bazel targets upstream CI uses (_tpu-v6e variant) plus the matching pytest files against the installed wheel. The bounded scalar cache test passes and a four-worker torchrun of the leak check gives every rank the same counts as the single chip; the max_unpool, eight-op and eleven-op dynamic tests pass; the split-compile test passes; the loader test passes; the memory APIs return per-process values; the example runs in all three modes at 2.0e-7. Logs are under results/v6e-4/. What went wrong. Partway through, the host-side TPU session manager on the first v6e-4 started refusing every process (tpunetd START_SESSION failed). A guest reboot did not clear it, nor did restarting the runtime service. Queued-resource nodes cannot be stopped and started, so the only remedy was to delete the node and queue another with the same persistent build disk attached. The bazel cache on that disk took the replacement from bare image to its first branch test in under 20 minutes, which is the strongest argument I have for keeping build state on a separate disk. Not covered: the v5e-8 that upstream CI runs on, and tests/distributed/. v5e capacity was unavailable all day and a v5litepod-8 exceeds the project's 4-chip per-zone serving quota. Learned: multi-chip attribution is worth it only when a single chip runs out of memory. Build from source on the host, with the bazel cache on a detachable disk. Expect a TPU host to need replacing at some point and plan the setup so that costs 20 minutes rather than a day. --- 9. What went upstream To torch_tpu, on the explainability-example branch: - examples/explainability/captum_attribution.py: one script with flags for the model (tiny or any Hugging Face causal LM), steps, runs, precision, --compile, --internal-batch-size and --profile-dir. It reports CPU reference time, first-call time, warm median, max relative error against CPU and peak HBM. - docs/source/case_studies_explainability.py and its Markdown export: a seven-section case study covering running Captum on a TPU, setting matmul precision, comparing execution modes, reading the XProf trace, bounding memory, checking convergence and running the example. Linked from the case-study index and the examples README; mkdocs build --strict passes. To Captum, tests/attr/test_accelerator.py (meta-pytorch/captum#1951), which re-runs Captum's 78 attribution configurations on CPU and on any accelerator and compares them, so TPU compatibility is covered by tests rather than by this post. --- 10. If you want to try this - Strict eager and DeferAndFuse work on torch_tpu 0.2.0 as released. Compiled mode needs two unmerged patches (Captum #1952, torch_tpu #6106). Until they land, DeferAndFuse is the no-changes option. - Set matmul precision to "highest" and check the convergence delta. - Five branches (missing-ops-batch, explainability-example, dynamism-denylist-triage, docs-memory-apis, loader-import-error-context) are pushed to the fork without upstream PRs yet. The repository allows users without write access three open PRs and the first three fill the quota. Each is tracked by an issue (#6111 to #6114, and a comment on #4413) so review can start before the PR exists. --- 11. Where everything lives !Contributions map Open upstream - google-pytorch/torch_tpu#6102: bounded scalar buffer cache (fixes #6036) - google-pytorch/torch_tpu#6103: issue, aten::max_unpool2d missing - google-pytorch/torch_tpu#6105: max_unpool2d / max_unpool3d lowerings - google-pytorch/torch_tpu#6106: fix compiled TPU partitions reading CPU view arguments past their buffer - meta-pytorch/captum#1951: device-generic accelerator comparison tests - meta-pytorch/captum#1952: Python ints for LayerIntegratedGradients output offsets Pushed, PR pending (fork joinalahmed/torch_tpu) - missing-ops-batch: google-pytorch/torch_tpu#6111 - explainability-example: google-pytorch/torch_tpu#6114 - dynamism-denylist-triage: google-pytorch/torch_tpu#6112 - docs-memory-apis: comment on google-pytorch/torch_tpu#4413 - loader-import-error-context: google-pytorch/torch_tpu#6113 In this repository (joinalahmed/captum-on-torchtpu) - images/: the figures in this post, generated by scripts/make_blog_images.py - FINDINGS.md: the measurements, in the order they were taken - scripts/: every benchmark, probe and reproduction above, plus tpu_vm.sh for creating, setting up and deleting the VM - captum_on_torchtpu.ipynb: notebook version of sections 2 to 6 - reports/Explainability tools on TPUs.md: the prior-art survey - results/v6e-4/: the four-chip matrix, the sharded IG numbers and every branch's multi-chip test log - tpu_agent_project/: the same workload dispatched through Google's TPU agent suite, with the XProf report quoted in section 6 Cost. The single-chip work ran on a v6e-1 created and deleted per work session; the four-chip work on one v6e-4 from a queued resource after several hours of waiting. VMs and disks bill while they exist; everything was deleted when the work finished. --- Written with Claude Code as the pair programmer. Commits in torch_tpu carry an Assisted-by trailer, as that project's contributing guide requests; AI is not listed as a co-author there.
@joefioti BTW, blog.ezyang.com was a nice post on this from a while ago (and follow on blog.ezyang.com, meta-pytorch/spmd_types)