Add meta_init, enable it as default init process

[lessw2020]

This PR enables meta_init functionality to avoid OOM'ing on cpu for larger models.
The core functionality is in meta_init.py, and a few changes in parallelization and train.py.
Key items:
1 - this is largely the same as the earlier PR I had for meta_init, but I did a new one b/c faster than reworking it with all the interim changes.
2 - to address feedback in previous PR:
a - why do we need meta_init.py, can't we just do:

with torch.device("meta"):
    model = Model.from_args(...)

Unfortunately this does not work b/c the rope embeddings are treated differently (buffer) and thus the simple lambda call from param_init_fn in FSDP (lambda module: module.to_device('cuda') ) will not invoke or move the rope embeddings and the model will fail on first forward.
This issue relates to the nn.embeddings not being moved, and that the device is referenced in the forward pass for the current rope class. Have opened #110 to track this and investigate while not holding up meta init that is working from landing.

b - per earlier feedback - meta init is now 'not optional' but simply the default. This should ensure all models leverage it and ensure we aren't missing things for future meta_init aspects.

3 - misc change - I switched the model_params to just do the normal all params count instead of 'unique params' b/c it does not mesh with what people perceive model size as.

Testing:
tested both debugmodel and 26B model with and without meta init to confirm same loss curves.
Note for future reference - if you get a bad init (meta init failure) you will simply not train (loss is same every iter).
If you fail to call reset params after FSDP, then you will train (b/c we default to torch.randn_like) but your starting loss will be 5x+ higher (telling you that you have not properly init'ed the model).

[wanchaol]

thus the simple lambda call from param_init_fn in FSDP (lambda module: module.to_device('cuda') ) will not invoke or move the rope embeddings

@lessw2020 curious is this because the buffer are created in the meta device however the module.to_device('cuda') won't move the buffer to GPU? iirc model.to("cuda") would move both params and buffers to cuda https://github.com/pytorch/pytorch/blob/main/torch/nn/modules/module.py#L848

[tianyu-l]

if you get a bad init (meta init failure) you will simply not train (loss is same every iter)

Would it be possible to just error out if meta init fails?

[wconstab]

i actually kinda like the separation of device-changes from parallelism application.

that leads me to think its nice to keep .cuda() or .to() or .reset_parameters() out of parallelize_llama() helper, and explicitly in the train loop. but i could be convinced otherwise.

[lessw2020]

thus the simple lambda call from param_init_fn in FSDP (lambda module: module.to_device('cuda') ) will not invoke or move the rope embeddings

@lessw2020 curious is this because the buffer are created in the meta device however the module.to_device('cuda') won't move the buffer to GPU? iirc model.to("cuda") would move both params and buffers to cuda https://github.com/pytorch/pytorch/blob/main/torch/nn/modules/module.py#L848

It's deeper than that - I've opened an issue for tracking this to resolve and thus simplify the whole meta init:
#110
For now, would like to land a working meta_init and then refine with the fix for rope embeddings.

[lessw2020]

re: meta init simplification - This issue relates to the nn.embeddings not being moved, and that the device is referenced in the forward pass for the current rope class. Have opened #110 to track this and investigate while not holding up meta init that is working from landing.

[tianyu-l]

LGTM. Thanks for enabling meta init! Great to see this it's working well. Had one inline comment.

Let's follow up on #110 to see if we can get a simpler solution.

[awgu]

I apologize that FSDP meta-device init is confusing, but I think this might not be fully correct.

1. The contract from PyTorch core is supposed to be that module.reset_parameters() only resets/initializes the parameters immediately owned by module, not those of its children/submodules. Here, model: Transformer is initializing all of its submodules' parameters. This contract is the only way for reset_parameters() to always work compositionally, otherwise (like on our case) we must assume Transformer to always be the root module.
2. When we call model.reset_parameters(), the parameters have already been flattened and sharded by FSDP. This means that any initialization method that depends on the tensor shape would be incorrect. That is why we would normally want users to do the correct initialization in the param_init_fn.

For this Llama case, it looks like perhaps the reset_parameters() have been written to not depend on the tensor shape directly?

[lessw2020]

correct that the init does not depend on tensor shape directly.
I see your point though about only resetting it's own params and not children, but let's meet to discuss as I also have questions on how meta_init should work for FSDP2 and we can get an updated impl.