Distributed checkpointing with mcore GPT
State: closed | Pull request: yes
Labels: name: core

name: NLP

name: CI

This PR needs mcore dist ckpt for GPT PR to be pushed before merging.

# What does this PR do ?

Adds distributed checkpointing when using mcore gpt.

Distributed checkpointing enables training runs to restart automatically with different model parallel configs. 
The checkpoint is saved to disk according to the `sharded_state_dict`:

Below is a sample of what the checkpoint looks like on disk.

``` bash
common.pt                                                     model.decoder.layers.self_attention.linear_qkv.weight                           optimizer.state.exp_avg.model.embedding.word_embeddings.weight                     optimizer.state.fp32_from_fp16.model.decoder.final_layernorm.bias
metadata.json                                                 model.embedding.position_embeddings.weight                                      optimizer.state.exp_avg.model.output_layer.weight                                  optimizer.state.fp32_from_fp16.model.decoder.final_layernorm.weight
model.decoder.final_layernorm.bias                            model.embedding.word_embeddings.weight                                          optimizer.state.exp_avg_sq.model.decoder.final_layernorm.bias                      optimizer.state.fp32_from_fp16.model.decoder.layers.input_layernorm.bias
model.decoder.final_layernorm.weight                          model.output_layer.weight                                                       optimizer.state.exp_avg_sq.model.decoder.final_layernorm.weight                    optimizer.state.fp32_from_fp16.model.decoder.layers.input_layernorm.weight
model.decoder.layers.input_layernorm.bias                     optimizer.state.exp_avg.model.decoder.final_layernorm.bias                      optimizer.state.exp_avg_sq.model.decoder.layers.input_layernorm.bias               optimizer.state.fp32_from_fp16.model.decoder.layers.mlp.linear_fc1.bias
model.decoder.layers.input_layernorm.weight                   optimizer.state.exp_avg.model.decoder.final_layernorm.weight                    optimizer.state.exp_avg_sq.model.decoder.layers.input_layernorm.weight             optimizer.state.fp32_from_fp16.model.decoder.layers.mlp.linear_fc1.weight
model.decoder.layers.mlp.linear_fc1.bias                      optimizer.state.exp_avg.model.decoder.layers.input_layernorm.bias               optimizer.state.exp_avg_sq.model.decoder.layers.mlp.linear_fc1.bias                optimizer.state.fp32_from_fp16.model.decoder.layers.mlp.linear_fc2.bias
model.decoder.layers.mlp.linear_fc1._extra_state              optimizer.state.exp_avg.model.decoder.layers.input_layernorm.weight             optimizer.state.exp_avg_sq.model.decoder.layers.mlp.linear_fc1.weight              optimizer.state.fp32_from_fp16.model.decoder.layers.mlp.linear_fc2.weight
model.decoder.layers.mlp.linear_fc1.weight                    optimizer.state.exp_avg.model.decoder.layers.mlp.linear_fc1.bias                optimizer.state.exp_avg_sq.model.decoder.layers.mlp.linear_fc2.bias                optimizer.state.fp32_from_fp16.model.decoder.layers.post_self_attn_layernorm.bias
model.decoder.layers.mlp.linear_fc2.bias                      optimizer.state.exp_avg.model.decoder.layers.mlp.linear_fc1.weight              optimizer.state.exp_avg_sq.model.decoder.layers.mlp.linear_fc2.weight              optimizer.state.fp32_from_fp16.model.decoder.layers.post_self_attn_layernorm.weight
model.decoder.layers.mlp.linear_fc2._extra_state              optimizer.state.exp_avg.model.decoder.layers.mlp.linear_fc2.bias                optimizer.state.exp_avg_sq.model.decoder.layers.post_self_attn_layernorm.bias      optimizer.state.fp32_from_fp16.model.decoder.layers.self_attention.linear_proj.bias
model.decoder.layers.mlp.linear_fc2.weight                    optimizer.state.exp_avg.model.decoder.layers.mlp.linear_fc2.weight              optimizer.state.exp_avg_sq.model.decoder.layers.post_self_attn_layernorm.weight    optimizer.state.fp32_from_fp16.model.decoder.layers.self_attention.linear_proj.weight
model.decoder.layers.post_self_attn_layernorm.bias            optimizer.state.exp_avg.model.decoder.layers.post_self_attn_layernorm.bias      optimizer.state.exp_avg_sq.model.decoder.layers.self_attention.linear_proj.bias    optimizer.state.fp32_from_fp16.model.decoder.layers.self_attention.linear_qkv.bias
model.decoder.layers.post_self_attn_layernorm.weight          optimizer.state.exp_avg.model.decoder.layers.post_self_attn_layernorm.weight    optimizer.state.exp_avg_sq.model.decoder.layers.self_attention.linear_proj.weight  optimizer.state.fp32_from_fp16.model.decoder.layers.self_attention.linear_qkv.weight
model.decoder.layers.self_attention.linear_proj.bias          optimizer.state.exp_avg.model.decoder.layers.self_attention.linear_proj.bias    optimizer.state.exp_avg_sq.model.decoder.layers.self_attention.linear_qkv.bias     optimizer.state.fp32_from_fp16.model.embedding.position_embeddings.weight
model.decoder.layers.self_attention.linear_proj._extra_state  optimizer.state.exp_avg.model.decoder.layers.self_attention.linear_proj.weight  optimizer.state.exp_avg_sq.model.decoder.layers.self_attention.linear_qkv.weight   optimizer.state.fp32_from_fp16.model.embedding.word_embeddings.weight
model.decoder.layers.self_attention.linear_proj.weight        optimizer.state.exp_avg.model.decoder.layers.self_attention.linear_qkv.bias     optimizer.state.exp_avg_sq.model.embedding.position_embeddings.weight              optimizer.state.fp32_from_fp16.model.output_layer.weight
model.decoder.layers.self_attention.linear_qkv.bias           optimizer.state.exp_avg.model.decoder.layers.self_attention.linear_qkv.weight   optimizer.state.exp_avg_sq.model.embedding.word_embeddings.weight
model.decoder.layers.self_attention.linear_qkv._extra_state   optimizer.state.exp_avg.model.embedding.position_embeddings.weight              optimizer.state.exp_avg_sq.model.output_layer.weight
```

Then inside a module directory we have the sharded tensor: 

``` bash
ls model.decoder.layers.mlp.linear_fc1.weight/
0.0.0  1.0.0  10.0.0  11.0.0  12.0.0  13.0.0  14.0.0  15.0.0  2.0.0  3.0.0  4.0.0  5.0.0  6.0.0  7.0.0  8.0.0  9.0.0
```
To implement distributed checkpointing for a model, the `sharded_state_dict` has to be defined.
This is done in megatron core so that in NeMo, if the module is from mcore, we only have to  call `module.sharded_state_dict()`.

**Collection**: NLP

# Usage

Usage is automatic when using mcore:

``` bash
model.mcore_gpt=True
```


# Before your PR is "Ready for review"
**Pre checks**:
- [ ] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [ ] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Related to # (issue)