Comfy-Org / Comfy-Org/ComfyUI

Bug: MixedPrecisionOps Embedding int8_tensorwise fails when DynamicVRAM offloads INT8 weight to GPU

Open
#15,545 1 comment 2 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.