Quantized weight dequant/cast order drops requested dtype (crashes when aimdo-offloaded + float32 input, e.g. MiniMax H3 vision-conditioned encode)
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 158
Description
### Summary
`resolve_cast_module_with_vbar()`'s `to_dequant()` helper and `cast_bias_weight()`'s inline equivalent both call `QuantizedTensor.to(dtype)` **before** `.dequantize()`. The requested target dtype passed to `.to()` is not honored by `.dequantize()`, which always emits the tensor's original quantization dtype (e.g. `bfloat16`) regardless of what was requested. Net effect: a `Linear` that needs a non-native compute dtype gets back a weight in the wrong dtype, and crashes with:
```
RuntimeError: expected mat1 and mat2 to have the same dtype, but got: float != c10::BFloat16
```
### Where it bites
MiniMax H3's text encoder (Qwen3-VL, vision-conditioned) builds its embeddings in `float32` (`comfy/sd1_clip.py`, `dtype=torch.float32` at the call site — the standard CLIP convention). `comfy/text_encoders/llama.py`'s `forward()` only applies the `dtype` argument on the token-id path (`self.embed_tokens(x, out_dtype=dtype)`); the `embeds is not None` branch (used by MiniMax H3 / any vision-conditioned encoder) passes `embeds` straight through unmodified, so `x` stays `float32` end to end.
When a downstream `Linear` layer is INT8/INT4-quantized **and** has been dynamically offloaded by the aimdo dynamic-VRAM allocator (`has_v` / vbar-backed), `comfy/ops.py`'s `resolve_cast_module_with_vbar()` resolves the weight through:
```python
def to_dequant(tensor, dtype):
tensor = tensor.to(dtype=dtype) # no-op on a still-quantized tensor
if isinstance(tensor, QuantizedTensor):
tensor = tensor.dequantize() # emits the tensor's native dtype, ignoring the .to() above
return tensor
```
`cast_bias_weight()` has the identical ordering bug:
```python
if weight_has_function or weight.dtype != dtype:
weight = weight.to(dtype=dtype)
if isinstance(weight, QuantizedTensor):
weight = weight.dequantize()
...
```
Since `.to()` runs before `.dequantize()`, the requested `dtype` is discarded and the dequantized weight always comes back in its native storage dtype (`bfloat16` in this case), while `input` is `float32` — hence the crash in `torch.nn.functional.linear`.
This only manifests once the weight is aimdo/vbar-offloaded (`has_v=True`), which is why it doesn't show up on a plain/low-memory text-only encode — it needs enough VRAM pressure (e.g. a full MiniMax H3 image/video workflow with the diffusion model + VAE + text encoder all resident) for aimdo to actually offload these particular layers.
### Repro
- Model: `qwen3vl_32b_minimax_h3_int4_convrot.safetensors` (INT8/INT4 mixed, TensorWiseINT8Layout on the affected layer) as the MiniMax H3 text encoder, loaded alongside the MiniMax H3 diffusion model + VAE so aimdo dynamic-VRAM offloading engages.
- Run any `MiniMaxH3ImageToVideo` (or other vision-conditioned MiniMax H3) workflow.
- Crash:
```
RuntimeError: expected mat1 and mat2 to have the same dtype, but got: float != c10::BFloat16
File "comfy/ops.py", line 1338, in _forward
return torch.nn.functional.linear(input, weight, bias)
```
Confirmed via a temporary debug print at the failure site:
```
input=torch.float32/Tensor weight=torch.bfloat16/Tensor layout_type=TensorWiseINT8Layout
has_v=True weight_shape=(8192, 5120) raw_weight_dtype=torch.bfloat16/QuantizedTensor
```
`weight` comes back as a plain `Tensor` (so dequantization did run) but in `bfloat16` instead of the requested `float32`.
### Fix
Dequantize first, then cast to the actually-requested dtype — swap the two lines in both spots:
```python
def to_dequant(tensor, dtype):
if isinstance(tensor, QuantizedTensor):
tensor = tensor.dequantize()
tensor = tensor.to(dtype=dtype)
return tensor
```
```python
if weight_has_function or weight.dtype != dtype:
if isinstance(weight, QuantizedTensor):
weight = weight.dequantize()
weight = weight.to(dtype=dtype)
for f in s.weight_function:
weight = f(weight)
```
Verified locally: this fully resolves the crash on the MiniMax H3 image/video workflow above, with no regression observed on the plain text-only path.
### Environment
- ComfyUI: v0.34.0 (master, commit `8a33128f`) — also reproduced on an older commit (`40e46c71`), so this is not a recent regression in `comfy/ops.py` itself.
- comfy-kitchen 0.2.31, comfy-aimdo 0.4.15
- torch 2.10.0+cu130, single RTX 3090
Happy to open a PR with this exact two-line diff if useful.
Contributor guide
Research direction
Start in comfy/ops.py at resolve_cast_module_with_vbar() and cast_bias_weight(), focusing on their quantized-weight dequantization paths. Reproduce with the MiniMax H3 vision-conditioned workflow and the listed quantized model, then verify that the requested input dtype is retained after dequantization and the workflow no longer raises the reported linear dtype mismatch.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 84/100