[model loading] don't gc.collect() if only 1 shard is used

[gante]

What does this PR do?

Fixes a test-related regression introduced by #36033. Despite speeding up from_pretrained in general, most tests that relied on from_pretrained became much slower (see tests section below).

Measuring time, the entire slowdown can be traced to a single line, a gc.collect(). This line was not added in #36033, it was moved. Before #36033, if the checkpoint was not sharded, a state_dict would have been passed to _load_pretrained_model, and the gc.collect() branch would not be reached.

This PR adds a tiny if to skip gc.collect() if the checkpoint is not sharded. Since many tests rely on unsharded from_pretrained, we can immediately observe faster tests.

Tests

py.test tests/models/gpt2/test_modeling_gpt2.py times, which includes a mix of tests with and without from_pretrained, on my machine:

• commit before [core] Large/full refactor of from_pretrained #36033: 25.41s
• main: 61.14s
• this PR: 21.78s 💪 (this means [core] Large/full refactor of from_pretrained #36033 resulted in 15%+ test speedup, after this fix)

[github-actions]

Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. When it is ready for review, please click the Ready for review button (at the bottom of the PR page).

[Cyrilvallez]

LGTM! Thanks!

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