fix: answer the bfloat16 dtype question without allocating

[lstein]

Summary

TorchDevice.choose_bfloat16_safe_dtype decided which dtype to use by building a one-element bfloat16 tensor on the target device. That put a real allocation between a caller and its dtype: on a loaded GPU the allocation fails, and since only TypeError was caught, the failure escaped as a RuntimeError from a function that had been asked a question about dtypes, before the caller had requested any memory of its own.

This is what turned a memory-starved macOS CI runner red:

RuntimeError: MPS backend out of memory (MPS allocated: 0 bytes, other allocations: 0 bytes,
max allowed: 7.93 GB). Tried to allocate 256 bytes on shared pool.

The probe now builds a zero-element tensor. It gets the same verdict from the same check, and asks the allocator for nothing:

• torch gates bfloat16 on the macOS version inside empty_mps with a TORCH_CHECK_TYPE on the dtype argument, above the allocator->allocate() call and with no reference to the element count — so TypeError still fires on macOS 13 and the float32 fallback is unchanged (verified in PyTorch v2.7.1 and v2.4.0 source, at the version pyproject.toml pins for darwin).
• MPSAllocator::allocate is nbytes > 0 ? _getAllocImpl().malloc(...) : nullptr, so zero bytes never reaches Metal.
• Measured here on ROCm: torch.empty(0, …) allocates 0 bytes and its data_ptr is null, against 512 bytes plus a 2 MiB reserved segment for the old probe. It is also ~30x faster per call on GPU, which matters on a loader path that calls it once per component.

The same allocating probe appears verbatim at three FLUX VAE loader sites (flux.py:158, :184, :284) — on a load path, where VRAM is tightest — and is fixed the same way, so the repo answers this question one way.

No dtype changes on any device: bfloat16 stays bfloat16 on CUDA/ROCm, CPU, MPS on macOS 14+ and XPU; macOS 13 still gets float32.

Related Issues / Discussions

Companion to #309, which reduces the memory pressure on the macOS runner. This PR removes the dependence on that memory in the first place; #309 is still worth having, since 6.70GB of workers on a 7GB runner has no headroom for anything else.

Merging this into main fixes the failing macOS jobs on the open LTX-2 PRs (#305 and the rest of the stack) with no change to their tests.

QA Instructions

• pytest tests/backend/util/ tests/backend/model_manager/load/ … — 1423 passed, 150 skipped, 6 xfailed. ruff check and format clean.
• The CI failure reproduces locally and the fix clears it. Driving the failing LTX-2 test from LTX-2: keyframe conditioning and video extend #305 against a device whose allocator refuses every non-empty request (what the runner did) gives the exact Tried to allocate 256 bytes on shared pool error on main's probe, and passes with this change applied and nothing else. The test itself is untouched.
• Mutation-tested the new tests against four defects: the old one-element probe, a probe that ignores the device it was handed, a genuinely allocating probe, and an implementation that swallows RuntimeError. Each is caught by at least one of the new tests.
• Not verifiable here: no MPS device on this machine (Linux/ROCm), so the macOS 13 → float32 path is established from PyTorch's source at the pinned tags, not by execution. The macOS jobs on this PR exercise the macOS 14+ path only.

Review

Two independent read-only reviews, both adversarial. Both found material problems in the first draft, which the diff now reflects:

• An earlier draft also caught RuntimeError and returned bfloat16. That was wrong twice over: once the probe stopped allocating, the branch could no longer catch the OOM it was written for, and what it could still catch was a device that cannot be reached at all — an out-of-range ordinal, or a backend this build lacks. Those used to raise immediately and by name; the draft turned them into a plausible dtype and a failure much later inside a loader. The branch is gone; a test now pins that an unreachable device still raises.
• The tests over-pinned the implementation and under-pinned the contract. One asserted that the backend must be asked at all, which would break a capability-API refactor that preserves every answer; none asserted that the probe targets the device it was handed — a probe that always asked the CPU passed the whole suite. Both fixed.
• The TypeError fallback test gained an xpu case: that arm is the else of the CUDA branch, so it covers XPU too.

Known limitations, stated rather than papered over: the CUDA → float16 arm is preserved but is not known to be reachable (torch does not reject bfloat16 as a storage dtype on any CUDA arch); the new tests capture the constructors a probe would plausibly use, so a probe spelled some other way records nothing rather than failing. Pre-existing and out of scope: test_devices.py's choose_anima_inference_dtype tests set config.precision on the global singleton without restoring it, which makes two unrelated tests in that file order-dependent.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Meaningful regression coverage added / updated where needed; obsolete tests/code removed
• Persisted-state and API changes include required migrations / compatibility validation
• Relevant performance/efficiency opportunities considered; material claims have evidence
• Material review findings resolved and relevant checks rerun
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

🤖 Generated with Claude Code