[loading] Really initialize on meta device for huge perf gains

[Cyrilvallez]

What does this PR do?

Follow-up to #42309 to really leverage meta device loading. Gives crazy speedups for loading some models, e.g. about 2.5x on gpt-oss 20b and about 3x on the 120b version

The issue at hand

Currently, during loading we initialize the model on meta device before loading weights, thanks to init_empty_weights from accelerate.
However, this context manager has BIG drawbacks:

• everything (every parameter and buffer) is first materialized on cpu before being moved to meta device
• this is extremely inefficient of course, as we only want them on meta -> it wastes time and memory
• all buffers stay on cpu (even the persistent ones, that we are loading again after anyway, so don't need to be there...)

For some models, e.g. gpt-oss, we have the following during loading:

Note how most of the loading time is BEFORE the actual loading of the weights (_load_pretrained call), just to initialize parameters that should be on meta device anyway....

What this PR is doing

This PR completely removes init_empty_weights in favor of torch.device("meta") to really start with a model on meta device, without first putting them on cpu.
We are free to do so since I've merged #42309 yesterday, to correctly handle re-initialization of the non-persistent buffers which are put on meta device as well.

The performance gains are immense, the same benchmark as before for gpt-oss shows it:

and we can now see that most of the time in from_pretrained is used for the actual weight loading (_load_pretrained call), as it should be.

Raw numbers for the following simple benchmark script (from which the above traces are from) on our cluster:

from transformers import AutoModelForCausalLM
import torch
import time
from viztracer import VizTracer

model_id = "openai/gpt-oss-20b"
device = 0

tracer = VizTracer()
tracer.start()

torch.cuda.synchronize()
t0 = time.time()
model = AutoModelForCausalLM.from_pretrained(model_id, device_map=device, dtype=torch.float16)
torch.cuda.synchronize()
dt = time.time() - t0
print(f"Took {dt:.2f} s")

tracer.stop()
tracer.save(output_file="../trace.json")

are the following:

BEFORE THIS PR -> ~9.1s
ON THIS PR -> ~3.7s

Which means a speedup of about 2.5x.

For the 120B gpt-oss version, we have:

BEFORE THIS PR -> ~32.4s
ON THIS PR -> ~10.7s

or a speedup of about 3x

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

[Cyrilvallez]

cc @SunMarc for this change, do you know if this switch on all the quantizers's replace_with_xxx is fine? Basically, the only difference is if the quantized layer registers some buffers, they would now be on meta as well. I checked for bnb and there it seems to be alright at least (no buffers)

[SunMarc]

as long as it's not a non persistant buffer, it should be fine !

[vadimkantorov]

Could be related minor issue:

• AutoModel.from_pretrained(...) (with explicit device_map unset) fails under with torch.device("meta") with PyTorch 2.6.0 and 2.7.0 #38066

[ArthurZucker]

An early Christmas gift for everyone

[ArthurZucker]

💌

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: chameleon, codegen, ctrl, deepseek_vl, deepseek_vl_hybrid, emu3, eomt, fastspeech2_conformer, gemma, gemma2, gptj, idefics2, idefics3, janus, layoutlmv3, llava_next