Converting trained llama 2 checkpoint to hf gives "invalid key" error
State: closed | Pull request: no
Labels: name: bug

name: 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



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.