[feat] cache allocator warmup for from_single_model

[sayakpaul]

What does this PR do?

fromdiffusersimportFluxTransformer2DModelimporttorchimporttimeckpt_path="https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/flux1-dev.safetensors"start=time.perf_counter()
transformer=FluxTransformer2DModel.from_single_file(
ckpt_path, torch_dtype=torch.bfloat16, device_map="cuda"
)
# transformer = FluxTransformer2DModel.from_single_file(# ckpt_path, torch_dtype=torch.bfloat16,# ).to("cuda")torch.cuda.synchronize()
end=time.perf_counter()
print(f"Take taken to initialize the model on CUDA: {(end-start)} seconds.")

Timing:

CUDA placement manual: 4.86 seconds
Device map: 1.09 seconds

(On an H100)

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

[sayakpaul]

@stevhliu it would be awesome to also add a doc entry for this.

@DN6 test needed to be added this PR?

[stevhliu]

Added a note about it in #12256 👍

[sayakpaul]

Included a test too.