Current best practices to initialize massive (50B parameter+) models

[azton]

Hi, I am working with GPT-style models and need to intitialize a model at the GPT-3 scale. Unfortunately, this means the model will run out of memory during initialization on CPU (or take an eternity to initialize layer-by-layer on cpu before shipping to GPU). In vanilla Pytorch I solved this using FSDP by initializing my models on the "meta" device, with full initialization on GPU afterward. What is the current best, most performant method to accomplish this with lightning?

Note: I found this, which references an init_meta_context(), but my pytorch-lightning (v1.9.0) has no such functionality:
https://devblog.pytorchlightning.ai/experiment-with-billion-parameter-models-faster-using-deepspeed-and-meta-tensors-2e9c255edd71

[carmocca]

Hi! The init_meta_context functionality was replaced with a torchdistx integration in #13868. You can do the following:

from torchdistx.deferred_init import deferred_init

model = deferred_init(YourLightningModule)

And we'll materialize it for you in the Trainer. This is very experimental, and you might encounter installation issues.

In the long term, we'll adopt the fake tensor mode from PyTorch: #16448.

Otherwise, for a stable(r) solution, you can use the DeepSpeed integration: https://pytorch-lightning.readthedocs.io/en/stable/advanced/model_parallel.html#deepspeed-zero-stage-3

[justusschock]

Hey, in addition to @carmocca response, for deepspeed you need to initialize the model within configure_sharded_model or you'll run in the same issues.

[carmocca]

Yes. This is described in https://pytorch-lightning.readthedocs.io/en/stable/advanced/model_parallel.html#shard-model-instantly-to-reduce-initialization-time-memory

[azton]

Thank you, I got this working with the Deepspeed integration--which was failing before because of my own errors. Just a small follow-up: does reloading the model from a checkpoint still function the same and utilize the configure_sharded_model logic to avoid memory overflow, or is there anything special that needs to happen on my end?