float8_e4m3fn Fails on MPS Due to Execution Order with Scaled Models
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 158
Description
After debugging this for a while, I think I came up with some good info on this. This was done with the help of Google Gemini Pro 2.5 and other AIs, and I guided them when they did AI things, so I'm confident about this, but it still has a chance of being wrong.
System Information:
* Hardware: Apple Silicon Mac (M1 Max chip)
* Backend: MPS (Metal Performance Shaders)
* Model: flux1-dev-kontext_fp8_scaled.safetensors
* ComfyUI Version: 0.3.43
### Summary
Loading fp8 models that use a "scaled" format, such as flux1-dev-kontext_fp8_scaled.safetensors, fails on the MPS backend. The root cause is an order-of-operations conflict in the model loading process. The generic ComfyUI framework attempts to cast or move the fp8 tensor to the MPS device before the model's specialized patching logic has a chance to execute, leading to a fatal, unsupported dtype error.
### Symptoms
1. Initial Error: Loading the model with an unmodified ComfyUI instance fails with the standard error: "Trying to convert Float8_e4m3fn to the MPS backend but it does not have support for that dtype." This occurs because the generic cast_to function calls a native PyTorch operation that is not implemented for fp8 on MPS.
2. Manual Conversion Failures: Attempts to work around this by writing a manual fp8 -> fp16 conversion function inside cast_to also fail, but produce different symptoms depending on the decoding logic:
* Clamped Output: Decoding the fp8 values and clamping them to a [-1, 1] range produces a recognizable but heavily distorted image with color shifting and saturation artifacts, similar to an "overbaked" LoRA.
* Unclamped Output: Faithfully decoding the fp8 values to their full "unscaled" magnitude (e.g., values like -26.0, 18.0) results in a completely noisy image, as these large values cause numerical instability in the model without their corresponding scaling factors.
### Root Cause Analysis
The model flux1-dev-kontext_fp8_scaled.safetensors does not use a standard fp8 format. It uses a two-part scaled system where the final weight for a layer is derived from:
1. An fp8 tensor of base values.
2. A separate fp16 tensor of scaling factors.
The logic to handle this is not in a central conversion function but is implemented as a patching mechanism within the comfy.ops.Linear class, which is used by the Flux model definition in comfy/ldm/flux/model.py. This custom Linear layer is designed to perform the fp8 * scale multiplication on the fly.
The failure occurs because of this sequence of events:
1. A model layer is loaded.
2. The generic cast_to function in model_management.py is called on the fp8 weight tensor.
3. The cast_to function returns, and the framework immediately attempts to move the tensor to the mps device.
4. This triggers the native PyTorch operation for device placement on an fp8 tensor, which is unsupported on MPS, causing the crash.
The specialized patching logic in the Flux model and comfy.ops never gets a chance to run because the process fails at the generic device placement step.
### Suggested Solution
This issue cannot be resolved by modifying the fp8 conversion logic itself. The fix requires a change to the core model loading execution flow.
The proposed solution is to defer device placement for fp8 tensors until after any model-specific patching has been applied. The framework needs to be aware of these special "scaled" model types and allow their patching mechanisms to execute first, which would correctly convert the weights to a supported fp16 or bfloat16 format before the tensor is moved to the final device.
Contributor guide
Assessment
This issue has not been assessed yet.