Bug: MixedPrecisionOps Embedding int8_tensorwise fails when DynamicVRAM offloads INT8 weight to GPU
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 6h
- Merged PRs (30d)
- 155
Description
## Description
When using an INT8 TensorWise quantized text encoder (e.g., MiniMax-H3's Qwen3-VL INT8 ConvRot, 25.77GB) with DynamicVRAM enabled, MixedPrecisionOps.Embedding.forward_comfy_cast_weights crashes:
```
NoCapableBackendError: No backend can handle 'dequantize_int8_embedding':
eager: q: dtype torch.bfloat16 not in {torch.int8}
```
The same model works fine on CPU. The issue is in comfy/ops.py, not comfy_kitchen.
## Root Cause
In MixedPrecisionOps.Embedding.forward_comfy_cast_weights (ops.py ~line 1643), when the Embedding weight is a QuantizedTensor:
```python
qdata, _, offload_stream = cast_bias_weight(self, device=input.device,
dtype=weight.dtype, offloadable=True)
if isinstance(qdata, QuantizedTensor):
# ... optimized dequantize_embedding path ...
else:
params = weight._params
scale = None
```
The else branch assumes qdata is non-quantized but still came from a normal cast path, so it assigns params from the original weight. Then code falls through to:
```python
if self.quant_format == "int8_tensorwise":
x = get_layout_class(self.layout_type).dequantize_embedding(qdata, params, input)
```
But qdata is now a plain bf16 tensor because DynamicVRAM's vbar path (resolve_cast_module_with_vbar) already dequantized the INT8 weight to bf16 during the GPU transfer. Calling dequantize_int8_embedding on bf16 data fails.
The bug only triggers when vbar decides to offload that specific Embedding layer — non-deterministic depending on VRAM pressure. This explains why the same model sometimes works and sometimes fails on the same setup.
## Environment
- ComfyUI v0.31.0 (Windows standalone)
- comfy_aimdo 0.4.13 (DynamicVRAM)
- comfy_kitchen 0.2.28
- CUDA backend: does NOT support dequantize_int8_embedding
- Eager backend: supports dequantize_int8_embedding but receives bf16 dtype
- Test model: linjian257/qwen3vl_32b_minimax_h3_int8_convrot_uncensored (INT8 TensorWise + ConvRot)
## Fix
Move the non-QuantizedTensor handling out of the else clause and into its own branch that just uses F.embedding (since vbar already baked in the scale during dequantization):
```python
if isinstance(weight, QuantizedTensor) and len(self.weight_function) == 0:
qdata, _, offload_stream = cast_bias_weight(self, device=input.device,
dtype=weight.dtype, offloadable=True)
if isinstance(qdata, QuantizedTensor):
# Standard path: still quantized, use optimized dequantize_embedding
params = qdata._params
scale = params.scale
qdata = qdata._qdata
if self.quant_format == "int8_tensorwise":
x = get_layout_class(self.layout_type).dequantize_embedding(
qdata, params, input)
uncast_bias_weight(self, qdata, None, offload_stream)
return x if out_dtype is None else x.to(dtype=out_dtype)
x = torch.nn.functional.embedding(
input, qdata, self.padding_idx, self.max_norm,
self.norm_type, self.scale_grad_by_freq, self.sparse)
uncast_bias_weight(self, qdata, None, offload_stream)
target_dtype = out_dtype if out_dtype is not None else weight._params.orig_dtype
x = x.to(dtype=target_dtype)
if scale is not None and scale != 1.0:
x = x * scale.to(dtype=target_dtype)
return x
else:
# DynamicVRAM (vbar) already dequantized the weight to raw bf16/fp16.
# Scale is already baked in by vbar's dequantization pass.
x = torch.nn.functional.embedding(
input, qdata, self.padding_idx, self.max_norm,
self.norm_type, self.scale_grad_by_freq, self.sparse)
uncast_bias_weight(self, qdata, None, offload_stream)
return x if out_dtype is None else x.to(dtype=out_dtype)
```
**Impact**: The fix is a pure no-op when DynamicVRAM does NOT touch the Embedding layer (standard path unchanged). When vbar does offload, we correctly use F.embedding instead of trying to call int8-specific dequantize on already-dequantized data. Tested and confirmed working on RTX 4070 Ti Super 16GB with the MiniMax-H3 INT8 text encoder.
Contributor guide
Research direction
Start in comfy/ops.py around MixedPrecisionOps.Embedding.forward_comfy_cast_weights and trace cast_bias_weight, including the resolve_cast_module_with_vbar path. Confirm the quantized path still uses dequantize_embedding, while a plain bf16/fp16 result uses F.embedding without int8 dequantization; validate with the reported INT8 TensorWise MiniMax-H3 setup under DynamicVRAM and without it.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 72/100