convert_hf_llama_to_nemo.py -> megatron_change_num_partitions.py on Llama 2 70b results in a crash.
State: closed | Pull request: no
Labels: name: bug

name: stale

**Describe the bug**

I'm trying to RLHF Llama 2 70b w/ trlX NeMo. A required step of this is to convert llama 2 70b to the nemo format and perform tensor parallel sharding.

Currently, conversion works fine. I can get a copy of llama 70b that loads within nemo, albeit without tensor parallel sharding. When I run `megatron_change_num_partitions.py` on my checkpoint though, and attempt to change the number of TP shards from 1 to 8, I get the following error:

```[NeMo I 2023-10-02 15:58:58 megatron_change_num_partitions:352] Writing pp rank 0 tp rank 7 to file /home-dir/nemo_checkpoints/SB2_8tp
Invalid access pattern for ShardedTensor(key='model.embedding.word_embeddings.weight'): tensor([[1],
        [0],
        [0],
        [0],
        [0],
        [0],
        [0],
        [0]], dtype=torch.int32)
Traceback (most recent call last):
  File "/home-dir/nemo-trlx/NeMo/examples/nlp/language_modeling/megatron_change_num_partitions.py", line 1502, in <module>
    main()
  File "/home-dir/nemo-trlx/NeMo/examples/nlp/language_modeling/megatron_change_num_partitions.py", line 1455, in main
    global_offset = split_partition(
  File "/home-dir/nemo-trlx/NeMo/examples/nlp/language_modeling/megatron_change_num_partitions.py", line 746, in split_partition
    write_tp_pp_split(model, splits, app_state, tp_size, pp_rank, write_path)
  File "/home-dir/nemo-trlx/NeMo/examples/nlp/language_modeling/megatron_change_num_partitions.py", line 353, in write_tp_pp_split
    model.save_to(write_path)
  File "/opt/NeMo/nemo/core/classes/modelPT.py", line 392, in save_to
    self._save_restore_connector.save_to(self, str(save_path))  # downstream tasks expect str, not Path
  File "/opt/NeMo/nemo/collections/nlp/parts/nlp_overrides.py", line 480, in save_to
    dist_checkpointing.save(sharded_state_dict=sharded_state_dict, checkpoint_dir=dist_ckpt_dir)
  File "/opt/Megatron-LM/megatron/core/dist_checkpointing/serialization.py", line 160, in save
    validate_sharding_integrity(sharded_tensors)
  File "/opt/Megatron-LM/megatron/core/dist_checkpointing/serialization.py", line 217, in validate_sharding_integrity
    _validate_sharding_for_key(shardings)
  File "/opt/Megatron-LM/megatron/core/dist_checkpointing/serialization.py", line 245, in _validate_sharding_for_key
    raise CheckpointingException(f'Invalid access pattern for {rank_sharding[0][1]}')
```

**Steps/Code to reproduce bug**

1. Download llama 2 70b
2. Create a copy of the llama 2 config that is provided [here](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/language_modeling/conf/megatron_llama_config.yaml) and tweak it for 70b as opposed to 7b. Details on how to do so are provided within the configuration file.
3. Run `convert_hf_llama_to_nemo.py` within NeMo's docker container (in my case, on a node of H100s).
4. Point `megatron_change_num_partitions.py` to the tar created by the conversion process. (Note that as it stands right now, you actually cannot run this script on extracted tars. It does not work. See [here](https://github.com/NVIDIA/NeMo/blob/f477e051ec68aaa909ca891c1605383f87e11fbb/examples/nlp/language_modeling/megatron_change_num_partitions.py#L986). It will attempt to extract the tar regardless of parameters.)
5. Get trace/error provided above during the saving sequence after sharding has been completed (?)


**Expected behavior**

I expect a fully sharded model to be produced at the end, however this obviously does not occur.

**Environment overview (please complete the following information)**

 - Environment location: Lambda Labs, NeMo GAI Docker Image (nvcr.io/ea-bignlp/ga-participants/nemofw-training:23.08.02) 
 - Method of NeMo install: [pip install or from source]. Docker