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

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

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

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