Converting trained llama 2 checkpoint to hf gives "invalid key" error
This issue is stale because it has been open for 30 days with no activity. Remove stale label or comment or this will be closed in 7 days. This issue was closed because it has been inactive for 7 days since being marked as stale.
**Describe the bug**

While converting a trained nemo checkpoint using `NeMo/scripts/checkpoint_converters/convert_llama_nemo_to_hf.py`, I get the following error
```
[rank0]: Traceback (most recent call last):
[rank0]:   File "/opt/megatron-lm/megatron/core/dist_checkpointing/strategies/tensorstore.py", line 123, in open_ts_array
[rank0]:     arr = ts.open(ts.Spec(spec), open=True).result()
[rank0]: ValueError: Error opening "zarr" driver: Error reading local file "../exp/nemo/exp_inhouse_7b/TP8_PP1/checkpoints/peep/model_weights/model.embedding.word_embeddings.weight/.zarray": Invalid key: "../exp/nemo/exp_inhouse_7b/TP8_PP1/checkpoints/peep/model_weights/model.embedding.word_embeddings.weight/.zarray" [source locations='tensorstore/kvstore/file/file_key_value_store.cc:659\ntensorstore/kvstore/kvstore.cc:268\ntensorstore/driver/driver.cc:114'] [tensorstore_spec='{\"context\":{\"cache_pool\":{},\"data_copy_concurrency\":{},\"file_io_concurrency\":{},\"file_io_sync\":true},\"driver\":\"zarr\",\"kvstore\":{\"driver\":\"file\",\"path\":\"../exp/nemo/exp_inhouse_7b/TP8_PP1/checkpoints/peep/model_weights/model.embedding.word_embeddings.weight/\"}}']

[rank0]: The above exception was the direct cause of the following exception:

[rank0]: Traceback (most recent call last):
[rank0]:   File "/workspace/repo/scripts/convert_llama_nemo_to_hf.py", line 222, in <module>
[rank0]:     convert(
[rank0]:   File "/workspace/repo/scripts/convert_llama_nemo_to_hf.py", line 85, in convert
[rank0]:     model = MegatronSparseFTModel.restore_from(
[rank0]:   File "/opt/NeMo/nemo/collections/nlp/models/nlp_model.py", line 478, in restore_from
[rank0]:     return super().restore_from(
[rank0]:   File "/opt/NeMo/nemo/core/classes/modelPT.py", line 468, in restore_from
[rank0]:     instance = cls._save_restore_connector.restore_from(
[rank0]:   File "/opt/NeMo/nemo/collections/nlp/parts/nlp_overrides.py", line 1339, in restore_from
[rank0]:     checkpoint = checkpoint_io.load_checkpoint(
[rank0]:   File "/usr/lib/python3.10/contextlib.py", line 79, in inner
[rank0]:     return func(*args, **kwds)
[rank0]:   File "/opt/NeMo/nemo/utils/callbacks/dist_ckpt_io.py", line 334, in load_checkpoint
[rank0]:     return dist_checkpointing.load(
[rank0]:   File "/opt/megatron-lm/megatron/core/dist_checkpointing/serialization.py", line 157, in load
[rank0]:     loaded_state_dict = sharded_strategy.load(sharded_state_dict, checkpoint_dir)
[rank0]:   File "/opt/megatron-lm/megatron/core/dist_checkpointing/strategies/tensorstore.py", line 43, in load
[rank0]:     dict_list_map_inplace(load_fn, sharded_state_dict)
[rank0]:   File "/opt/megatron-lm/megatron/core/dist_checkpointing/dict_utils.py", line 180, in dict_list_map_inplace
[rank0]:     x[k] = dict_list_map_inplace(f, v)
[rank0]:   File "/opt/megatron-lm/megatron/core/dist_checkpointing/dict_utils.py", line 180, in dict_list_map_inplace
[rank0]:     x[k] = dict_list_map_inplace(f, v)
[rank0]:   File "/opt/megatron-lm/megatron/core/dist_checkpointing/dict_utils.py", line 184, in dict_list_map_inplace
[rank0]:     return f(x)
[rank0]:   File "/opt/megatron-lm/megatron/core/dist_checkpointing/strategies/tensorstore.py", line 80, in _load_from_array
[rank0]:     x = _load_regular_chunk(sharded_tensor, checkpoint_dir)
[rank0]:   File "/opt/megatron-lm/megatron/core/dist_checkpointing/strategies/tensorstore.py", line 91, in _load_regular_chunk
[rank0]:     arr = open_ts_array(checkpoint_dir / sharded_tensor.key)
[rank0]:   File "/opt/megatron-lm/megatron/core/dist_checkpointing/strategies/tensorstore.py", line 125, in open_ts_array
[rank0]:     raise CheckpointingException(f'Array {arr_path} could not be loaded. Error: {e}') from e
[rank0]: megatron.core.dist_checkpointing.core.CheckpointingException: Array ../exp/nemo/exp_inhouse_7b/TP8_PP1/checkpoints/peep/model_weights/model.embedding.word_embeddings.weight could not be loaded. Error: Error opening "zarr" driver: Error reading local file "../exp/nemo/exp_inhouse_7b/TP8_PP1/checkpoints/peep/model_weights/model.embedding.word_embeddings.weight/.zarray": Invalid key: "../exp/nemo/exp_inhouse_7b/TP8_PP1/checkpoints/peep/model_weights/model.embedding.word_embeddings.weight/.zarray" [source locations='tensorstore/kvstore/file/file_key_value_store.cc:659\ntensorstore/kvstore/kvstore.cc:268\ntensorstore/driver/driver.cc:114'] [tensorstore_spec='{\"context\":{\"cache_pool\":{},\"data_copy_concurrency\":{},\"file_io_concurrency\":{},\"file_io_sync\":true},\"driver\":\"zarr\",\"kvstore\":{\"driver\":\"file\",\"path\":\"../exp/nemo/exp_inhouse_7b/TP8_PP1/checkpoints/peep/model_weights/model.embedding.word_embeddings.weight/\"}}']
```

**Steps/Code to reproduce bug**
1. Train a megatron llama model as usual. Here I used TP size=8, PP size=1 on 4 nodes.
2. Call `python NeMo/scripts/checkpoint_converters/convert_llama_nemo_to_hf.py --input_name_or_path $PATH_TO_SAVED_NEMO --output_path $PATH_TO_CONVERTED_HF.bin`


**Environment details**
Image used: nvcr.io/nvidia/nemo:24.05.llama3.1

