huggingface / huggingface/diffusers

pipe.to(device) corrupts custom quantized (uint8 packed) module buffers

Aperta
#14,449 1 commento 0 reazioni 0 assegnatari Vedi su GitHub
bug needs-env-info pipelines quantization
Lingua principale
Python
Stelle
34.5k
Fork
7.3k
Merge medio
3g 3h
PR unite (30g)
91

Descrizione

## Description

`StableDiffusionXLPipeline.to(device)` re-casts all module parameters and buffers, which corrupts packed `uint8` buffers used by custom Int4 quantized layers. The buffers store two 4-bit weights per byte in `uint8` format — when the pipeline moves them through `.to(device)` or casts dtype, the packed representation is destroyed.

## Reproduction

```python
# Custom Int4 layer stores weights as packed uint8 (2x int4 per byte)
# with float16 scale/zero_point per group
class Int4LinearMPS(nn.Module):
def __init__(self, ...):
self.register_buffer('packed_weight', torch.zeros(..., dtype=torch.uint8))
self.register_buffer('scales', torch.ones(..., dtype=torch.float16))

# After quantizing UNet and injecting LoRA:
unet = ... # Int4 quantized, on MPS, working

# This WORKS:
pred = unet(latent, t, encoder_hidden_states=enc) # ✓ Valid output

# This BREAKS the uint8 buffers:
pipe = StableDiffusionXLPipeline.from_pretrained(model_name, unet=unet, ...)
pipe.to('mps') # ← corrupts packed_weight buffers
# All subsequent generations produce NaN/black images
```

## Root Cause

`pipe.to(device)` calls `.to()` on all submodules recursively. For standard fp16/fp32 parameters this is fine. For packed `uint8` buffers (used in Int4/Int2 quantization), the `.to()` call may:
1. Attempt dtype conversion (uint8 → float16)
2. Move buffers through an intermediate state that corrupts packing

## Workaround

Use a manual inference loop instead of the pipeline:
```python
# Pre-compute text embeddings, offload encoders
# Run UNet denoising loop manually
# Decode with VAE separately in float32
```

## Expected Behavior

`pipe.to(device)` should respect buffer dtypes and not re-cast `uint8` buffers. Custom quantized modules with non-standard buffer dtypes should be preserved during device transfer.

## Environment
- diffusers 0.39.0
- PyTorch 2.13
- Apple M1, MPS backend
- macOS

## Relevance

This affects anyone using custom quantization (Int4, Int2, GPTQ-style packed weights) with diffusers pipelines. As quantization becomes more common for edge deployment, this incompatibility will affect more users.

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

Start at StableDiffusionXLPipeline.to(device) and trace how recursive PyTorch module transfers handle registered buffers. Reproduce with the custom Int4LinearMPS packed_weight and scales buffers on MPS, then verify that device transfer preserves uint8 packing and that subsequent pipeline generation does not produce NaN or black images.

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à
Tranquilla
Chiarezza
Abbastanza chiara
Idoneità per principianti
45/100

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.