huggingface / huggingface/diffusers

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

Open
#14,449 1 comment 0 reactions 0 assignees View on GitHub
bug needs-env-info pipelines quantization
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.