megatron.core.dist_checkpointing.core.CheckpointingException: Object shard /ckpt/model_weights/model.decoder.layers.self_attention.core_attention._extra_state/shard_0_80.pt not found
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. @OrenLeung Hello Have you fixed this problem im stuck with the same..!

**Describe the bug**

I got MLPerf LLama2 LoRA working with 24.04-py3 pytorch image with the [same modifications](https://github.com/mlcommons/training_results_v4.0/blob/main/NVIDIA/benchmarks/llama2_70b_lora/implementations/nemo/Dockerfile) with NVIDIA's submission. The model weights are converted using the following [checkpoints conversion script](https://github.com/mlcommons/training_results_v4.0/blob/main/NVIDIA/benchmarks/llama2_70b_lora/implementations/nemo/scripts/convert_model.py#L182-L258) 

While trying to update to the latest Nemo, Megatron and Pytorch image (24.08), i ran into the following issue saying that NeMo checkpoint restore `self_attention.core_attention._extra_state` but we don't have that as for llama2 lora, core attention does not have any weights. I have checked that our self attention already have qkv projection and output projection weights.

What is the resolution towards fixing this? is there a certain thing I should put to stop restore from trying to get core attention even tho it doesnt have it?

This is my config file https://github.com/mlcommons/training_results_v4.0/blob/main/NVIDIA/benchmarks/llama2_70b_lora/implementations/nemo/conf/base/model_config.yaml

```
6: [rank6]:     instance = cls._save_restore_connector.restore_from(
6: [rank6]:   File "/workspace/ft-llm/NeMo/nemo/collections/nlp/parts/nlp_overrides.py", line 1333, in restore_from
6: [rank6]:     checkpoint = checkpoint_io.load_checkpoint(
6: [rank6]:   File "/usr/lib/python3.10/contextlib.py", line 79, in inner
6: [rank6]:     return func(*args, **kwds)
6: [rank6]:   File "/workspace/ft-llm/NeMo/nemo/utils/callbacks/dist_ckpt_io.py", line 340, in load_checkpoint
6: [rank6]:     return dist_checkpointing.load(
6: [rank6]:   File "/usr/local/lib/python3.10/dist-packages/megatron/core/dist_checkpointing/serialization.py", line 155, in load
6: [rank6]:     sharded_objects = common_strategy.load_sharded_objects(
6: [rank6]:   File "/usr/local/lib/python3.10/dist-packages/megatron/core/dist_checkpointing/strategies/common.py", line 107, in load_sharded_objects
6: [rank6]:     return dict_list_map_inplace(load_sharded_object, sharded_objects_state_dict)
6: [rank6]:   File "/usr/local/lib/python3.10/dist-packages/megatron/core/dist_checkpointing/dict_utils.py", line 180, in dict_list_map_inplace
6: [rank6]:     x[k] = dict_list_map_inplace(f, v)
7: [rank7]:   File "/usr/local/lib/python3.10/dist-packages/megatron/core/dist_checkpointing/dict_utils.py", line 180, in dict_list_map_inplace
7: [rank7]:     x[k] = dict_list_map_inplace(f, v)
7: [rank7]:   File "/usr/local/lib/python3.10/dist-packages/megatron/core/dist_checkpointing/dict_utils.py", line 184, in dict_list_map_inplace
7: [rank7]:     return f(x)
7: [rank7]:   File "/usr/local/lib/python3.10/dist-packages/megatron/core/dist_checkpointing/strategies/common.py", line 104, in load_sharded_object
7: [rank7]:     raise CheckpointingException(err_msg) from e
7: [rank7]: megatron.core.dist_checkpointing.core.CheckpointingException: Object shard /ckpt/model_weights/model.decoder.layers.self_attention.core_attention._extra_state/shard_0_80.pt not found
```

**Expected behavior**

checkpoints should work between versions

**Environment overview (please complete the following information)**

 - Environment location: [Bare-metal, Docker, Cloud(specify cloud provider - AWS, Azure, GCP, Collab)]
 - Method of NeMo install: [pip install or from source]. Please specify exact commands you used to install.
 - If method of install is [Docker], provide `docker pull` & `docker run` commands used

**Environment details**

Nvidia pytorch 24.08 Docker

With the following package updates
```
ARG NEMO_REVISION=r2.0.0rc1

RUN git clone https://github.com/NVIDIA/NeMo.git && \
    cd NeMo && \
    echo NEMO_REVISION=${NEMO_REVISION} && \
    git checkout ${NEMO_REVISION} && \
    echo NEMO_COMMIT_HASH=$(git rev-parse HEAD) && \
    pip install --no-build-isolation -e ".[nlp]"

## Megatron-core
ARG MEGATRON_REVISION=core_r0.8.0

RUN pip uninstall -y megatron-core && \
    git clone https://github.com/NVIDIA/Megatron-LM.git && \
    cd Megatron-LM && \
    echo MEGATRON_REVISION=${MEGATRON_REVISION} && \
    git checkout ${MEGATRON_REVISION} && \
    echo MEGATRON_COMMIT_HASH=$(git rev-parse HEAD) && \
    pip install . && \
    cd megatron/core/datasets && \
    make
```

```
$ ls /mnt/nvme/llama2/model/model_weights/
common.pt                                              model.decoder.layers.self_attention.linear_proj._extra_state
metadata.json                                          model.decoder.layers.self_attention.linear_proj.weight
model.decoder.final_layernorm.weight                   model.decoder.layers.self_attention.linear_qkv._extra_state
model.decoder.layers.mlp.linear_fc1._extra_state       model.decoder.layers.self_attention.linear_qkv.layer_norm_weight
model.decoder.layers.mlp.linear_fc1.layer_norm_weight  model.decoder.layers.self_attention.linear_qkv.weight
model.decoder.layers.mlp.linear_fc1.weight             model.embedding.word_embeddings.weight
model.decoder.layers.mlp.linear_fc2._extra_state       model.output_layer.weight
model.decoder.layers.mlp.linear_fc2.weight
```