huggingface / huggingface/diffusers
pipe.to(device) corrupts custom quantized (uint8 packed) module buffers
- Dominant language
- Python
- Stars
- 34.5k
- Forks
- 7.3k
- Avg merge
- 3d 3h
- Merged PRs (30d)
- 91
Description
## 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.
Contributor guide
Research direction
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.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100