Segfault on macOS (MPS) when loading a bf16 checkpoint with dtype=torch.float32 and device_map="auto" (async loading)

[sergiopaniego]

System Info

• transformers version: 5.15.0
• Platform: macOS-26.3.1-arm64-arm-64bit-Mach-O
• Python version: 3.13.5
• PyTorch version (GPU?): 2.13.0 (False)
• Accelerate version: 1.14.0
• MPS available: True
• Using distributed or parallel set-up in script?: No

(also reproduced on transformers 5.14.1 and Python 3.11)

Who can help?

@Cyrilvallez

Information

• My own modified scripts

Tasks

• My own task or dataset (give details below)

Reproduction

import torch
from transformers import AutoModelForCausalLM

AutoModelForCausalLM.from_pretrained(
    "trl-internal-testing/tiny-Qwen3ForCausalLM",  # any checkpoint stored in bf16
    dtype=torch.float32,
    device_map="auto",  # resolves to MPS; {"": "mps"} also crashes
)
# -> SIGSEGV (exit 139). faulthandler places the crash in the worker threads of
# core_model_loading._materialize_copy.

All four conditions are required:

1. checkpoint stored in bf16 (a checkpoint stored in f32 loads fine, e.g. hf-internal-testing/tiny-random-LlamaForCausalLM)
2. requested dtype=torch.float32 (forces a dtype conversion; loading as bf16 is fine)
3. MPS target (device_map="auto" or {"": "mps"}; device_map="cpu" is fine)
4. async loading active (the default). HF_DEACTIVATE_ASYNC_LOAD=1 avoids the crash.

What we ruled out (isolation attempts that do NOT reproduce):

• pure torch: ThreadPoolExecutor workers doing t.to(torch.float32).to("mps") -> OK
• mmap-backed safetensors (get_tensor) converted to mps in workers -> OK
• the exact _materialize_copy pattern (get_slice, tensor[...] inside the worker, combined .to(device="mps", dtype=torch.float32)) with 1 and 8 workers -> OK

So the crash seems to be an interaction inside the full loading pipeline (workers materializing while the main thread also touches MPS), not a broken primitive.

Expected behavior

Model loads, or a clean error.

Note this combination is common in practice: TRL's GRPOTrainer(model="...") passes dtype=float32 + device_map="auto" by default, so it segfaults on macOS for most modern (bf16) checkpoints.

Possible fix: convert_and_load_state_dict_in_model already disables the thread pool for disk offload and on-the-fly quantization (each with its own reason). Adding an MPS carve-out to that same condition avoids the crash:

or any((d == "mps" or getattr(d, "type", None) == "mps") for d in device_map.values())

Happy to send that PR if the carve-out is acceptable, or to test a deeper fix.

[Anirud007]

Can confirm this independently on different hardware/Python version, Happy to add data pointes if useful

macOS Apple Silicon, python - 3.12.13
transformers==5.15.0
Reproduced with the exact repro script in the issue (trl-internal-testing/tiny-Qwen3ForCausalLM, dtype=torch.float32, device_map="auto")

Result - segmentation fult (exit 139) right at the start of the weight loading

Tested the proposed carve-out fix (adding an MPS check to the existing thread-pool-disable condition in convert_and_load_state_dict_in_model):

if (
         is_env_variable_true("HF_DEACTIVATE_ASYNC_LOAD")
         or "disk" in device_map.values()
         or has_on_the_fly_quantization
        or any((d == "mps" or getattr(d, "type", None) == "mps") for d in device_map.values())
     ):
         thread_pool = None
     else:
         thread_pool = ThreadPoolExecutor(max_workers=GLOBAL_WORKERS)

[Cyrilvallez]

Hey @sergiopaniego. I cannot repro the issue on mac. What version of safetensors are you using?