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

[DN6]

Would add HF_PARALLEL_LOADING_WORKERS as a constant at the top of the file for consistency.

[DN6]

Since the same arguments are used across the two loading functions, it's a good candidate for functools.partial

        load_fn = partial(
            load_shard_files_with_threadpool if is_parallel_loading_enabled else load_shard_file,
            model=model,
            model_state_dict=model_state_dict,
            device_map=device_map,
            dtype=dtype,
            hf_quantizer=hf_quantizer,
            keep_in_fp32_modules=keep_in_fp32_modules,
            dduf_entries=dduf_entries,
            loaded_keys=loaded_keys,
            unexpected_keys=unexpected_keys,
            offload_index=offload_index,
            offload_folder=offload_folder,
            state_dict_index=state_dict_index,
            state_dict_folder=state_dict_folder,
            ignore_mismatched_sizes=ignore_mismatched_sizes,
            low_cpu_mem_usage=low_cpu_mem_usage,
        )

        if is_parallel_loading_enabled:
            offload_index, state_dict_index, _mismatched_keys, _error_msgs = load_fn(
                resolved_model_file,
            )
            error_msgs += _error_msgs
            mismatched_keys += _mismatched_keys
        else:
            shard_files = resolved_model_file
            if len(resolved_model_file) > 1:
                shard_files = logging.tqdm(resolved_model_file, desc="Loading checkpoint shards")

            for shard_file in resolved_model_file:
                offload_index, state_dict_index, _mismatched_keys, _error_msgs = load_fn(shard_file)
                error_msgs += _error_msgs
                mismatched_keys += _mismatched_keys

[sayakpaul]

@stevhliu, could you help add docs for this PR (separate PR is fine)? I think we could have some guidance on how to load a DiffusionPipeline faster when the hardware allows for it (like directly loading all model components on the accelerator).

#11904 could also be mentioned in the document.

Then we're working on #12122

[sayakpaul]

@DN6 thanks a lot for your thoughtful suggestions. I have reflected them and I have added a test case, as well.

LMK what you think.

[DN6]

Thanks!

[DN6]

I meant to run the env check here

HF_ENABLE_PARALLEL_LOADING  = os.environ.get("HF_ENABLE_PARALLEL_LOADING", "").upper() in ENV_VARS_TRUE_VALUES

Then import the constant into modeling_utils.

[sayakpaul]

PR: #12137