Finish removing the MPS autocast workaround (#49157)
#49006 removed the obsolete `device.type != "mps" -> "cpu"` fallback so RoPE
stays float32 inside MPS autocast contexts, but it replaced one exact spelling,
so 24 call sites that feed an autocast were left behind in 21 files. They are
the same workaround written differently:
- without the `isinstance` guard (`esmfold`, `falcon`, `gpt_neo`, `gptj`,
`hy_v4`, `nemotron`),
- spread over several lines (`ernie4_5_moe`, `ernie4_5_vl_moe`, `neucodec`,
`nemotron_asr_streaming`, `parakeet`, `xcodec2`),
- with the exact swept spelling but another variable name (`pi0` uses `time`,
the OpenVINO exporter uses `x`; the exporter landed about an hour after
#49006 merged).
Measured on an M-series machine with torch 2.14, replicating the `hy_v4`
block inside `torch.autocast("mps", dtype=torch.float16)`:
device_type resulting context dtype
"cpu" (before) nullcontext torch.float16
"mps" (after) autocast torch.float32
`maybe_autocast` returns `nullcontext()` when `torch.is_autocast_enabled(
device_type)` is False and `enabled` is False, so passing `"cpu"` while running
on MPS makes the block a no-op: the comment above it reads "Key difference is
to force fp32 in any case", and on MPS it forced nothing. The same holds for
`exporter_openvino.py`, which calls `torch.autocast` directly: it disabled CPU
autocast while MPS autocast stayed on. In `gptj` and `gpt_neo` the consequence
differs -- those query `torch.is_autocast_enabled(device_type)` to pick a
target dtype, so on MPS they queried the wrong device.
Deliberately not touched:
- the nine `compute_dtype = torch.float64 if ... else torch.float32` sites
(`aimv2`, `d_fine`, `deimv2`, `pp_doclayout_v2/v3`, `rt_detr`, `rt_detr_v2`,
`vit_mae`), which work around MPS lacking float64 and are unrelated to
autocast,
- `generation/utils.py:437`, which is device-specific generation logic.
Both `modular_*.py` sources and their generated `modeling_*.py` files are
updated, and every `device_type` line in a modular file appears verbatim in its
generated counterpart.
`examples/modular-transformers/modeling_super.py` also carries the old spelling
but is left alone: its `modular_super.py` source does not contain the line, so
it comes from inherited code and a manual edit there would be overwritten by the
next regeneration.
No tests, matching #49006: the behaviour needs an MPS device and CI has no MPS
runner. The measurement above is reproducible on any Apple Silicon machine.
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
