Deepspeed sharding and load from checkpoint with custom lightning module - setup() not called during checkpoint loading about lightning-transformers

[uakarsh]

Hi @maxzvyagin, what I can understand from your question is, performing the enable_, on the custom Lightning Module. For that, I think one simple and straight strategy would be inheriting the TaskTransformer class and modifying the initialize_model method, as mentioned here.

Do correct, me if I did not understand your question correctly.

Regards,
Akarsh

[maxzvyagin]

Hi Akarsh, thanks for checking this out! I guess my question is partially what the reason is for calling enable_transformers_pretrained_deepspeed_sharding(self) in the setup() function vs in the constructor, and if there's a detrimental effect by calling it in the constructor if we're loading from a pre-trained checkpoint. Because otherwise, we're not able to load from our own PyTorch checkpoint file and continue training with a sharded DeepSpeed approach.