Support offloading feature in diffuser

[jerryzh168]

1. diffuser currently is using older torchao versions, we should update it to stable 0.14.1 (and use nightly for testing)

2. Repro error and add missing support torchao to unskip the following tests:

Note: right now they are using int8 weight only: https://github.com/huggingface/diffusers/blob/d5da453de56fe73e0cfd26204ccca441af568ca1/tests/quantization/torchao/test_torchao.py#L660, but we haven't fully migrated this config yet, maybe we can start with float8 weight only config which is using the new Float8Tensor design and expand to other quant types.

https://github.com/huggingface/diffusers/blob/d5da453de56fe73e0cfd26204ccca441af568ca1/tests/quantization/torchao/test_torchao.py#L664

https://github.com/huggingface/diffusers/blob/d5da453de56fe73e0cfd26204ccca441af568ca1/tests/quantization/torchao/test_torchao.py#L673

int8 weight only is still more important, but migration is in progress: https://github.com/pytorch/ao/pull/3241/files, we can use float8 to figure out what fix is needed first and then do the same fix on the migrated int8 tensor.

cc @sayakpaul

[jerryzh168]

repro 1 (pin memory):

import torch 
from torchao.quantization import quantize_, Int8WeightOnlyConfig 


module = torch.nn.Linear(3072, 3072)
quantize_(module, Int8WeightOnlyConfig())

for p in module.parameters():
    p.data.cpu().pin_memory()

Traceback (most recent call last):
  File "/fsx/sayak/diffusers/check_pin_torchao.py", line 9, in <module>
    p.data.cpu().pin_memory()
  File "/fsx/sayak/miniconda3/envs/diffusers/lib/python3.10/site-packages/torchao/utils.py", line 658, in _dispatch__torch_function__
    return func(*args, **kwargs)
  File "/fsx/sayak/miniconda3/envs/diffusers/lib/python3.10/site-packages/torchao/utils.py", line 678, in _dispatch__torch_dispatch__
    raise NotImplementedError(
NotImplementedError: AffineQuantizedTensor dispatch: attempting to run unimplemented operator/function: func=<OpOverload(op='aten.is_pinned', overload='default')>, types=(<class 'torchao.dtypes.affine_quantized_tensor.AffineQuantizedTensor'>,), arg_types=(<class 'torchao.dtypes.affine_quantized_tensor.AffineQuantizedTensor'>,), kwarg_types={}

[jerryzh168]

repro 2 (swap_tensor):

from diffusers import DiffusionPipeline, TorchAoConfig
from torchao.quantization import Int8WeightOnlyConfig
import torch


pipe = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image",
    quantization_config=TorchAoConfig(Int8WeightOnlyConfig()),
    torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()
if getattr(pipe.transformer, "_repeated_blocks"):
    pipe.transformer.compile_repeated_blocks(fullgraph=True)
else:
    pipe.transformer.compile()

# small resolutions to ensure speedy execution.
pipe("a dog", num_inference_steps=2, max_sequence_length=16, height=256, width=256)

[sayakpaul]

FYI, the pin memory issue happens with FP8, too.