Improved model initialization API for Fabric

[awaelchli]

What does this PR do?

Adds a new context manager in Fabric that allows you to init your model:

• directly on the right device
• with the right device type
• getting sharded instantly
• allow for empty weight initialization (meta device) or deferred init (torchdistx) (Adopt FakeTensorMode for FSDP #16448)

This context manager genralized the previous sharded_model manager from the LightningLite-days (we don't need both).

Example 1:

Init the model in the GPU instantly and with weights in half precision (your model may not fit in float32).

fabric = Fabric(accelerator="cuda", precision="bf16-true")

with fabric.init_module():
    model = MyModel()

See #17287 for half support.

Example 2:
Init FSDP model on the meta device (or using torchdistx).

fabric = Fabric(accelerator="cuda", strategy="fsdp")

with fabric.init_module():
    model = MyModel()  # model params are now on meta device (no memory allocated)

model = fabric.setup(model)  # params get sharded and put on device

Example 3:
(Future) Init the model with empty weights explicitly (no memory allocated) if you later need to overwite by loading a checkpoint anyway:

fabric = Fabric(...)

with fabric.init_module(empty_weights=True):
    model = MyModel()
    model.load_state_dict(...)

This API can be used for #16448 or to support torchdistx.

cc @Borda @carmocca @justusschock @awaelchli

[carmocca]

Great job. Since you might prefer to address some of my comments in a follow-up, approving to unblock.

We should also add this to the fabric docs

[carmocca]

This might also be used in the future for things that arent modules, maybe we should choose a different name that doesnt explicitly state "module"

[lantiga]

init_module is fine IMO, but if we want to discuss names :-) I find that the name doesn't convey what is different from just initializing outside the context manager.

Maybe direct_init would be an alternative? I'm good with init_module though, it probably sounds less cryptic.

[awaelchli]

I was thinking mentioning the "module" in the name is safer because it expresses more precisely what we should do under this context manager. This hopefully helps prevent users from accidentally doing weird stuff under this context manager and getting in trouble, but of course, we can never completely prevent that.

[carmocca]

I would vote for something like efficient_init or similar

[carmocca]

For example, this context manager woudl be useful for non-modules too:

import torch
from lightning import Fabric

fabric = Fabric(accelerator="cuda", precision="64-true")
with fabric.efficient_init():
    x = torch.zeros(1)
    print(x.device)  # cuda
    print(x.dtype)  # 64

[awaelchli]

It is very important that this context manager be only used for model initialization. For everything else, the user should use fabric.device.

The main motivation for this is to load and shard large models efficiently, or to provide a convenient way to cast the model to the desired dtype for inference without code changes.

[lantiga]

Looks great

[lantiga]

There wouldn't be any impact for current Fabric end-users (not developers) right? If so, I agree we should remove.

[lantiga]

init_module is fine IMO, but if we want to discuss names :-) I find that the name doesn't convey what is different from just initializing outside the context manager.

Maybe direct_init would be an alternative? I'm good with init_module though, it probably sounds less cryptic.

[codecov]

Codecov Report

Merging #17462 (758fbd7) into master (d48ec08) will decrease coverage by 24%.
The diff coverage is 89%.

Additional details and impacted files

@@            Coverage Diff            @@
##           master   #17462     +/-   ##
=========================================
- Coverage      83%      59%    -24%     
=========================================
  Files         415      410      -5     
  Lines       31649    31594     -55     
=========================================
- Hits        26405    18708   -7697     
- Misses       5244    12886   +7642