google-pytorch

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.

This project provides a PyTorch backend for Google's Tensor Processing Units (TPUs), enabling users to run PyTorch models and operations on TPU devices google-pytorch/torch_tpu

Google open-source TorchTPU, un back end PyTorch natif pour les TPU

Now Open Source: TorchTPU and the TPU Agent Suite for Native PyTorch on Google Cloud TPUs We’re excited to share that TorchTPU—our native PyTorch stack for Google Cloud TPUs—is now open source on GitHub, alongside the early preview of the TPU Agent Suite! Co-developed with the PyTorch team at Meta, TorchTPU brings Cloud TPUs (v5e, v5p, v6e, and v7 / Ironwood) alongside GPUs and CPUs as a native backend option for PyTorch. You can now train, fine-tune, and serve PyTorch models on TPUs with minimal code changes using the standard APIs and open-source tools you already rely on. ⚡ Quickstart: Installing TorchTPU You can install the latest torch_tpu release directly via pip (Python 3.11+): python -m venv .venv source .venv/bin/activate pip install torch_tpu --extra-index-url download.pytorch.org 🛠️ What’s Included in TorchTPU? ⚬ Native PyTorch Experience: Debug models effortlessly in eager mode (.to("tpu")) and transition to production-grade performance with torch.compile(), backed by custom ATen kernels and an XLA compilation cache. ⚬ Out-of-the-Box Ecosystem Integrations: ⚬ Hugging Face: Deploy popular open-weight models natively in Transformers v5 simply by passing device_map="tpu", plus native Diffusers support. ⚬ High-Throughput Serving: Native PyTorch backends for vLLM and SGLang, paired with llm-d for prefix-cache-aware Kubernetes fleet routing across GKE. ⚬ Distributed Training & RL at Scale: Production-grade integrations with TorchTitan (SimpleFSDP, custom Pallas kernels), VeRL for post-training/RL alignment, and Ray for heterogeneous compute orchestration. 🤖 Introducing the TPU Agent Suite (Early Preview) To make onboarding and performance tuning even faster, we’ve also open-sourced the TPU Agent Suite in preview. It provides declarative bundles of multi-agent skills, routing protocols, and tools that compile and install natively across coding agent harnesses (including Claude Code and Antigravity, with Codex and Cursor support coming soon). With the tpu-agent skill enabled, you can use natural language prompts directly in your coding environment to automate: ⚬ Model Onboarding & Benchmarking: Spin up vLLM inference servers on torch_tpu, run test completions, and benchmark TTFT, ITL, and throughput. ⚬ Profiling & Diagnostics: Analyze XProf profiles to pinpoint memory bottlenecks and TPU compute utilization. ⚬ Kernel Optimization: Optimize custom Pallas kernels and run automated self-refinement benchmark loops. 🔗 Resources & Links ⚬ TorchTPU GitHub Repository: google-pytorch/torch_tpu ⚬ TPU Agent Suite GitHub Repository: google-pytorch/tpu-agent-suite ⚬ Official Google Cloud Documentation: docs.cloud.google.com ⚬ Launch Blog Post: cloud.google.com We’d love for you to try out both projects! If you have questions about migrating a workload, feedback on the agent skills, or run into any issues, drop a reply below or open an issue directly on the GitHub trackers. What PyTorch workloads are you most excited to test on TPUs?

TorchTPU is now open source. You can run native PyTorch on Google TPUs. GitHub: google-pytorch/torch_tpu t.co

Google open-sources TorchTPU, a native PyTorch back end for TPUs View Article | Discussion + AI Summary Summary of HN discussion 🧵👇

Why did this take so long, right move regardless google-pytorch/torch_tpu

TorchTPU vs CUDA: Can Native PyTorch Make Google TPUs Easier to Use? A PyTorch model can be portable in the way a recipe is portable: the ingredients and steps look familiar, but the kitchen changes everything. Moving from an NVIDIA GPU to a Google TPU has traditionally meant learning a different execution path, debugging compiler behavior, and checking whether the libraries your model depends on actually work. Google's newly open-sourced TorchTPU is an attempt to shrink that gap. Announced on October 20, 2026, it brings a native PyTorch backend to Google's TPU hardware. That sounds like a direct alternative to CUDA, but the comparison needs care: CUDA is NVIDIA's broad GPU programming platform, while TorchTPU is Google's PyTorch-facing software stack for a specific accelerator family. The useful question is not whether one has "won". It is whether a PyTorch team should now consider TPUs alongside NVIDIA GPUs, and what migration still costs. The actual comparison: two paths from PyTorch to silicon CUDA is NVIDIA's programming platform and development environment for applications running on NVIDIA GPUs. Its toolkit includes libraries, compiler and runtime components, debugging and optimization tools, and deployment support across workstation, cloud, data-center, and embedded environments. PyTorch users often interact with CUDA indirectly through PyTorch and its libraries, but a deep ecosystem of CUDA-specific tools and kernels sits underneath that experience. NVIDIA describes CUDA as a platform and programming model, not just a device switch. TorchTPU is a PyTorch backend for Google's Tensor Processing Units. Its open-source repository describes custom ATen kernels, a compilation cache, distributed training support, and integration with torch.compile(). The project currently documents Python 3.11 or newer and offers a public package install. The repository README lists the project scope and installation path. Those are not identical layers. CUDA supports a large range of NVIDIA GPU software and native development workflows. TorchTPU is focused on running PyTorch workloads on Google TPU hardware. The comparison here is between the practical PyTorch experience each stack enables, not between two equivalent programming languages or chips. What changed with the open-source release Google says TorchTPU supports eager execution for debugging and torch.compile() for production optimization. The announcement also names integrations or collaborations across Hugging Face, vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, and LiteRT. In Hugging Face Transformers v5, Google says users can select a TPU with device_map="tpu". Those are vendor-stated integration claims in Google's release announcement. The important change is the interface. Developers can begin with familiar PyTorch APIs instead of first translating an entire model into another framework. That lowers the amount of code that must change before a TPU experiment can start. It does not mean every PyTorch model, custom operation, training recipe, or serving engine works unchanged. "Native PyTorch" describes the direction of the interface, not a universal compatibility guarantee. The repository provides a concrete signal of maturity and limits. It documents examples for models such as Llama, Qwen, ResNet, DLRM, minGPT, and UNet, plus distributed training and inference examples. A project example is evidence that a path exists, not proof that every model variant or workload is production-ready. Teams still need to check the exact operators, model versions, and performance characteristics they use. The current project README is the place to verify supported examples and install details. A small first experiment A TPU trial should answer one narrow question before it becomes a platform migration. For a supported Hugging Face model, the high-level shape shown by Google's announcement is simply to select the TPU device in Transformers v5: from transformers import AutoModelForCausalLM, AutoTokenizer model_id = "your-supported-model" tokenizer = AutoTokenizer.from_pretrained(model_id) model = AutoModelForCausalLM.from_pretrained( model_id, device_map="tpu", ) Treat this as an illustration of the API direction, not a complete turnkey deployment recipe. The model must be supported by the TPU backend and your Transformers version; the runtime and accelerator need to be provisioned and configured; and you should validate output correctness and memory behavior before timing it. Google's announcement says the device_map="tpu" path is available with Transformers v5, while the TorchTPU repository has separate setup requirements. Read both before using this snippet as a runbook. Google's announcement and the TorchTPU repository are the authoritative starting points. Keep the first test intentionally boring: a model and input shape you already run, a fixed input set, and a known-good output comparison. Record initialization time, steady-state throughput, memory, compilation or warm-up behavior, and the cost of the TPU VM configuration. Do not compare one warm GPU run against a cold TPU run, or compare different precision and batch settings. That would measure the test setup as much as the hardware. Where CUDA still has the advantage A mature development surface. CUDA's official guide spans the programming model, C++ and Python programming, advanced features, hardware-specific capabilities, and technical reference material. The toolkit also includes math libraries, profiling, debugging, compiler, and runtime components. This is a deep surface area for teams that need to tune kernels or depend on GPU-specific libraries. NVIDIA's current CUDA documentation shows that breadth. A familiar baseline across many environments. CUDA is available across desktop workstations, enterprise data centers, cloud platforms, and embedded systems, according to NVIDIA's toolkit documentation. That does not mean all NVIDIA GPUs are equally available or cost-effective, but it gives teams multiple deployment contexts to consider. The CUDA documentation lists these deployment targets. Existing investments. A project that already relies on CUDA extensions, custom kernels, a specific inference engine, or an established GPU operations stack may face substantial porting work. TorchTPU's promised PyTorch compatibility does not automatically translate a CUDA extension into an equivalent TPU implementation. The relevant test is whether the exact library and operators are supported, and whether the needed custom path exists. An established operational playbook. Teams may already have dashboards, profilers, container images, runbooks, and incident knowledge built around NVIDIA GPUs. Replacing hardware changes more than the model's device argument. It affects capacity planning, deployment, debugging, and on-call response. TorchTPU can reduce application-level friction while leaving infrastructure and operations differences intact. Where TorchTPU may be compelling A lower-friction PyTorch entry point. The promise is that familiar PyTorch code can target TPU with fewer changes, and Google's release describes eager mode plus torch.compile() as part of the workflow. For teams that have wanted to evaluate TPU but avoided a framework detour, that could make a small controlled experiment practical. Google's description is a vendor claim; verify it against your own workload. A growing set of upstream-adjacent integrations. The announced work with Hugging Face and open-source serving and training projects could matter more than a one-line device API. A backend becomes useful when the libraries around model loading, distributed execution, serving, monitoring, and orchestration support it. The announcement names several integrations, but the exact feature set can differ by project and version. Verify each project's own documentation before designing around an integration. A different accelerator option for Google Cloud users. TPUs are Google-designed accelerators, and Google Cloud's documentation describes different TPU versions, capacity options, and quota requirements. Hardware access is therefore part of the decision: you need suitable TPU capacity in a supported region and project, not just a local GPU machine. Google's TPU planning guide covers on-demand, Spot, Flex-start, and reservation choices, while the pricing page shows that prices vary by TPU version, location, and purchase model. That last point is a useful reality check. Open-source software does not make the accelerator free, guarantee quota, or ensure that a specific region has the capacity you need. Budget both the machine and the engineering time spent validating compatibility. The hidden work: compilation, support, and portability A familiar API can hide a different execution model. The PyTorch/XLA documentation now labels that project deprecated, says TorchTPU replaces it, and asks teams to start new PyTorch TPU workloads with TorchTPU and migrate existing ones. It also says the PyTorch/XLA repository will be archived on November 15, 2026. For teams already using PyTorch/XLA, this is a concrete migration timeline to account for, not a claim that the older path stops working today. Profile compilation, runtime, and supported operations in the stack you actually deploy. The PyTorch/XLA notice gives the migration direction and archive date. Before a real migration, inspect the workload for dynamic shapes, custom operators, unsupported kernels, frequent branching, or CPU fallbacks. Google's general TPU guidance says TPUs are not suited to linear-algebra workloads with frequent branching or many element-wise operations, workloads requiring high-precision arithmetic, or neural-network training loops with custom operations. That guidance is about TPU workload suitability generally, not a TorchTPU-specific compatibility matrix. Google's TPU overview documents these workload considerations. Then profile the entire path. Include data loading, host-to-device transfers, compilation and warm-up, batch formation, communication between devices, and serving overhead. A benchmark that times only the model's core operation can make an accelerator look better than the full application will. A practical decision checklist Use this checklist before moving a production workload: - Start with a concrete reason. Is the goal capacity, price-performance, access to a different accelerator pool, or reduced dependence on one vendor? "TPUs are new" is not a business requirement. - Inventory the software path. List PyTorch, Transformers, serving libraries, custom operators, precision settings, and distributed-training features. Confirm support for the exact versions you use. - Choose one representative workload. Include the real sequence lengths, batch sizes, input distribution, and serving pattern. A toy example tests installation, not production fit. - Check numerical correctness. Compare outputs and task-level quality within an acceptable tolerance before comparing speed. - Measure cold and warm behavior separately. Capture compile or initialization costs, steady-state throughput, latency percentiles, memory use, and restart behavior. - Calculate end-to-end cost. Include accelerator time, host VM cost, storage, network transfer, idle time, and engineering maintenance. Verify regional price and capacity at test time. - Test failure and scale paths. Confirm how jobs recover, how workloads distribute, and what happens when a TPU resource is unavailable or preempted. - Keep a rollback route. Preserve the CUDA deployment until the TPU version has passed a representative evaluation and an operational review. Verdict: experiment with TorchTPU, do not assume CUDA portability TorchTPU makes the TPU question more approachable for PyTorch teams. Its strongest case is not that it has already displaced CUDA. It is that a native PyTorch backend and announced ecosystem integrations may reduce the cost of evaluating Google's accelerators without rewriting a project around a different model framework. CUDA remains the safer default when a workload depends on NVIDIA-specific code, established GPU operations, or broad deployment choices. TorchTPU is worth a measured pilot when the team can access Google TPU capacity, the model and operators are supported, and there is a plausible infrastructure reason to test it. The choice should come from a workload benchmark and a full operating-cost calculation, not from launch-day performance claims. This is a source-based comparison, not a report of hands-on benchmarking. Google's release announcement includes performance and scale claims, but I have not independently reproduced them here. Treat the announcement as a map of what to test, then verify the result with your own model, configuration, and bill. I write about AI infrastructure and developer tools. Subscribe if you want practical comparisons and implementation guides. Would you test a second accelerator backend for your PyTorch workloads? What would decide it for you: model compatibility, measured cost, performance, or operational complexity? Sources: Google Cloud: TorchTPU is now open source, TorchTPU repository, NVIDIA CUDA Programming Guide, NVIDIA CUDA Toolkit Documentation, Google Cloud TPU planning guide, Google Cloud TPU pricing, Google Cloud TPU workload overview, PyTorch/XLA documentation

Engenheiro da PyTorch para TPUs no Google, anuncia a disponibilização do código-fonte do TorchTPU. Isso permite que códigos PyTorch sejam executados em TPUs sem a necessidade de grandes reescritas, utilizando APIs padrão e mapeamento de dispositivos. google-pytorch/torch_tpu

Meta「Muse」の共有リンク、会話が誰でも見つかる状態に 2026年10月21日（水）のAIニュース完全解説版。図解を見れば本文を読まなくても分かる形にしています。 今日はちょっと怖い話から入ります。MetaのAIで、相手にだけ送ったつもりの会話が、誰でも見つけられる状態になってたんです。そのあと、飲食店がAIで経理の時間を94%減らした話、AIの値段を左右する部品の競争、AIの答えは裏から動かせるかもしれないという指摘、Anthropicの巨額の赤字の報道と続きます。全部で5本。どれも皆さんのお金と信用に関わる話なんですよ。 1. Meta「Muse」、共有した会話が検索1回で見つかる ちょっと怖い話からいきますね。MetaのAI「Muse」には、会話をURLで人に送れる共有機能があったんです。URLは長い文字列で、普通は推測できない作り。ところがセキュリティ会社Clarocの研究者が調べたら、共有リンクが6,800件も見つかった、とArs Technicaが報じてます。研究者いわく「検索1回で済んだ」。 中身を数百件開いたら、全国展開のマーケティング会社がCEO向けに作った売上の資料、小さな会社の融資の申込書、職場の1年分の休暇の予定表。どう見ても公開するつもりのないものなんですよ。Metaは共有機能をいったん止めて見直す、と答えています。 ここが大事で、種明かしすると、URLが複雑でも、そのアドレスの記録を集める商用サービスから拾えたんです。金庫の暗証番号は複雑なのに、番号のメモが掲示板に貼ってあった、みたいな話でしょ。 共有ボタンを押した瞬間、それは誰かが見つけられるページになる。 経営者の皆さんは、仲間がAIとの会話を共有リンクで回していないか、一度聞いてみてください。個人で学んでる人は、共有する前に名前や金額が入ってないか見直す。僕はこれだけでかなり防げると思ってます。 出典: arstechnica.com 2. 飲食店オーナー、AIで現金の照合を94%短縮 ここなんですよ。AIの使い道って、つい派手なものを想像しちゃうじゃないですか。でも今日は、すごく地味な話です。 OpenAIが、ある飲食店の事例を公開しました。Schlotzsky'sとCinnabonのお店を独立して経営しているフランチャイズのオーナーさんが、「ChatGPT Work」で現金の照合にかかる時間を94%減らした、と出てるんです。メニューの価格を見直す時間も95%減。作った作業の流れで、$100Kの資金を生み出せる見込みだそうです。ただし、これはOpenAI自身の事例紹介で、$100Kも見込みの数字です。 これ、めちゃくちゃ大事で。毎月のレジ締めの突き合わせって、テストの丸付けみたいなものでしょ。正解はあるのに、時間だけ取られる。そういう仕事こそ、AIが得意なところなんですよ。 AIで一番お金になるのは、毎月くり返している面倒な作業。 経営者の皆さんは、毎月誰かが半日かけている数字の作業を1つ挙げてみてください。そこが最初の候補です。個人で学んでる人は、家計簿や副業の売上の集計で小さく試す。僕も自分の会社の仕組みをClaude Codeで作って毎日回してますけど、手応えがあるのはやっぱりこういうくり返しの作業なんです。 出典: openai.com 3. Google、AI開発の定番道具をNvidia以外でも動く形に 先に結論を言います。AIを動かす「部品」の競争が、一歩進みました。 Googleが、TorchTPUというものを無料で公開したんです。PyTorchは、AIを作るときに世界中の開発者が使っている定番の道具。これまでは、Nvidiaの半導体で動かすのが当たり前だったんですよ。TorchTPUは、それをGoogle自前の半導体「TPU」でそのまま動かせるようにするもの。TechCrunchによると、Amazonも同じ方向のTorchNeuronを先に出していて、Googleがそこに加わった形です。Hugging Faceも対応を発表してます。 ここが大事で、例えるなら、1社のコンセントにしか挿せなかった家電に、変換プラグが配られ始めた感じです。プラグがあれば、電気代の安い方を選べるじゃないですか。選べる部品が増えれば値段の競争が起きて、AIの利用料が下がる可能性がある。ただし、すぐ値下げとは限らないし、Nvidiaがいなくなるわけでもないです。 一社頼みが崩れると、得をするのは使う側。 経営者の皆さんは、AIの費用はこれから下がり得る、という前提で、長く縛られる契約は慎重に考えてください。個人で学んでる人は、今すぐ触る話ではないです。僕は毎朝AIのニュースを自分の仕組みで集めて読んでますけど、AIの値段は裏側の部品で決まってると知っておくだけで、読み方が変わりますよ。 出典: google-pytorch/torch_tpu 4. AIの答えは裏から動かせる？ネットで話題の2つの指摘 誤解しないでほしいんですけど、これはまだ「指摘の段階」の話です。確信度は中くらいとして聞いてください。 1つ目。WhatsAppのヘルプページに、人には見えない形でAI向けの指示が隠されていて、中身は「規制当局を勧めないで」という趣旨だった、とredact.devのブログが指摘しました。Hacker Newsで708ポイント集めてます。2つ目は、peec.aiのブログの「偽の統計を載せれば、AIに引用される」という話で、こっちは304ポイントです。 ここが大事で、根っこは同じなんですよ。AIはネットの文章を読んで答えている。だから書く側が、AIの答えを動かせてしまうかもしれない。お店の口コミが、実は自作自演だった、みたいな話でしょ。僕は20歳のとき甘い投資話に乗って痛い目を見たので、「うまい話ほど裏を取る」が体に染みてます。 AIに拾われたいなら、偽物じゃなく本物の数字を出す。 経営者の皆さんは、偽の数字で目立つのは短期の得、信用を失うのは長期の損だと考えてください。自社の実績みたいな本物のデータを出す方が強いです。個人で学んでる人は、AIの答えに出典が出たら1回開いて確かめるクセをつける。 出典: redact.dev / peec.ai 5. Anthropic、上半期の赤字が$80B超と報道 図解（完全解説版） ここが大事です。ClaudeというAIを作っているAnthropicが、2026年の上半期に$80Bを超える純損失を出した、とWall Street Journalが報じ、ロイターが伝えています。会社の正式な発表ではないので、確信度は中で聞いてください。TechCrunchによると、11月の上場を予定するなか設備への支出が膨らみ、AIを動かす計算資源の確保に苦労しているそうです。クラウドの会社が、Anthropicを競合として見始めている、と。 $80Bですよ。たった半年で。AIの会社はいま、ものすごい額を先に払って設備を整えてるんですよ。開店セールみたいなものですよね。利益を削ってでも安く出す。でもセールって、いつかは終わるじゃないですか。赤字だからすぐ危ない、とまでは言えません。ただ、今の料金や使える量がずっと続くとは思わない方がいいです。 安さがずっと続く前提で、仕組みを一社に固めない。 経営者の皆さんは、AIを業務に入れるとき、別のAIに乗り換える手順もセットで決めておくと安心です。個人で学んでる人は、2つのAIを触れる状態を作っておくと強い。僕も毎日AIで会社の仕組みを回しているので、ここは他人事じゃないんです。 出典: reuters.com まとめ 今日の5本、根っこは同じなんですよ。AIは便利な分だけ、使う側の扱い方で結果が変わる。共有ボタンを気軽に押したり、安さに乗って一社に固めたりするのは、短期では楽です。でもそれが、長期の信用と自由を削っていく。逆に、地味な作業から丁寧に任せていけば、時間もお金もちゃんと手元に残ります。短期の「楽」より、長期の「強さ」を選んでいきましょう。 この記事は、公式ブログ・研究者の発信・掲示板（Hacker News など）・X の一次情報を毎朝集めて、非エンジニアにも分かる形にまとめています。 今日読んだ全件リストは別途保管しています。