fp8 comfy_quant models fail on MPS: emulated dequant needs an fp8 cast the backend doesn't implement
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 158
Description
### Custom Node Testing
Reproduced with **no custom nodes involved at all** — the repro below drives `comfy.ops` directly from a Python script, so the node system never runs.
### Expected Behavior
A checkpoint carrying `comfy_quant` fp8 weights should run on MPS (Apple Silicon). When a device can't execute fp8 natively, ComfyUI already intends to fall back to an emulated path, and that fallback should produce images rather than an exception.
### Actual Behavior
Every fp8 layer raises `RuntimeError: Undefined type Float8_e4m3fn` at the first forward pass. `PYTORCH_ENABLE_MPS_FALLBACK=1` does not help — this is a `TORCH_CHECK` inside libtorch's MPS dtype mapping, not a missing-op dispatch, so there's nothing for the fallback to intercept.
**Mechanism**
1. `pick_operations` sees `supports_fp8_compute()` is false and adds `float8_e4m3fn` to `disabled` (`comfy/ops.py:1663`).
2. `disabled` only sets `_full_precision_mm = True` (`comfy/ops.py:1153`). That upcasts the *matmul*, but the weight is still stored in fp8 on the compute device (`comfy/ops.py:1229`).
3. At forward, `comfy_kitchen/backends/eager/quantization.py:63` runs `x.to(dtype=output_type) * scale.to(dtype=output_type)` on an fp8 MPS tensor and raises.
So `disabled` conflates "cannot execute this format natively" with "can emulate this format". On MPS neither holds, and there is no third state.
`unet_dtype()` is not involved — it correctly returns `float16` here, but the `comfy_quant` loader ignores compute dtype and stores `qconfig["storage_t"]`.
**What MPS actually supports.** fp8 tensors can be allocated and moved; every cast and arithmetic op fails:
| op | result |
|---|---|
| `torch.empty(fp8, device="mps")` | OK |
| `.clone()`, `.view()`, `.cpu()` | OK |
| `fp8_mps.to(torch.float16)` | **RuntimeError: Undefined type Float8_e4m3fn** |
| `fp8_mps * 2`, `.zero_()`, `torch.zeros(fp8, "mps")` | **same** |
| `fp8_cpu.to(device="mps", dtype=fp16)` | OK — the cast happens on CPU |
That last row is the escape hatch the proposed fix uses.
### Steps to Reproduce
No model download needed. From the ComfyUI root on an Apple Silicon Mac:
```python
# repro.py -> PYTHONPATH=. python repro.py
import json, torch, comfy.ops as ops, comfy.model_management as mm
dev = mm.get_torch_device()
print("torch", torch.__version__, "device", dev)
# disabled={...} is what pick_operations() passes on any non-CUDA device
Ops = ops.mixed_precision_ops(quant_config={"any": 1}, compute_dtype=torch.float16,
disabled={"float8_e4m3fn", "float8_e5m2"})
lin = Ops.Linear(4, 4, bias=False, device=dev, dtype=torch.float16)
lin._load_from_state_dict({
"l.weight": torch.ones(4, 4, dtype=torch.float16).to(torch.float8_e4m3fn),
"l.weight_scale": torch.tensor(1.0, dtype=torch.float32),
"l.comfy_quant": torch.tensor(list(json.dumps({"format": "float8_e4m3fn"}).encode()), dtype=torch.uint8),
}, "l.", {}, True, [], [], [])
print("loaded weight:", type(lin.weight.data).__name__)
print(lin.forward_comfy_cast_weights(torch.eye(4, dtype=torch.float16, device=dev)))
```
It also reproduces with any real fp8 `comfy_quant` checkpoint. Mine was [`SwarmUI_Z-Image-Turbo-FP8Mix.safetensors`](https://huggingface.co/mcmonkey/swarm-models/blob/main/SwarmUI_Z-Image-Turbo-FP8Mix.safetensors) — Z-Image Turbo, 170 fp8 layers tagged `{"format": "float8_e4m3fn"}`, no `weight_scale`.
**Environment:** ComfyUI `30bdda1e`, torch 2.14.0, macOS 25.6 (Apple Silicon, 36 GiB unified).
### Debug Logs
```powershell
torch 2.14.0 device mps
loaded weight: QuantizedTensor
Traceback (most recent call last):
File "repro.py", line 17, in
print(lin.forward_comfy_cast_weights(torch.eye(4, dtype=torch.float16, device=dev)))
File "ComfyUI/comfy/ops.py", line 1345, in forward_comfy_cast_weights
return self._forward(input, weight, bias)
File "ComfyUI/comfy/ops.py", line 1328, in _forward
return torch.nn.functional.linear(input, weight, bias)
File "comfy_kitchen/tensor/base.py", line 362, in __torch_dispatch__
return op_handlers[parent_cls](qt, args, kwargs)
File "comfy_kitchen/tensor/fp8.py", line 145, in _handle_fp8_linear
return torch.nn.functional.linear(*dequantize_args((input_tensor, weight, bias)))
File "comfy_kitchen/tensor/base.py", line 386, in dequantize_args
return type(args)(dequantize_args(a) for a in args)
File "comfy_kitchen/tensor/base.py", line 382, in dequantize_args
return args.dequantize()
File "comfy_kitchen/tensor/base.py", line 291, in dequantize
full = self.layout_cls.dequantize(qdata, self._params)
File "comfy_kitchen/tensor/fp8.py", line 71, in dequantize
return ck.dequantize_per_tensor_fp8(qdata, params.scale, params.orig_dtype)
File "comfy_kitchen/__init__.py", line 273, in dequantize_per_tensor_fp8
return torch.ops.comfy_kitchen.dequantize_fp8(x, scale, dtype_code)
File "torch/_ops.py", line 1350, in __call__
return self._op(*args, **kwargs)
File "torch/_library/custom_ops.py", line 442, in backend_impl
result = self._backend_fns[device_type](*args, **kwargs)
File "comfy_kitchen/backends/eager/quantization.py", line 461, in _op_dequantize_fp8
return impl(**kwargs)
File "comfy_kitchen/backends/eager/quantization.py", line 63, in dequantize_per_tensor_fp8
dq_tensor = x.to(dtype=output_type) * scale.to(dtype=output_type)
RuntimeError: Undefined type Float8_e4m3fn
```
(Paths shortened for readability; the run is otherwise verbatim.)
### Other
**Proposed fix.** Probe whether the device can cast fp8 at all, and if not, dequantize on CPU at load time and hand back a plain `compute_dtype` weight. Against `30bdda1e`:
```diff
diff --git a/comfy/model_management.py b/comfy/model_management.py
index d56e69e4..3350cc6c 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -1980,6 +1980,35 @@ def supports_fp8_compute(device=None):
return True
+FP8_CAST_SUPPORT = {}
+
+def supports_fp8_cast(device=None):
+ """Whether the device can convert an fp8 tensor to a wider dtype at all.
+
+ Distinct from supports_fp8_compute, which asks whether fp8 matmuls run
+ natively. This asks a much weaker question, and some backends answer no:
+ MPS can allocate, copy and reshape fp8 tensors but implements no fp8
+ kernels, so `fp8_tensor.to(torch.float16)` raises
+ "Undefined type Float8_e4m3fn". The emulated fp8 path is built on exactly
+ that cast, so on such a device fp8 weights have to be upcast on CPU at
+ load time instead of being dequantized inside the op.
+
+ Probed once per device type rather than allowlisted, so a torch release
+ that adds the missing kernels re-enables the normal path with no change
+ here.
+ """
+ if device is None:
+ device = get_torch_device()
+ dev = torch.device(device)
+ if dev.type not in FP8_CAST_SUPPORT:
+ try:
+ torch.empty((1,), dtype=torch.float8_e4m3fn, device=dev).to(torch.float32)
+ FP8_CAST_SUPPORT[dev.type] = True
+ except Exception:
+ logging.info("Device {} cannot cast fp8 tensors; fp8 weights will be upcast at load time.".format(dev.type))
+ FP8_CAST_SUPPORT[dev.type] = False
+ return FP8_CAST_SUPPORT[dev.type]
+
def supports_nvfp4_compute(device=None):
if not is_nvidia():
return False
diff --git a/comfy/ops.py b/comfy/ops.py
index ff64aad5..83e97974 100644
--- a/comfy/ops.py
+++ b/comfy/ops.py
@@ -1143,6 +1143,21 @@ def _load_quantized_module(module, super_load, state_dict, prefix, local_metadat
if layer_conf is not None:
layer_conf = json.loads(layer_conf.numpy().tobytes())
+ if layer_conf is not None and layer_conf.get("format") in ("float8_e4m3fn", "float8_e5m2") \
+ and not comfy.model_management.supports_fp8_cast():
+ # The emulated fp8 path dequantizes inside the op, which needs an
+ # fp8 -> compute_dtype cast on the compute device. Backends without fp8
+ # kernels (MPS) cannot do that cast at all, so dequantize here on CPU and
+ # hand back a plain compute_dtype weight. Costs the memory saving but is
+ # the difference between running and not.
+ scale_key = f"{prefix}weight_scale"
+ scale = state_dict.pop(scale_key, None)
+ weight = weight.cpu().to(dtype=compute_dtype)
+ if scale is not None:
+ manually_loaded_keys.append(scale_key)
+ weight = weight * scale.cpu().to(dtype=compute_dtype)
+ layer_conf = None
+
if layer_conf is None:
module.weight = torch.nn.Parameter(weight.to(device=device, dtype=compute_dtype), requires_grad=False)
else:
```
The probe is cached per device type rather than allowlisting MPS by name, so a torch release that adds the missing kernels restores the normal path with no code change here.
**Trade-off.** The memory saving is lost — the Z-Image checkpoint goes from ~6 GB to ~12 GB resident. On a device that cannot execute fp8 in any form that seems the right call, since the alternative is not running at all, but if you would rather it were opt-in behind a flag, or a clear refusal at load time instead of a silent upcast, say so and I'll rework it.
**Verification.** Same script, same seed, patch on and off:
| | QuantizedTensor weights after load | result |
|---|---|---|
| unpatched | 170 | `RuntimeError: Undefined type Float8_e4m3fn` |
| patched | 0 | 4 steps in ~43 s, finite latents, correct image |
Full pipeline (qwen_3_4b text encoder + Flux VAE, euler/simple, 4 steps, cfg 1, 512x512) produces the expected image for "a red cube on a white table, studio photo". The dequant arithmetic is separately bit-exact against `fp8_ref.float() * scale`, tested both with and without a `weight_scale`.
**Deliberately not covered**, to keep the change reviewable — same defect, same cause:
- `MixedPrecisionOps.Embedding._load_from_state_dict` (`comfy/ops.py:1565`) has its own fp8 loader with the identical problem.
- `nvfp4` stores `uint8` data but a `float8_e4m3fn` block scale, and `mxfp8` uses a `float8_e8m0fnu` scale. Both should fail on MPS the same way. Untested — I have no such checkpoint to hand.
- I have no NVIDIA hardware, so this change being a no-op on CUDA is reasoned (`supports_fp8_cast` probes true, the guard never fires), not executed. Worth a second pair of eyes.
**Separate thing noticed in passing**, not investigated and possibly not a real bug: if an fp8 layer has *no* `weight_scale`, `Params.clone()` raises `AttributeError: 'NoneType' object has no attribute 'clone'` (`comfy_kitchen/tensor/base.py:73`) when `torch.nn.Parameter()` detaches the `QuantizedTensor`. It fires in the isolated repro above but *not* when loading the real checkpoint, which also lacks scales — so something about the real load path avoids it and I haven't chased down what. Mentioning it only in case it's meaningful to you; it's why the repro above passes an explicit scale of 1.0.
Contributor guide
Research direction
Start with the repro.py script and trace comfy/model_management.py:supports_fp8_compute alongside _load_quantized_module in comfy/ops.py. Check the existing comfy_kitchen eager dequantization path and the separate Embedding loader mentioned in the issue. Done means fp8 comfy_quant weights run on MPS without the Float8_e4m3fn error, while CUDA behavior remains unchanged.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- macos, python, pytorch
- Domain
- backend, machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 52/100