Fix sharding when no device_map is passed

[SunMarc]

What does this PR do?

This PR fixes the loading for sharded checkpoint when no device_map is passed. Currently, the following doesn't work:

from diffusers import UNet2DConditionModel, StableDiffusionXLPipeline 
import torch

unet = UNet2DConditionModel.from_pretrained(
    "sayakpaul/sdxl-unet-sharded", torch_dtype = torch.float16
)

You can have more details here.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[SunMarc]

There is still a path where sharding is not handled. It happens when low_cpu_mem_usage=False. I see that by default, low_cpu_mem_usage is set to True, it is the case for most models ? cc @sayakpaul

[yiyixuxu]

thank you!
very nice tests too:)

is it possible to explain device_map=None in the doc string for device_map too?

[SunMarc]

is it possible to explain device_map=None in the doc string for device_map too?

Done !

[sayakpaul]

This test is already here:

diffusers/tests/models/test_modeling_common.py

Line 874 in a899e42

def test_sharded_checkpoints(self):

Is it different?

[yiyixuxu]

he renamed this test to test_sharded_checkpoints_device_map because in that test it loads with device_map='auto' flag; this is a new test testing default value for device_map

[sayakpaul]

Thanks for explaining.

[SunMarc]

Yes, I renamed the tests since it makes more sense this way

[sayakpaul]

Thanks so much, Marc. I think there's some confusion in the tests as they are existing in the main already. Am I missing out on something?

[sayakpaul]

Alright then! Let’s merge this.

[sayakpaul]

Thanks!