how to use device_map in diffusers.stablediffusionpipeline.from_pretrained

[chaowenguo]

❓ Questions and Help

https://huggingface.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map

i try device_map=auto. but it seems not recongize the xla device tpu all goto cpu. please help and provide the minimum working example. @JackCaoG

[zpcore]

The feature sounds like host offloading. It's currently in experimental stage by @tengyifei .

[chaowenguo]

@zpcore @tengyifei could you give me more detailed example？how to use host offloading？

[miladm]

You should be able to find an example of this implementation here: https://github.com/AI-Hypercomputer/tpu-recipes/tree/main/training/trillium/Llama3-405B-PyTorch

On a related note, we are working on rolling out a much cleaner reference model.