is torchao compatible with cpu offload?

[zhangvia]

i run the script below:

from diffusers import FluxFillPipeline
from torchao.quantization import autoquant


import torch
p = FluxFillPipeline.from_pretrained("FLUX.1-Fill-dev/",torch_dtype=torch.bfloat16)
p.enable_model_cpu_offload()
autoquant(p.transformer)



from PIL import Image
image = p(prompt="a girl",image = Image.new('RGB',(3072,2048)),mask_image=Image.new('L',(3072,2048)),height=2048,width=3072)

but get this error in model.to():

Traceback (most recent call last):
  File "/media/74nvme/research/test.py", line 322, in <module>
    image = p(prompt="a girl",image = Image.new('RGB',(3072,2048)),mask_image=Image.new('L',(3072,2048)),height=2048,width=3072)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_fill.py", line 916, in __call__
    noise_pred = self.transformer(
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1844, in _call_impl
    return inner()
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1769, in inner
    args_kwargs_result = hook(self, args, kwargs)  # type: ignore[misc]
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torchao/quantization/autoquant.py", line 1314, in autoquant_prehook
    real_model.forward(*args, **kwargs)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/accelerate/hooks.py", line 179, in new_forward
    args, kwargs = module._hf_hook.pre_forward(module, *args, **kwargs)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/accelerate/hooks.py", line 714, in pre_forward
    module.to(self.execution_device)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/diffusers/models/modeling_utils.py", line 1077, in to
    return super().to(*args, **kwargs)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1340, in to
    return self._apply(convert)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 900, in _apply
    module._apply(fn)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 900, in _apply
    module._apply(fn)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 900, in _apply
    module._apply(fn)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 928, in _apply
    p_should_use_set_data = compute_should_use_set_data(param, param_applied)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 903, in compute_should_use_set_data
    if torch._has_compatible_shallow_copy_type(tensor, tensor_applied):
TypeError: _has_compatible_shallow_copy_type(): argument 'from' (position 2) must be Tensor, not NoneType

[sayakpaul]

What happens when you quantize first and then CPU offload?

[sayakpaul]

I think autoquant() needs a GPU to work.

[zhangvia]

What happens when you quantize first and then CPU offload?

same error happen

Traceback (most recent call last):
  File "/media/74nvme/research/test.py", line 322, in <module>
    image = p(prompt="a girl",image = Image.new('RGB',(3072,2048)),mask_image=Image.new('L',(3072,2048)),height=2048,width=3072)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_fill.py", line 916, in __call__
    noise_pred = self.transformer(
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1844, in _call_impl
    return inner()
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1769, in inner
    args_kwargs_result = hook(self, args, kwargs)  # type: ignore[misc]
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torchao/quantization/autoquant.py", line 1314, in autoquant_prehook
    real_model.forward(*args, **kwargs)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/accelerate/hooks.py", line 179, in new_forward
    args, kwargs = module._hf_hook.pre_forward(module, *args, **kwargs)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/accelerate/hooks.py", line 714, in pre_forward
    module.to(self.execution_device)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/diffusers/models/modeling_utils.py", line 1077, in to
    return super().to(*args, **kwargs)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1340, in to
    return self._apply(convert)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 900, in _apply
    module._apply(fn)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 900, in _apply
    module._apply(fn)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 900, in _apply
    module._apply(fn)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 928, in _apply
    p_should_use_set_data = compute_should_use_set_data(param, param_applied)
  File "/media/74nvme/software/miniconda3/envs/repose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 903, in compute_should_use_set_data
    if torch._has_compatible_shallow_copy_type(tensor, tensor_applied):
TypeError: _has_compatible_shallow_copy_type(): argument 'from' (position 2) must be Tensor, not NoneType

I think autoquant() needs a GPU to work.

i know that, but just like optimum-quanto, i can use quantization and cpu offload to make flux run on 24GB gpu even with 2408*3072 input. and if i can use torchao and cpu offload simultaneously, i can try more optmizations like torch.compile or something else

[sayakpaul]

Then I would suggest running autoquant first on a GPU which should give you an idea about the quantization plan it is picking up for your model. Then, you can use a similar config through the torchao support we provide in diffusers:
https://huggingface.co/docs/diffusers/main/en/quantization/torchao