switch to new meta device init method

[lchu-ibm]

This PR makes a new and improved version of model init which is more efficient.

Full details: #64

Performance on model init time (model creation + fsdp wrapping)

We have benchmarked the model init performance on 16 aws p4de nodes:

	old implementation	new implementation
7b	100s	0.2s
13b	193s	0.3s
34b	550s	0.6s
70b	1146s	1.5s

Validation test

Since unittest is hard to capture true-multi-node runs, we have ran external tests to make sure the new implementation pass validation test and yields true init. Specifically, we added the following plugin after FSDP call for testing purpose

with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, FullStateDictConfig(offload_to_cpu=True, rank0_only=True)):
    model_state = model.state_dict()

if rank == 0:
    model2 = LLaMA(llama_config)
    model2.load_state_dict(model_state)

    model2.check_parameters()

and check_parameters() is defined here

[nairbv]

it looks like this is depending on code that's still on a branch? should this PR be a draft?

[lchu-ibm]

@nairbv it does not. that extra piece of code is just for a one-time regression/validation test that performed externally and will not be repeated. we might re-run the same test next time if we modify this but this will likely be a good final version.

[lchu-ibm]

Although - that function we added as a one-time thing might be potentially useful to add in FMS too @daviswer . And it does not depend on anything else. If we do so, @daviswer just remember to check in your first version (the one I linked above, without the later "wrong fix" like adding .data)

[daviswer]

@lchu-ibm Sure, just add it in that format to this branch?

[lchu-ibm]

@daviswer yes, maybe rename it as validate_parameters()