huggingface / huggingface/diffusers
Helios schedulers raise `TypeError` on MPS: `float64` without the device guard
- Ngôn ngữ chính
- Python
- Star
- 34.5k
- Fork
- 7.3k
- Merge trung bình
- 3 ngày 3 giờ
- Pull request đã merge (30 ngày)
- 91
Mô tả
### Describe the bug
Both Helios schedulers raise `TypeError` on MPS (Apple Silicon) because they build `float64` tensors without the device guard the rest of the codebase uses. `HeliosPyramidPipeline` calls `set_timesteps(..., device=device)`, so on an MPS device the pipeline fails at the first scheduler call, before any model runs.
Two sites, both verified by traceback on `main` (0.40.0.dev0):
**1. `set_timesteps` — affects both schedulers**
```
scheduling_helios.py:238 self.timesteps = torch.from_numpy(timesteps).to(device=device)
scheduling_helios_dmd.py:216 self.timesteps = torch.from_numpy(timesteps).to(device=device)
```
`timesteps` is a NumPy `float64` array, so `torch.from_numpy` yields `float64` and `.to(device="mps")` raises.
**2. `convert_flow_pred_to_x0` — `HeliosDMDScheduler` only, reached from `step()`**
```
scheduling_helios_dmd.py:278 flow_pred, xt, sigmas, timesteps = (x.double().to(device) for x in (...))
```
### Reproduction
CPU only for the passing case, MPS for the failing one — no model weights needed:
```python
import torch
from diffusers import HeliosScheduler, HeliosDMDScheduler
# 1. set_timesteps on an MPS device — both schedulers
for name, cls in (("HeliosScheduler", HeliosScheduler), ("HeliosDMDScheduler", HeliosDMDScheduler)):
try:
cls().set_timesteps(4, device="mps", stage_index=0)
print(name, "ok")
except TypeError as e:
print(name, "TypeError:", str(e)[:60])
# 2. HeliosDMDScheduler.step() with tensors already on MPS
d = HeliosDMDScheduler()
a = torch.randn(1, 4, 2, 8, 8, device="mps")
b = torch.randn(1, 4, 2, 8, 8, device="mps")
ts = torch.tensor([1000., 750., 500., 250.], device="mps")
sg = torch.tensor([1., .75, .5, .25], device="mps")
d.step(model_output=a, timestep=1000.0, sample=b, cur_sampling_step=3,
dmd_noisy_tensor=torch.zeros_like(b), dmd_sigmas=sg,
dmd_timesteps=ts, all_timesteps=ts)
```
```
HeliosScheduler TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS
HeliosDMDScheduler TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS
Traceback (most recent call last):
scheduling_helios_dmd.py:298 in step
scheduling_helios_dmd.py:278 in convert_flow_pred_to_x0
TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64.
```
Both schedulers work correctly on CPU with the same inputs.
### Expected behavior
The intent at `scheduling_helios_dmd.py:277` is stated as *"use higher precision for calculations"* — the goal is precision, not `float64` specifically, so a downcast on backends that lack it preserves the intent.
The codebase already has the mechanism. `maybe_adjust_dtype_for_device` in `utils/torch_utils.py` maps `float64 -> float32` for `mps`/`npu`/`neuron` and is used at ~53 call sites. Neither Helios scheduler imports it. A sibling scheduler handles the same problem explicitly:
```python
# scheduling_flow_map_euler_discrete.py:169-173
# ...the final tensors to the requested device (with a float32 downcast for MPS / NPU).
is_mps = device_obj is not None and device_obj.type == "mps"
is_npu = device_obj is not None and device_obj.type == "npu"
out_dtype = torch.float32 if (is_mps or is_npu) else torch.float64
```
### Scope
Verified at the scheduler level only. I have not run `HeliosPyramidPipeline` end to end on MPS — that needs the Helios-14B weights, which I cannot download here — so I have not established whether the pipeline would otherwise work on MPS once the schedulers are fixed. What is established is that `pipeline_helios_pyramid.py:941` passes `device=device` into `set_timesteps`, so this is on the pipeline's path rather than an artificial call.
Also not tested: `npu` and `neuron`, which `_DTYPE_UNSUPPORTED_DEVICES` lists alongside `mps` for `float64`. They may be affected identically.
Unrelated to #14353, which covers `HeliosDMDScheduler`'s ignored config options in the same file — different defect, no overlap in the fix.
### System Info
- diffusers version: 0.40.0.dev0 (`main`, commit `a8345366e`)
- Platform: macOS 26.0.1, arm64 (Apple Silicon)
- Python version: 3.11
- PyTorch version: 2.13.0, `torch.backends.mps.is_available() == True`
### Who can help?
@dg845 (`git blame` points to #13208 for both files) and @yiyixuxu
Hướng dẫn đóng góp
Hướng nghiên cứu
Bắt đầu với scheduling_helios.py và scheduling_helios_dmd.py tại các vị trí set_timesteps và convert_flow_pred_to_x0 được báo cáo, sau đó đọc maybe_adjust_dtype_for_device trong utils/torch_utils.py và phần xử lý tương ứng trong scheduling_flow_map_euler_discrete.py. Chạy các bản tái hiện issue của các scheduler trên CPU và MPS, bao gồm HeliosDMDScheduler.step(). Hoàn thành khi cả hai scheduler đều tránh các phép toán float64 không được MPS hỗ trợ, đồng thời vẫn giữ các phép tính có độ chính xác cao hơn khi thiết bị hỗ trợ chúng.
Do mô hình lập chỉ mục viết ra từ nội dung của issue.
Đánh giá
- Công nghệ
- python, pytorch
- Lĩnh vực
- machine-learning
- Loại issue
- Lỗi
- Độ khó
- 3/5
- Thời gian dự kiến
- 1-2 ngày
- Mức độ hoạt động
- Sôi nổi
- Độ rõ ràng
- Đặc tả rõ ràng
- Mức phù hợp với người mới
- 74/100