[wip][core] parallel loading of shards

[sayakpaul]

What does this PR do?

Similar to huggingface/transformers#36835.

`main`: time: 8.162s
this branch: time: 5.663s

code

import time
t_ini = time.time()

import torch
import os
from diffusers import DiffusionPipeline, AutoModel
print(f"import time: {time.time() - t_ini:.3f}s")

os.environ["HF_ENABLE_PARALLEL_LOADING"] = "YES"
os.environ["HF_PARALLEL_LOADING_WORKERS"] = "12"
model_id = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"

t0 = time.time()
torch.cuda.synchronize()
print(f"CUDA sync time: {time.time() - t0:.3f}s")

print("starting model load")
t1 = time.time()
transformer = AutoModel.from_pretrained(
    model_id, subfolder="transformer", torch_dtype=torch.bfloat16, device_map="cuda"
)
torch.cuda.synchronize()
t2 = time.time()

diff = t2 - t1
print(f"time: {diff:.3f}s")

[sayakpaul]

Moved it here from modeling_utils.py.

[sayakpaul]

Same. Moved it out of modeling_utils.py.

[sayakpaul]

This has been moved to load_shard_file().

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