Improve performance of `load_state_dict`
Author: woct0rdho
Pull request: yes

# What does this PR do?

We avoid executing `get_slice` unless `map_location == "meta"` to improve the performance when loading a model with a large number of tensors.

Even though we avoid the dtype check in Python, the dtype will be checked at https://github.com/huggingface/safetensors/blob/7d5af853631628137a79341ddc5611d18a17f3fe/bindings/python/src/lib.rs#L1186

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes #37887


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker
- vision models: @amyeroberts, @qubvel
- speech models: @eustlb
- graph models: @clefourrier

Library:

- flax: @gante and @Rocketknight1
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @zach-huggingface and @SunMarc
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @SunMarc @zach-huggingface
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber

Documentation: @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @Rocketknight1
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->
@Rocketknight1

Comments:
Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the `Ready for review` button (at the bottom of the PR page). This will assign reviewers and trigger CI.

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_37902). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

In my case the time to load [Unsloth's Qwen3-30B-A3B](https://huggingface.co/unsloth/Qwen3-30B-A3B-bnb-4bit) reduced from 15 min to 2 min. But loading a dense model of the same size takes only seconds, and only one core is used when loading, so maybe it's still worth considering how to further speed up this.

@woct0rdho it could definitely be worth optimizing this further. In the past, most models were either dense or MoEs with a small number of experts (e.g. Mixtral-8x7B), and so the slowdown here wasn't very important. However, I expect that models like Qwen3 or Deepseek-V3 will be more common in future, with a huge number of experts but only a small number of activated experts. These models combine high capacity with fast inference, which is very important for "reasoning" training pipelines that involve RL. On a fast SSD, a 30B model should not take 2 minutes to load, so we should consider parallel weight loading or other speedups! 

cc @narsil, is there a recommended way to speed up multiple calls to `f.get_tensor()`? Is there a method like `get_tensors()` that can efficiently get multiple tensors more quickly than repeatedly calling `get_tensor()`, or can we use multithreading if it releases the GIL?

Hi @woct0rdho, we're investigating the slow loading issue. Can you give us more details on the disc you're loading from? For example, is it a local NVMe drive or a remote mount?

@woct0rdho Are you running on a mounted network disk by any chance ? Mounted network disks do no play well with memory mapping, and every read incurs a quite large overhead (depending on how it's mounted you can see the sort of latencies you see here).

On a AWS NVMe the model loads in ~12s for me (4xL4).
It takes ~500ms to load everything on CPU, and the rest to move to CUDA.
Fwiw, it takes ~1-2 mn to download the entire model on that machine from S3 (with HF transfer)

I created a small repro script outside of transformers:

```python
from huggingface_hub import hf_hub_download

import torch

from safetensors.torch import load_file
import datetime

torch.zeros((2, 2)).cuda()  # Initialize cuda runtime to skip that in measurements

filenames = []
for i in range(16):
    rfilename = f"model-{i + 1:05d}-of-00016.safetensors"
    filename = hf_hub_download("Qwen/Qwen3-30B-A3B", rfilename)
    filenames.append(filename)

start = datetime.datetime.now()
all_data = {}
for i, filename in enumerate(filenames):
    data = load_file(filename, device="cpu")
    all_data.update(data)

print(f"CPU Load Took {datetime.datetime.now() - start}")

start = datetime.datetime.now()
all_data = {}
for i, filename in enumerate(filenames):
    data = load_file(filename, device=f"cuda:{i % 4}")
    all_data.update(data)

print(f"Took {datetime.datetime.now() - start}")
```

Results (4xL4)
```
CPU Load Took 0:00:00.466379
Took 0:00:11.799003
```

The machinery in transfomers, follows pretty closely what `safetensors` does internally (with some light additional bookkeeping, but it shouldn't be that bad).

Sidenote: HDD can also suffer from slow reads, and that can be fixed by changing `f.keys()` to `f.ordered_keys()` (because HDD prefer sequential reads and are much more sensitive to it than SSDs/NVMe).

Good point. I was running it on AutoDL (a cloud GPU provider). There is a folder `/root/autodl-tmp/` for storage. It does not look like an individual disk or partition mounted by NFS, but it may be somehow virtualized and I don't know the underlying physical setup.

Update: I modified your repro script a bit to load only 4 files on 1 GPU:
```python
#!/usr/bin/env python3

import os

os.environ["HF_HOME"] = "/root/autodl-tmp/.cache"

import datetime

import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

torch.zeros((2, 2)).cuda()  # Initialize cuda runtime to skip that in measurements

filenames = []
for i in range(4):
    rfilename = f"model-{i + 1:05d}-of-00016.safetensors"
    filename = hf_hub_download("Qwen/Qwen3-30B-A3B", rfilename)
    # rfilename = f"model-{i + 1:05d}-of-00004.safetensors"
    # filename = hf_hub_download("unsloth/Qwen3-30B-A3B-bnb-4bit", rfilename)
    filenames.append(filename)

start = datetime.datetime.now()
all_data = {}
for i, filename in enumerate(filenames):
    data = load_file(filename, device="cpu")
    all_data.update(data)

print(f"CPU Load Took {datetime.datetime.now() - start}")

start = datetime.datetime.now()
all_data = {}
for i, filename in enumerate(filenames):
    data = load_file(filename, device="cuda:0")
    all_data.update(data)

print(f"GPU Load Took {datetime.datetime.now() - start}")
```
The results are:
```
# Qwen/Qwen3-30B-A3B
CPU Load Took 0:00:00.142864
GPU Load Took 0:00:04.896957
# unsloth/Qwen3-30B-A3B-bnb-4bit
CPU Load Took 0:00:02.586595
GPU Load Took 0:00:12.870413
```