int8_tensorwise Embedding breaks when cast returns a dequantized weight: "No backend can handle 'dequantize_int8_embedding'"
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 158
Description
**What happens**
With an `int8_tensorwise` / convrot-quantized text encoder on a setup where the TE gets offloaded/staged (typical on ≤12 GB cards), any embedding lookup fails:
```
NoCapableBackendError: No backend can handle 'dequantize_int8_embedding': eager: q: dtype torch.bfloat16 not in {torch.int8}
```
Regression relative to pre-memory-management behaviour; fp8 TEs are unaffected. Reported downstream with two affected users: lbouaraba/comfyui-krea2edit#19.
**Root cause**
`cast_bias_weight` deliberately returns an already-**dequantized plain tensor** on some paths (e.g. CPU-resident weights: `weight = weight.dequantize()` in the CPU branch). The quantized-Embedding `forward_comfy_cast_weights` handles that case for its params/scale bookkeeping, but the `int8_tensorwise` branch still unconditionally routes the now-plain bf16 tensor into `dequantize_embedding` → the kitchen op's dtype constraint fails.
The int8 **Linear** path already has exactly the guard this path is missing, with a comment documenting that casts can hand back dequantized weights ("A LoRA weight_function, or activations whose dtype differs from the weight's, make the cast hand back a dequantized tensor.").
**Minimal repro (CPU-only, no model needed)**
```python
import torch
from comfy_kitchen.tensor.int8 import TensorWiseINT8Layout
class P: pass
p = P(); p.scale = torch.tensor(1.0); p.orig_dtype = torch.bfloat16
TensorWiseINT8Layout.dequantize_embedding(
torch.randn(100, 64, dtype=torch.bfloat16), p, torch.tensor([1, 2, 3]))
# -> NoCapableBackendError, exact message above
```
End-to-end: any int8_tensorwise TE plus conditions that stage/offload it.
**Fix**
Gate the int8 gather on the cast having returned a `QuantizedTensor`; plain data falls through to the existing `F.embedding` path. PR incoming, referencing this issue.
Contributor guide
Research direction
Start with TensorWiseINT8Layout.dequantize_embedding and the quantized Embedding forward_comfy_cast_weights path, then compare them with the existing int8 Linear guard described in the issue. Run the CPU-only minimal reproduction first. Done means plain dequantized weights use the existing F.embedding path while QuantizedTensor weights retain the int8 gather path without the backend dtype error.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 52/100