Performance of `load_state_dict` with large number of tensors (Qwen3 MoE)
Author: woct0rdho
Pull request: no
Labels: name: Feature request

### Feature request

It takes a long time to load a model with a large number of tensors. Can we improve its performance?

### Motivation

For example, [Unsloth's Qwen3-30B-A3B](https://huggingface.co/unsloth/Qwen3-30B-A3B-bnb-4bit) has 31337 tensors in a file, because there are 128 experts in each layer, and it takes more than 15 minutes to load it on my computer, which significantly slows down the development around this model.

### Your contribution

Most of the time is spent in this loop: https://github.com/huggingface/transformers/blob/481de7204c6e065616d4f848ea7f69a2287727df/src/transformers/modeling_utils.py#L509

A simple improvement is to move
```python
k_dtype = f.get_slice(k).get_dtype()
if k_dtype in str_to_torch_dtype:
    dtype = str_to_torch_dtype[k_dtype]
else:
    raise ValueError(f"Cannot load safetensors of unknown dtype {k_dtype}")
```
into `if map_location == "meta":` so it's not executed when loading the weights normally. This reduces the loading time from 15 min to 2 min for me. But loading a dense model of the same size takes only seconds.

Then maybe it boils down to the performance of `get_tensor` but I can't see a simple way to improve it.

Even if the performance cannot be improved, it feels better if we add `logging.tqdm` to this loop.

Comments:
Hi @woct0rdho, can you tell me what version of `transformers` you're using? We made some updates to this code to improve performance recently!

I'm using transformers 4.51.3

<details>
<summary>Full results of pip list</summary>

```
absl-py                        2.1.0
accelerate                     1.6.0
aiohappyeyeballs               2.4.6
aiohttp                        3.11.13
aiosignal                      1.3.2
anaconda-anon-usage            0.4.4
annotated-types                0.7.0
anyio                          4.6.2.post1
archspec                       0.2.3
argon2-cffi                    23.1.0
argon2-cffi-bindings           21.2.0
arrow                          1.3.0
asttokens                      2.4.1
async-lru                      2.0.4
attrs                          24.2.0
babel                          2.16.0
beautifulsoup4                 4.12.3
bitmath                        1.3.3.1
bitsandbytes                   0.45.5
bleach                         6.2.0
boltons                        23.0.0
Brotli                         1.0.9
cachetools                     5.5.2
certifi                        2024.2.2
cffi                           1.16.0
chardet                        4.0.0
charset-normalizer             2.0.4
click                          8.1.8
comm                           0.2.2
conda                          24.4.0
conda-content-trust            0.2.0
conda-libmamba-solver          24.1.0
conda-package-handling         2.2.0
conda_package_streaming        0.9.0
contourpy                      1.3.1
cryptography                   42.0.5
cut-cross-entropy              25.1.1
cycler                         0.12.1
datasets                       3.5.1
debugpy                        1.8.9
decorator                      5.1.1
defusedxml                     0.7.1
deprecation                    2.1.0
diffusers                      0.33.1
dill                           0.3.8
distro                         1.9.0
docker-pycreds                 0.4.0
docstring_parser               0.16
executing                      2.1.0
fastjsonschema                 2.20.0
filelock                       3.16.1
fonttools                      4.55.0
fqdn                           1.5.1
frozenlist                     1.5.0
fsspec                         2024.10.0
gitdb                          4.0.12
GitPython                      3.1.44
grpcio                         1.68.0
h11                            0.14.0
hbutils                        0.11.0
hf_transfer                    0.1.9
hf-xet                         1.1.0
hfutils                        0.10.0
httpcore                       1.0.7
httpx                          0.27.2
huggingface-hub                0.30.2
idna                           3.7
importlib_metadata             8.6.1
ipykernel                      6.29.5
ipython                        8.29.0
ipywidgets                     8.1.5
isoduration                    20.11.0
jedi                           0.19.2
Jinja2                         3.1.4
json5                          0.9.28
jsonpatch                      1.33
jsonpointer                    2.1
jsonschema                     4.23.0
jsonschema-specifications      2024.10.1
jupyter_client                 8.6.3
jupyter_core                   5.7.2
jupyter-events                 0.10.0
jupyter-lsp                    2.2.5
jupyter_server                 2.14.2
jupyter_server_terminals       0.5.3
jupyterlab                     4.3.1
jupyterlab-language-pack-zh-CN 4.2.post3
jupyterlab_pygments            0.3.0
jupyterlab_server              2.27.3
jupyterlab_widgets             3.0.13
kiwisolver                     1.4.7
libmambapy                     1.5.8
Markdown                       3.7
markdown-it-py                 3.0.0
MarkupSafe                     3.0.2
matplotlib                     3.9.2
matplotlib-inline              0.1.7
mdurl                          0.1.2
menuinst                       2.0.2
mistune                        3.0.2
mpmath                         1.3.0
multidict                      6.1.0
multiprocess                   0.70.16
natsort                        8.4.0
nbclient                       0.10.0
nbconvert                      7.16.4
nbformat                       5.10.4
nest-asyncio                   1.6.0
networkx                       3.4.2
notebook_shim                  0.2.4
numpy                          2.2.5
nvidia-cublas-cu12             12.4.5.8
nvidia-cuda-cupti-cu12         12.4.127
nvidia-cuda-nvrtc-cu12         12.4.127
nvidia-cuda-runtime-cu12       12.4.127
nvidia-cudnn-cu12              9.1.0.70
nvidia-cufft-cu12              11.2.1.3
nvidia-curand-cu12             10.3.5.147
nvidia-cusolver-cu12           11.6.1.9
nvidia-cusparse-cu12           12.3.1.170
nvidia-cusparselt-cu12         0.6.2
nvidia-nccl-cu12               2.21.5
nvidia-nvjitlink-cu12          12.4.127
nvidia-nvtx-cu12               12.4.127
overrides                      7.7.0
packaging                      23.2
pandas                         2.2.3
pandocfilters                  1.5.1
parso                          0.8.4
peft                           0.15.2
pexpect                        4.9.0
pillow                         11.0.0
pip                            25.1
platformdirs                   3.10.0
pluggy                         1.0.0
prometheus_client              0.21.0
prompt_toolkit                 3.0.48
propcache                      0.3.0
protobuf                       3.20.3
psutil                         6.1.0
ptyprocess                     0.7.0
pure_eval                      0.2.3
pyarrow                        20.0.0
pycosat                        0.6.6
pycparser                      2.21
pydantic                       2.10.6
pydantic_core                  2.27.2
Pygments                       2.18.0
pyparsing                      3.2.0
PySocks                        1.7.1
python-dateutil                2.9.0.post0
python-json-logger             2.0.7
pytimeparse                    1.1.8
pytz                           2025.1
PyYAML                         6.0.2
pyzmq                          26.2.0
random-user-agent              1.0.1
referencing                    0.35.1
regex                          2024.11.6
requests                       2.32.3
rfc3339-validator              0.1.4
rfc3986-validator              0.1.1
rich                           13.9.4
rpds-py                        0.21.0
ruamel.yaml                    0.17.21
safetensors                    0.5.3
Send2Trash                     1.8.3
sentencepiece                  0.2.0
sentry-sdk                     2.22.0
setproctitle                   1.3.5
setuptools                     75.8.2
shtab                          1.7.1
six                            1.16.0
smmap                          5.0.2
sniffio                        1.3.1
soupsieve                      2.6
stack-data                     0.6.3
supervisor                     4.2.5
sympy                          1.13.1
tensorboard                    2.18.0
tensorboard-data-server        0.7.2
terminado                      0.18.1
tinycss2                       1.4.0
tokenizers                     0.21.1
torch                          2.6.0
torchvision                    0.21.0
tornado                        6.4.2
tqdm                           4.67.1
traitlets                      5.14.3
transformers                   4.51.3
triton                         3.2.0
trl                            0.15.2
truststore                     0.8.0
typeguard                      4.4.2
types-python-dateutil          2.9.0.20241003
typing_extensions              4.12.2
tyro                           0.9.16
tzdata                         2025.1
tzlocal                        5.3
unsloth                        2025.4.3
unsloth_zoo                    2025.4.2
uri-template                   1.3.0
urllib3                        2.1.0
URLObject                      2.4.3
wandb                          0.19.10
wcwidth                        0.2.13
webcolors                      24.11.1
webencodings                   0.5.1
websocket-client               1.8.0
Werkzeug                       3.1.3
wheel                          0.43.0
widgetsnbextension             4.0.13
xformers                       0.0.29.post3
xxhash                         3.5.0
yarl                           1.18.3
zipp                           3.21.0
zstandard                      0.22.0
```
</details>

Checked it out and you're right, this is still quite slow! I was thinking of #35493, which was a similar problem but touching different code.

The solution to move the dtype check into `if map_location == "meta":` makes sense, but isn't a complete solution because the meta device is often used. Instead, we probably need to speed this up, which means either speeding up `get_tensor()` or calling it in parallel.

cc @Narsil @ArthurZucker, is there anything you can think of to improve safetensors loading here? Multiprocessing is out (I don't even want to think about the memory access issues there), but loading with multithreading might be possible, particularly if `get_tensor()` releases the GIL while it's loading. That should improve load performance a lot on fast storage, WDYT?

Actually, thinking about it more, because we only do tensor assignment when `map_location != "meta"`, the meta path will generally be faster anyway, so moving things into the `if map_location == "meta":` block is probably an important change we should do quickly. @woct0rdho would you be willing to make that PR?

Ok see #37902

I have done a rough implementation too, do let me know if it works


Hi @demoncoder-crypto, thank you but we can only accept a single PR for this issue! I'm going to focus on #37902 since it was opened first by the issue author.

Understood @Rocketknight1, Thanks for reviewing it.


thanks both for reporting and helping! 🤗