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

快速结论：这个 SIGSEGV 通常出现在 macOS Apple Silicon（MPS）上，用 device_map="auto" 加载 bf16 权重、同时强制 dtype=torch.float32 触发权重转换时。异步加载线程池内有多个 worker 并发做 MPS 转换会踩到竞态；优先排查是否命中“bf16 权重 + f32 + MPS + 异步加载”这四个条件同时成立。

适用环境：Issue 中报告并复现的环境：macOS（Apple Silicon，arm64），MPS 可用；transformers 5.15.0（5.14.1 也复现，current main 复现）；Python 3.13.5（3.11、3.12.13、3.13.14 也复现）；PyTorch 2.13.0；Accelerate 1.14.0；safetensors 0.8.0。未使用分布式。

最快修复方案：暂无确认的一步修复方案。Issue 中已验证可绕过崩溃的方式是设置环境变量 HF_DEACTIVATE_ASYNC_LOAD=1，禁用异步加载。

注意事项：该环境变量只是绕过异步加载路径，不是根因修复，可能带来加载速度下降，且 Issue 未评估其其它副作用。讨论中提出的“在 convert_and_load_state_dict_in_model 中为 MPS 增加线程池禁用分支”以及“只对 MPS 转换加进程级锁”的补丁均未在正式版本中确认合入。

问题场景

用户在 macOS Apple Silicon 上使用 Transformers 的 AutoModelForCausalLM.from_pretrained() 加载一个 bf16 存储的 checkpoint（复现用 trl-internal-testing/tiny-Qwen3ForCausalLM），同时指定 dtype=torch.float32 和 device_map="auto"（解析到 MPS）。加载权重时进程直接段错误退出（exit 139），faulthandler 定位在 core_model_loading._materialize_copy 的 worker 线程内。相同组合在 GPU ? False、无分布式、MPS 可用的环境下触发。

该组合在现实中很常见：TRL 的 GRPOTrainer(model="...") 默认传 dtype=float32 + device_map="auto"，因此在 macOS 上加载大多数现代 bf16 checkpoint 会直接段错误。

报错原文

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

# -> SIGSEGV (exit 139). faulthandler places the crash in the worker threads of
# core_model_loading._materialize_copy.

# faulthandler: pool worker at core_model_loading.py:1221
tensor.to(device=device, dtype=dtype)

原因分析

Issue 中未给出最终确认的根因，但排查已经收窄了范围。以下四个条件必须同时成立才会崩溃，任一条件不满足都能正常加载：

checkpoint 以 bf16 存储（f32 存储的 checkpoint，例如 hf-internal-testing/tiny-random-LlamaForCausalLM，加载正常）
请求 dtype=torch.float32，即强制发生 dtype 转换（按 bf16 加载正常）
目标是 MPS（device_map="auto" 或 {"": "mps"}；device_map="cpu" 正常）
异步加载处于激活状态（默认行为；HF_DEACTIVATE_ASYNC_LOAD=1 可避免崩溃）

用户已排除的“纯原语”问题：纯 torch 中用 ThreadPoolExecutor worker 执行 t.to(torch.float32).to("mps") 正常；mmap 后端 safetensors（get_tensor）在线程中转到 MPS 正常；单独复刻 _materialize_copy 的模式（worker 内 get_slice、tensor[...]，配合 .to(device="mps", dtype=torch.float32)），1 个和 8 个 worker 都正常。因此疑似不是原语本身坏了，而是完整加载管线内部的交互问题——worker 在物化权重的同时主线程也在接触 MPS。

另一个复现者补充：线程数影响结果，但不稳定。workers=1 在 40 次连续运行中 0 次失败；workers=2 在 40 次中全部失败；workers=4 在 40 次中失败 33 次，且第二次采样中出现了 7/20 的成功。失败表现为 SIGSEGV、SIGABRT、SIGTRAP、SIGBUS 或挂起（60s 超时未完成）五种不同结局，更像竞态而非确定性的坏指针。由于 workers=1 同样跑在主线程之外却从不失败，问题可能是并发而非“后台线程访问 MPS”本身。

环境排查

确认操作系统与芯片：macOS Apple Silicon（arm64），Issue 在 macOS 26.x 上复现。
确认 MPS 是否可用（MPS available: True）。
确认 transformers 版本：5.15.0、5.14.1 及 current main（commit 0c92811846）均复现。
确认 Python 版本：3.11、3.12.13、3.13.5、3.13.14 均复现。
确认 PyTorch 版本：2.13.0。
确认 accelerate 版本：1.14.0。
确认 safetensors 版本：0.8.0（首帖与补充复现者一致）。
确认 checkpoint 的存储 dtype 是否为 bf16，以及是否传入 device_map（"auto" 或 {"": "mps"}）。
确认是否设置了 HF_DEACTIVATE_ASYNC_LOAD，以及是否可通过覆盖 GLOBAL_WORKERS 改变 worker 数量来观察稳定性差异。

解决步骤

先完整记录崩溃条件：确认脚本同时满足 bf16 存储的 checkpoint、dtype=torch.float32、MPS 目标（device_map="auto" 或 {"": "mps"}）以及异步加载默认开启。任一条件不满足时应能正常加载，可作为对照。
可优先尝试在运行前设置环境变量 HF_DEACTIVATE_ASYNC_LOAD=1，禁用异步加载，绕开崩溃路径。Issue 中已验证该方式可避免 SIGSEGV，这是当前最直接的验证手段。
如果只是要加载模型而不需要强制 f32 转换，可改为按 bf16 加载（不指定 dtype=torch.float32），或改用 device_map="cpu"，Issue 表格显示这两种配置均正常。
如果你在排查该缺陷或测试补丁，可按讨论中的思路验证：在 convert_and_load_state_dict_in_model 中，把已有的线程池禁用条件（HF_DEACTIVATE_ASYNC_LOAD 为真、disk 卸载、on-the-fly 量化）扩展到 MPS 目标，例如加入对 device_map.values() 中 MPS 设备的判断。注意这是讨论中的候选修复，未确认合入。
另一种讨论中的更窄方案是：保留多 worker 加载，仅在 _materialize_copy 内对目标为 MPS 的 tensor.to(...) 转换使用进程级锁串行化。该方案同样未确认合入。
不要仅凭“降低 worker 数”来当修复：GLOBAL_WORKERS=1 在采样中 40 次全过，但 4 workers 也有成功运行，说明该现象是概率性的竞态，靠调 worker 数无法稳定规避。

验证方法

用原始复现脚本再次加载 bf16 checkpoint，确认进程不再以 exit 139（SIGSEGV/SIGABRT/SIGTRAP/SIGBUS）退出且模型能正常完成加载。由于 Issue 显示该问题存在概率性（例如 4 workers 时并非每次都失败），单次成功不足以证明修复有效，建议连续多次运行并对比通过率；同时可保留“任一必要条件被破坏即正常”的对照表来确认你验证的就是同一条崩溃路径。