Support empty weight initialization in Fabric.init_module()

[awaelchli]

What does this PR do?

Fixes #17616

Adds a toggle on the Fabric.init_module that speeds up initialization and memory allocation for a large model.

with fabric.init_module(empty_weights=True):
    # it is very fast to initialize, and depending on the strategy allocates no memory, or uninitialized memory
    model = MyModel()
# weights get loaded into the model
model.load_state_dict(checkpoint["state_dict"])

Useful for finetuning / loading weights into a large model.

See how Fabric.init_module(empty_weights=True) can be applied in lit-llama to minimize boilerplate logic: Lightning-AI/lit-llama#360

cc @Borda @carmocca @justusschock @awaelchli

[codecov]

Codecov Report

Merging #17627 (0c7fda9) into master (420eb6f) will decrease coverage by 23%.
The diff coverage is 79%.

❗ Current head 0c7fda9 differs from pull request most recent head 4d71cdd. Consider uploading reports for the commit 4d71cdd to get more accurate results

Additional details and impacted files

@@            Coverage Diff            @@
##           master   #17627     +/-   ##
=========================================
- Coverage      84%      61%    -23%     
=========================================
  Files         419      415      -4     
  Lines       31721    31662     -59     
=========================================
- Hits        26634    19382   -7252     
- Misses       5087    12280   +7193     

[carmocca]

Implementation looks good.

I would personally go for empty_init.

[carmocca]

@awaelchli We might want to reconsider doing empty_init=True by default on FSDP. I just went down a rabbit hole chasing NaNs on FSDP when all I needed to change was empty_init=False. Since the default was to skip initialization and I wasn't loading a full checkpoint, the weights were initialized to garbage.

Since there's no way for us to know if the user will load a full checkpoint after initialization, I find it safer to not do this by default.

context: Lightning-AI/litgpt#193

[awaelchli]

That wasn't the intention, sorry. The plan was to do this automatically if we have fake tensor / support for materialization (torchdistx-style). In the meantime, it would be safer to do:

Suggested change

	_EmptyInit(enabled=(empty_init is not False)) if _TORCH_GREATER_EQUAL_1_13 else nullcontext()
	_EmptyInit(enabled=bool(empty_init) if _TORCH_GREATER_EQUAL_1_13 else nullcontext()

[awaelchli]

The comment above there

    # TODO: Use the meta device and reset parameters after https://github.com/pytorch/pytorch/issues/90465

is what I ultimately wanted to achieve

[robieta]

Same. While it's a cool feature, having this default to true is very disruptive and hard to debug.