huggingface / huggingface/diffusers

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

Offen
#12,610 5 Kommentare 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
bug stale
Vorherrschende Sprache
Python
Sterne
34.5k
Forks
7.3k
Ø Merge
3 T. 3 Std.
Gemergte PRs (30 T.)
91

Beschreibung

### 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

Beitragsleitfaden

Beitragsleitfaden öffnen

Rechercherichtung

Beginnen Sie mit der bereitgestellten Python-Reproduktion und verfolgen Sie die Verschiebung des Diffusers-Parameters für group-offload zusammen mit Quantos Verhalten von WeightQBytesTensor.to. Untersuchen Sie die relevanten Offload- und Quantisierungs-Einstiegspunkte und überprüfen Sie anschließend, dass der Forward-Pass eines quantisierten Modells Gewichte und Eingaben auf demselben Gerät belässt, und fügen Sie eine Regressionstestabdeckung für die gemeldete Abweichung hinzu.

Vom Indexierungsmodell aus dem Issue-Text verfasst.

Bewertung

Tech-Stack
python, pytorch
Bereich
machine-learning
Issue-Typ
Bug
Schwierigkeit
4/5
Geschätzter Aufwand
3-5 Tage
Aktivitätsstatus
Veraltet
Klarheit
Muss geklärt werden
Anfängerfreundlichkeit
35/100

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.