feat: TPU support via PyTorch/XLA (auto-detection, FSDP, XLA-safe merge/export)

[instax-dutta]

Summary

Adds TPU as a supported accelerator via PyTorch/XLA. Execution-only: every change is gated behind runtime TPU detection and is a no-op on CPU/CUDA/MPS - what a run computes (response length, prompt counts, trial counts, scorers, abliteration parameters) is untouched.

Validated end-to-end on a Kaggle TPU v5e-8 (torch 2.8.0+cpu / torch_xla 2.8.0) with a full 210-trial abliteration run of LiquidAI/LFM2.5-2.6B using the exact GPU recipe from reproduce/config.toml (same pinned model commit 403f92a, same seed 1552829245, train[:400] residuals, test[:100] eval, max_response_length=100, <think></think> prefix).

What's included

• system.py: torch_xla detection (detect_tpu), env-first SPMD setup for multi-chip FSDP (calling use_spmd() after XLA client init segfaults in PjRtComputationClient::ExecuteReplicated on torch_xla 2.8; setting XLA_USE_SPMD=1 before any device access avoids it), core-count resolution from PJRT env vars, mark_step(wait=True) so tensor device data is materialized before use, XLA branch in empty_cache, TPU accelerator info.
• fsdp_utils.py (new): FSDP wrapping across all physical cores via SPMD mesh for single-process hosts (Kaggle/Colab-style VMs), automatic layer-class discovery, output sharding for causal-LM heads.
• model.py: fit-based auto core selection (meta-device footprint estimate, fewest power-of-two cores that fit); CPU-load + FSDP shard before LoRA attach; bf16 forcing with quantization off; .eval() before any forward (dropout/RNG ops in the graph change the HLO per step and exhaust HBM); fixed-size sliding-window greedy decode loop _get_responses_xla; chunked get_responses_batched (20-prompt chunks are the measured HBM ceiling at 100-token responses; batching-only change, same responses); direct forward-pass paths for residuals/logits with mandatory CPU offload; FSDP-tolerant LoRA target collection; XLA-safe merge (materialize to CPU, strip FSDP _orig_module. key pollution, recompute LoRA targets against the CPU reload) so exported weights match base structure exactly; XM RNG reseed before randomized svd_lowrank.
• config.py: tpu_cores / tpu_use_fsdp / tpu_fsdp_config settings + adjust_for_tpu validator (bfloat16, disable bnb_4bit, resolve parallelism). Excluded from settings snapshots.
• main.py: early TPU environment setup before any device access.
• pyproject.toml: optional tpu extra (torch_xla).

Verification

On CPU/GPU paths the port is byte-for-byte identical to upstream: tests/run_tests.py produces identical model.safetensors hashes from a pristine upstream/master checkout and this branch (5/5 files compared on macOS/arm64). pytest tests/test_config.py passes.

TPU parity evidence (LFM2.5-2.6B, identical recipe):

	GPU reference (A100)	TPU v5e-8
Base refusals	98/100	84-89/100
Abliterated	4/100 @ KL 0.032	11/100 @ KL 0.142

The residual gap traces to bf16-XLA vs CUDA numerics: the baseline itself differs (84-89 vs 98) under the same commit/template, so borderline refusal decisions flip and the optimizer lands on a slightly different optimum. Same computation, different floating-point substrate.

Notes

• torch_xla 2.8 is the tested pairing (2.9 has an unrelated std::bad_alloc in from_pretrained on v5e).
• Single-core runs intentionally stay in plain (non-SPMD) mode: the SPMD virtual device breaks memory probing and accumulates null-data tensor fetches.

[instax-dutta]

On PR size

Flagging this upfront: the PR is large (+1037/-79 across 6 files). If you have a size/maintainability preference for how hardware support lands, I'm happy to rework it - including splitting or slimming.

Why it's this size:

• ~360 lines are comments/docstrings documenting torch_xla 2.8 failure modes that are invisible from the code alone (SPMD-after-client-init segfaults in PjRtComputationClient::ExecuteReplicated, HBM exhaustion at specific batch/chunk sizes, null-data tensor fetches after async mark_step). The intent is that future changes don't silently reintroduce them.
• Multi-core execution requires the FSDP path (fsdp_utils.py, 227 lines). Single-core plain-mode XLA turned out to be unusable on v5e-8 in practice (device-data fetches crash with null tensors within a few steps), so there is no smaller working configuration to ship first.
• The remainder is TPU-gated branches inside existing functions (decode loop, residuals/logits paths, merge/export). Heretic has no hardware-abstraction seam comparable to the scorer/plugin system, so the branches live inline - each one is if self._is_tpu: and a no-op on CPU/CUDA/MPS.

If incremental landing works better for you, it decomposes cleanly:

1. Detection + environment + config (~300 lines): detect_tpu, env-first SPMD setup, core-count resolution, settings fields.
2. FSDP multi-core execution (~350 lines): fsdp_utils.py + model loading/sharding + generation paths.
3. XLA-safe merge/export + chunked eval (~400 lines): CPU-reload merge with adapter-key fixes, chunked batching, offload paths.

Each stage is independently reviewable and leaves CPU/CUDA behavior byte-identical to master (verified: tests/run_tests.py produces identical weight hashes from a pristine checkout vs this branch).

Or if you'd prefer a different shape entirely - thinner integration, separate package, whatever fits the project - say the word and I'll adapt.