huggingface / huggingface/diffusers

Quanto + Group Offload causes device mismatch error (weights on cpu, mat1 on gpu)

Aperta
#12,610 5 commenti 0 reazioni 0 assegnatari Vedi su GitHub
bug stale
Lingua principale
Python
Stelle
34.5k
Fork
7.3k
Merge medio
3g 3h
PR unite (30g)
91

Descrizione

### Describe the bug

When using quanto int8 and Diffusers group offload, we're seeing device mismatch on forward pass.

### Reproduction

```py
"""
This demonstrates the core issue: when .to(device) is called on WeightQBytesTensor,
the wrapper reports the new device but internal components may not be moved properly
when accessed through parameter references.
"""

import torch
from optimum.quanto import freeze, qint8, quantize

class SimpleLinear(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(64, 64)

def forward(self, x):
return self.linear(x)

def demonstrate_device_issue():
compute_device = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cuda")
print(f"Using compute device: {compute_device}")
# Create and quantize
model = SimpleLinear()
quantize(model, weights=qint8)
freeze(model)

print("=" * 70)
print("Initial state (after quantization)")
print("=" * 70)
weight_param = model.linear.weight
print(f"Parameter device: {weight_param.device}")
print(f"Parameter.data type: {type(weight_param.data).__name__}")
print(f"Parameter.data device: {weight_param.data.device}")
if hasattr(weight_param.data, "_data"):
print(f"Parameter.data._data device: {weight_param.data._data.device}")
print(f"Parameter.data._scale device: {weight_param.data._scale.device}")

# Simulate what group offload hook does
print("\n" + "=" * 70)
print(f"Simulating group offload: moving to {compute_device}")
print("=" * 70)
print("Executing: param.data = param.data.to(device)")

# Save references BEFORE the move
old_data_id = id(weight_param.data)
if hasattr(weight_param.data, "_data"):
old_data_backing_id = id(weight_param.data._data)

# This is what the hook does
weight_param.data = weight_param.data.to(compute_device)

# Check AFTER the move
new_data_id = id(weight_param.data)
if hasattr(weight_param.data, "_data"):
new_data_backing_id = id(weight_param.data._data)

print(f"\nAfter .to({compute_device}):")
print(f"Parameter device: {weight_param.device}")
print(f"Parameter.data device: {weight_param.data.device}")
if hasattr(weight_param.data, "_data"):
print(f"Parameter.data._data device: {weight_param.data._data.device}")
print(f"Parameter.data._scale device: {weight_param.data._scale.device}")

print(f"\nObject identity changed:")
print(f" WeightQBytesTensor wrapper: {old_data_id} -> {new_data_id} ({'NEW' if old_data_id != new_data_id else 'SAME'})")
if hasattr(weight_param.data, "_data"):
print(
f" Backing _data tensor: {old_data_backing_id} -> {new_data_backing_id} ({'NEW' if old_data_backing_id != new_data_backing_id else 'SAME'})"
)

# Now try to use it
print("\n" + "=" * 70)
print("Attempting forward pass")
print("=" * 70)

input_tensor = torch.randn(2, 64, device=compute_device)
print(f"Input device: {input_tensor.device}")

try:
output = model(input_tensor)
print(f"✓ Success! Output device: {output.device}")

print("\n" + "=" * 70)
print("Analysis")
print("=" * 70)
print(
"If this succeeded, the parameter reassignment properly updated all references."
)

except RuntimeError as e:
print(f"✗ Failed with device mismatch!")
print(f"\nError: {e}")

print("\n" + "=" * 70)
print("Bug Confirmed")
print("=" * 70)
print("The parameter reassignment created a new WeightQBytesTensor")
print(
"with components on the target device, but the module's parameter"
)
print("reference is pointing to old tensor data.")

if __name__ == "__main__":
print("=" * 70)
print("Quanto WeightQBytesTensor Device Mismatch Demonstration")
print("=" * 70)
print("\nThis simulates what happens during Diffusers group offload hooks\n")

demonstrate_device_issue()

print("\n")
print("=" * 70)

```

### Logs

```shell
======================================================================
Quanto WeightQBytesTensor Device Mismatch Demonstration
======================================================================

This simulates what happens during Diffusers group offload hooks

Using compute device: mps
======================================================================
Initial state (after quantization)
======================================================================
Parameter device: cpu
Parameter.data type: WeightQBytesTensor
Parameter.data device: cpu
Parameter.data._data device: cpu
Parameter.data._scale device: cpu

======================================================================
Simulating group offload: moving to mps
======================================================================
Executing: param.data = param.data.to(device)

After .to(mps):
Parameter device: mps:0
Parameter.data device: cpu
Parameter.data._data device: cpu
Parameter.data._scale device: cpu

Object identity changed:
WeightQBytesTensor wrapper: 13019515792 -> 13019516672 (NEW)
Backing _data tensor: 6162576992 -> 4347766144 (NEW)

======================================================================
Attempting forward pass
======================================================================
Input device: mps:0
✗ Failed with device mismatch!

Error: Tensor for argument #2 'mat2' is on CPU, but expected it to be on GPU (while checking arguments for mm)

======================================================================
Bug Confirmed
======================================================================
The parameter reassignment created a new WeightQBytesTensor
with components on the target device, but the module's parameter
reference is pointing to old tensor data.
```

### System Info

- 🤗 Diffusers version: 0.35.2
- Platform: macOS-15.5-arm64-arm-64bit
- Running on Google Colab?: No
- Python version: 3.12.12
- PyTorch version (GPU?): 2.9.0 (False)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Huggingface_hub version: 0.36.0
- Transformers version: 4.57.1
- Accelerate version: 1.11.0
- PEFT version: 0.17.1
- Bitsandbytes version: not installed
- Safetensors version: 0.6.2
- xFormers version: not installed
- Accelerator: Apple M3 Max
- Using GPU in script?:
- Using distributed or parallel set-up in script?:

Also occurs using NVIDIA CUDA or AMD ROCm.

### Who can help?

@DN6 @sayakpaul

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

Inizia con la riproduzione Python fornita e traccia lo spostamento del parametro group-offload di Diffusers insieme al comportamento di Quanto per WeightQBytesTensor.to. Esamina i punti di ingresso rilevanti per l’offload e la quantizzazione, quindi verifica che il forward pass di un modello quantizzato mantenga pesi e input sullo stesso dispositivo e aggiungi una copertura di regressione per la discrepanza segnalata.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Valutazione

Stack tecnologico
python, pytorch
Ambito
machine-learning
Tipo di issue
Bug
Difficoltà
4/5
Tempo stimato
3-5 giorni
Stato di attività
Ferma
Chiarezza
Da chiarire
Idoneità per principianti
35/100

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.