fix(training): select XLAStrategy for multi-device XLA/TPU training - #1427

Summary

Asking for more than one TPU chip fails before any accelerator work happens:

ValueError: The `XLAAccelerator` can only be used with a `SingleDeviceXLAStrategy`
or `XLAStrategy`, found DDPStrategy.

XLAAccelerator pairs only with SingleDeviceXLAStrategy or XLAStrategy. But build_trainer's own strategy branch turns strategy="auto" into DDPStrategy(find_unused_parameters=True) as soon as more than one device is requested:

elif strategy_name == "ddp" or (strategy_name == "auto" and distributed_requested):
    ...
    strategy = _DDPStrategy(find_unused_parameters=True)

TrainConfig.devices defaults to 1, a default TPU run never reaches that branch and works fine. But if you ask for the whole slice, the run dies during Trainer construction. Since it fails before touching a chip, it looks like a configuration mistake and not a missing feature, and the one-word workaround (strategy="xla") is nowhere in the error message.

This change selects "xla" when the accelerator resolves to XLA/TPU, the strategy is still "auto", and more than one device or node is requested. It's the same kind of gap precision already needed handling for on this path: XLA needs its own object, the generic default doesn't work.

Two more paths hit the same crash and need the same fix:

• accelerator="auto", which is TrainConfig.accelerator's own default. build_trainer decides the strategy before Trainer runs its own accelerator resolution, so checking only for the literal strings "xla"/"tpu" misses the case where the caller never names an accelerator at all — this is the documented default on purpose (TrainConfig.accelerator's own docstring: "device is intentionally absent: PTL auto-detects accelerator via Trainer(accelerator="auto")"). xla_accelerator is now computed by a helper, _accelerator_resolves_to_xla, that also checks XLAAccelerator.is_available() for "auto" — it follows the same priority order (XLAAccelerator checked first) that lightning_fabric.utilities.device_parser._select_auto_accelerator uses.
• num_nodes > 1 with a single device per host (a multi-host XLA topology). The guard's own multi-device check, _requests_multiple_devices(devices, accelerator), only looked at devices. The pre-existing distributed_requested computation a few lines below already treats num_nodes > 1 as distributed too, so now the guard checks the same condition.

Scope

• Single-device, single-node XLA strategy selection stays the same — "auto" never reaches the DDP branch there, so it still resolves to SingleDeviceXLAStrategy. The precision-plugin branch (XLAPrecision vs. a raw precision= string) now also activates for accelerator="auto" resolving to XLA on a single device, because it reads the same xla_accelerator value as the strategy guard. This was already needed (XLAStrategy/SingleDeviceXLAStrategy reject a raw precision string) and gives the same resolved precision value either way.
• An explicitly passed strategy is never overridden. An explicit "ddp" still resolves to the same DDPStrategy object as before — that strategy_name == "ddp" branch already existed, is unrelated to this change, and fires the same with or without XLA.
• Non-XLA accelerators don't change: the guard requires xla_accelerator, and _accelerator_resolves_to_xla returns False for any accelerator other than "xla"/"tpu"/an XLA-resolving "auto".
• Keypoint models are excluded from this fix. They train under manual optimization, and the find_unused_parameters=True handling that the DDP branch relies on has no validated XLA equivalent. So use_grouppose_keypoints=True configs still hit the pre-existing DDPStrategy/XLAAccelerator mismatch, instead of silently running a combination nobody validated.
• Segmentation models are not excluded — the guard doesn't special-case segmentation_head, so a segmentation config reaches "xla" the same way plain detection does.

Validation

TestAcceleratorResolvesToXLA unit-tests _accelerator_resolves_to_xla directly (mocking XLAAccelerator.is_available, no torch_xla needed): explicit "tpu"/"xla" (case-insensitive) are always XLA; "auto" matches whatever is_available() reports; "cpu"/"gpu" are never XLA no matter what. 7/7 passed.

TestMultiDeviceXLAStrategy covers build_trainer's actual strategy dispatch. The first four cases need real torch_xla (marked xla, importorskip-guarded) because build_trainer builds a real XLAPrecision on this path. The newer cases patch XLAPrecision instead (and, for accelerator="auto", XLAAccelerator.is_available), the same way TestBuildTrainerPrecision's existing XLA tests already do, so they run on every CI lane:

Reverting only the source fix (the whole if (...): strategy = "xla" block disabled, tests kept) turns all three into failures — strategy stays a DDPStrategy object instead of "xla". So each test is checking the guard it targets, not passing by accident.

On a Cloud TPU v6e-1 (southamerica-west1-a, runtime v2-alpha-tpuv6e), torch 2.9.1, torch_xla 2.9.0, libtpu 0.0.21, pytorch-lightning 2.6.5, against this base commit, with accelerator="tpu" passed explicitly:

That last row is a different mutation on purpose. Removing the entire guard leaves keypoint configs on the pre-existing DDP branch anyway, so that assertion would hold either way and wouldn't tell us anything about the keypoint exclusion specifically. Removing only the has_keypoints check isolates it instead. Both mutations ran locally against the CPU-mocked XLAPrecision path (_mc(use_grouppose_keypoints=True) etc.) with the same result, then got re-verified with real torch_xla on the arm shown above.

On a Cloud TPU v6e-4 (4 chips, xla:0-xla:3), same base commit, accelerator="tpu" and devices=4 passed explicitly (not the accelerator="auto" default — that path is only covered by the CPU-mocked test above, not yet by real hardware):

The unpatched run prints rf-detr's own line right before it fails, and that's where the DDPStrategy comes from:

[INFO] rf-detr - strategy='auto' with distributed execution → DDPStrategy(find_unused_parameters=True).

test_all_gather_multiprocess_xla_collective_routing also passes on that slice in 18.11 s, so the four chips are doing real multi-replica work.

Training end to end across replicas was measured separately on 2026-08-24, on an earlier develop commit, on a v5litepod-4, with strategy="xla" passed explicitly (the pre-existing workaround, not touched by this change — not the accelerator="auto" auto-detection path this PR adds): the same configuration went from failing in 6.3 s to completing 2 epochs over 468 images in 919.2 s. This PR makes no multi-chip throughput claim — on that small dataset four chips weren't faster than one, since each replica pays its own warm-up.

Still not verified on real TPU silicon: accelerator="auto" resolving to XLA, num_nodes > 1, and a segmentation config reaching the guard. All three are covered by the CPU-mocked tests above, but not by a real multi-chip/multi-host run.

Refs #1058.