Large Model Initialization with FSDP

[jthomy]

With strategies like FSDP: If I have a model that is too large to fit into CPU memory (especically not to fit n_gpus many times into memory), just instantiating in the configure_model() hook will run out of memory.

What is lightning's intended way to initialize the model such that I only load the full model once overall? Optimally I would do this directly on the GPU to have faster startup times. I would load the model from a non-distributed checkpoint, but happy to head any suggestions.

[llllvvuu]

related:
#17043
pytorch/pytorch#104755
pytorch/pytorch#104026
https://huggingface.co/docs/transformers/v4.51.3/en/models#big-model-inference