Fix serialization for offloaded model

[SunMarc]

What does this PR do ?

This PR fixes the serialization of offloaded model. With this PR, the logic behind saving offloaded model is triggered only when we use more than one device that contain either "cpu" or "disk" device. Moreover, we fix a small bug that was introduced with this PR. We rename the state_dict variable to shard_state_dict. I also fixed the test test_save_offloaded_model.

Example:

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", device_map={"":"cpu"})
model.save_pretrained("llama")

Fixes #31685

[LysandreJik]

Could you implement a test that ensures that the previous faulty behavior has been fixed? :)

[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]

Could you implement a test that ensures that the previous faulty behavior has been fixed? :)

Done !

Do you know where I can see the results of the test for test_modeling_utils.py file @ydshieh ? I didn't see the following test test_save_offloaded_model in the daily report. On my local machine, it was failing so I was expecting to see it in the report.

[ydshieh]

Hi @SunMarc Let me check, but we might have a problem of some test files not included 😭

[SunMarc]

Hi @SunMarc Let me check, but we might have a problem of some test files not included 😭

Sounds good ! Maybe we can put all the utils python file into a separate folder to not have them mixed with the common test files.

[ydshieh]

I confirmed that we have some tests not running on daily CI.
(We also have push-ci where we try to collect the impacted files, but for daily CI, tests are collected from the subfolders of tests)

Maybe we can put all the utils python file into a separate folder to not have them mixed with the common test files.

Sound a quick solution. It's just we have to make sure there is not subclass of TestCase in other files under tests. But for now let's do what you suggested above.

[SunMarc]

Sound a quick solution. It's just we have to make sure there is not subclass of TestCase in other files under tests. But for now let's do what you suggested above.

Great !

[ydshieh]

Here it is #31730 but I have to see how CI is going.

[LysandreJik]

Looks good, thanks for the changes @SunMarc