[performance] fusing zero.Init and pretrained weights loading

I'm thinking there should be some kind of fused zero.Init and pretrained weights model loading function.

Currently with zero.Init we:

1. partition the model and scatter the weights and then for each layer:
2. gather the weights of one layer
3. load weights of that layer
4. scatter the weights of that layer

So gather+scatter happens 3 times, rather than 1.

this feels wasteful and probably does add quite an overhead for large models.

So perhaps there should be a way to hook torch.load into zero.Init. I haven't thought it through, but wanted to start the discussion.

edit: pt-1.9.0/1.9.1 now supports postponed module init where no storage is materialized until it's needed:
https://pytorch.org/tutorials/prototype/skip_param_init.html#implementation-details
So if this is used the solution can be made much simpler by not needing zero.Init (which forces the first wasteful scatter+gather) and perform the first scattering during load_state_dict.

checkpoint loading: need to investigate this one too in the same context of allocating and scattering weights only once (probably need a separate issue to track this)

Here is the corresponding issue on the HF side: https://github.com/huggingface/transformers/issues/12273

Thank you.

Start by reading the zero.Init flow alongside torch.load and load_state_dict, then trace the current gather, weight loading, and scatter sequence described in the issue. Compare the proposed path with PyTorch's postponed module initialization. Done means the design and implementation avoid redundant weight movement during pretrained loading, with checkpoint loading considered separately as noted.