TPU Training Support

Summary

RF-DETR can be made trainable on Google TPUs (v3, v5e, v6e) with a small set of targeted fixes. PyTorch Lightning's XLAStrategy handles all TPU infrastructure automatically — distributed init, data loading, checkpointing. The remaining work is making the RF-DETR model code XLA-compatible and fixing two known performance bottlenecks.

All changes are device-gated — no GPU or CPU behaviour changes.

Motivation

• Free TPU quota on Colab (v5e-1, v6e-1) and Kaggle (v3-8) makes TPU training accessible to the community without GPU costs.
• TPU v5e/v6e offer significantly higher bf16 throughput than comparable GPU instances.
• Several community members have requested TPU support (see #XXX).

Required Changes

Phase 0 — Test infrastructure (prerequisite)

Add @pytest.mark.tpu pytest marker and update CI filters.

• Register marker in pyproject.toml alongside existing gpu and flaky markers.
• Update CPU CI filter from -m "not gpu" to -m "not gpu and not tpu" so TPU tests are skipped in standard CI but remain runnable explicitly with -m tpu.

Phase 1 — XLA compatibility (7 tasks)

Tasks 1.1, 1.2, 1.3, 1.6, 1.7 are independent and can be developed in parallel.

Task 1.1 — XLA branch in _bilinear_grid_sample (tensors.py:228)

F.grid_sample backward is a known CPU-fallback op on XLA. The gather-based path already used for MPS is fully XLA-compatible. One-line fix:

# before
if input.device.type != "mps":
# after
if input.device.type not in ("mps", "xla"):

Task 1.2 — Remove .item() from num_boxes (criterion.py:497)

.item() forces a device-to-host sync, cutting the XLA lazy graph on every training step. Scalar tensor division works identically on CPU, CUDA, and XLA.

# before
num_boxes = torch.clamp(num_boxes / get_world_size(), min=1).item()
# after
num_boxes = torch.clamp(num_boxes / get_world_size(), min=1)  # keep as tensor

Note: dice_loss and sigmoid_ce_loss are @torch.jit.script functions with num_masks: float signatures — these must be updated to num_masks: torch.Tensor in the same PR.

Task 1.3 — Fix all_gather hardcoded device probe (distributed.py:79) ⚠ BLOCKER

all_gather hardcodes device = torch.device("cuda" if torch.cuda.is_available() else "cpu"). On TPU this creates gather buffers on CPU, causing a hang or silent corruption in multi-replica training. Fix: derive device from the distributed backend instead.

def all_gather(data: Any, device: torch.device | None = None) -> list[Any]:
    ...
    if device is None:
        backend = torch.distributed.get_backend() if is_dist_avail_and_initialized() else "cpu"
        device = torch.device("cuda" if backend == "nccl" else "cpu")

This also fixes the same latent bug for CPU-only DDP and other non-CUDA distributed backends. dist.all_reduce in criterion.py does not need changing — ProcessGroupXla already delegates it to xm.all_reduce() automatically.

Task 1.4 — Validate NestedTensor bool mask on XLA (module_data.py:489)

XLA materialises torch.bool tensors as uint8 internally. The bool mask is used only in shape-preserving ops (mask.float(), ~mask, masked_fill) so correctness risk is low, but needs explicit smoke-test validation. If dtype changes on transfer, a .to(torch.bool) cast is the fix.

Task 1.5 — Multi-scale: static-shape recommendation + compile note (module_model.py)

multi_scale=True triggers F.interpolate with varying (H, W) on each step. Every distinct shape causes XLA graph recompilation (~60–90 s each), adding 8–15 minutes of warm-up overhead on the first epoch.

• multi_scale=False (recommended for TPU): all images resized to a single fixed resolution — zero recompilations after the first batch.
• do_random_resize_via_padding=True: skips F.interpolate but the collate function still pads to the per-batch max size, so some recompilation remains.
• For multi-scale TPU training, shape buckets with block_size rounding (already supported in collate_fn) can cap the number of distinct shapes to a small finite set.

Recommended first TPU recipe:

model.train(
    accelerator="tpu",
    multi_scale=False,
    amp=True,
    ...
)

Task 1.6 — XLAPrecision plugin in build_trainer (trainer.py)

PTL's XLAStrategy rejects any standard precision= string and raises:

TypeError: The XLA strategy can only work with the XLAPrecision plugin

For TPU, omit the precision kwarg and attach XLAPrecision("bf16-true") via plugins instead:

if tc.accelerator in ("xla", "tpu"):
    from lightning.pytorch.plugins import XLAPrecision

    plugins = list(trainer_kwargs.get("plugins", [])) + [XLAPrecision("bf16-true")]
    trainer_kwargs["plugins"] = plugins
    trainer_kwargs.pop("precision", None)

Task 1.7 — Guard explicit augmentation_backend="gpu" on XLA (module_data.py)

augmentation_backend="auto" already avoids Kornia on non-CUDA hardware via _has_cuda_device(). However, an explicit augmentation_backend="gpu" bypasses this check and would select Kornia on TPU. One-line guard:

if backend == "gpu" and not _has_cuda_device():
    return "cpu"

Phase 2 — Smoke tests

Phase 2a — Colab TPU v5e-1 / v6e-1 (single chip, free)

Minimal checklist:

1. XLA device active (xla:0)
2. MXU utilization > 5 TFLOPS on a 2048×2048 matmul
3. RF-DETR runs 2 training steps (multi_scale=False, max_steps=2)
4. Loss finite after step 1
5. samples.mask.dtype == torch.bool inside training_step
6. Zero aten:: fallback ops in XLA metrics report
7. No grid_sampler_2d_backward CPU fallback in output

Phase 2b — Kaggle TPU v3-8 (8 chips, free)

Validates Task 1.3 in real multi-replica context:

1. 8 replicas initialised (WORLD_SIZE = 8)
2. Training completes 2 steps without distributed hang (< 5 min)
3. Loss consistent across ranks
4. No NCCL/Gloo backend error
5. Rank-0 checkpoint saved

Phase 3 — Convergence validation

Train rfdetr-small on COCO-2017 for 10 epochs on GPU (A100, seed=42) and TPU v3-8 (seed=42, multi_scale=False). Pass: val/mAP_50_95 within ±0.5 AP after epoch 10.

Also: run 100 GPU steps before/after Task 1.2 and assert train/loss identical within atol=1e-5.

Optional — Performance (post-launch)

CI Strategy

TPU hardware is not available in standard CI. Options for ongoing TPU validation:

Recommendation: Option A for v1. Revisit per-PR gate after initial implementation proves stable.

Risk Register

Acceptance Criteria

• @pytest.mark.tpu registered; CPU CI filter updated to not gpu and not tpu
• All 7 Phase 1 tasks merged; CPU/GPU CI 100% green
• Phase 2a Colab smoke: all 7 checklist items pass
• Phase 2b Kaggle multi-replica smoke: all 5 checklist items pass
• num_boxes GPU regression: loss identical within atol=1e-5
• COCO 10-epoch mAP: TPU within ±0.5 AP of GPU reference

TPU training support is not considered complete until all acceptance criteria pass.