Megatron -> .nemo checkpoint conversion script `megatron_lm_ckpt_to_nemo.py` fails
State: closed | Pull request: no
Labels: name: bug

name: stale

I have trained a Llama 3.1 8B model using the Nemo 24.05.llama3.1 container. It has generated intermediate Megatron checkpoints, and I want to convert them to `.nemo`. Apparently, `megatron_lm_ckpt_to_nemo.py` is supposed to accomplish that task. **I have tried both 24.05.llama3.1 and 24.07 containers for the conversion.**

Inside the script it says that I need to set the `checkpoint_name` to the folder name of the checkpoint.
![Image](https://github.com/user-attachments/assets/eae430d0-8601-46d5-9e57-8403900644be)

However, if I do this, I get  `IsADirectoryError`. So, clearly, either the documentation or the script is wrong.

To fix this, I append `/common.pt` to the `checkpoint_name`. This fixes the previous issue, and it is able to load the `.pt` file as a dictionary. However, it fails in the following way:

```
[rank0]:     checkpoint, consumed, steps, version = load_from_checkpoint(
[rank0]:   File "/workspace/scripts/megatron_lm_ckpt_to_nemo.py", line 315, in load_from_checkpoint
[rank0]:     old_checkpoint['model'],
[rank0]: KeyError: 'model'
```

Apparently, the dictionary does not have the `'model'` key, which should save the weights of the model or the path to them or something. And if they are not defined inside the .pt file, then I don't understand how it's supposed to load the weights. 

How can I make this work? Is appending `common.pt` good? Shouldn't this file have some way of accessing the weights? Is the `common.pt` wrong?

This is how I'm executing the script:

```
python -m torch.distributed.launch --nproc_per_node=1 /.../megatron_lm_ckpt_to_nemo.py \
  --checkpoint_folder=/.../checkpoints \
  --checkpoint_name=megatron_llama3.1_8b--val_loss=1.13-step=7500-consumed_samples=1920000.0/common.pt \
  --nemo_file_path=/.../abc.nemo \
  --model_type=gpt \
  --hparams_file /.../hparams.yaml \
  --tensor_model_parallel_size=1 \
  --pipeline_model_parallel_size=1 \
  --gpus_per_node 1 \
  --mcore_input
```

Structure of the checkpoint folder:

![Image](https://github.com/user-attachments/assets/3508e61e-a312-4492-bcd8-15a92de33484)


Full traceback of the first `IsADirectoryError`:

```
[rank0]: Traceback (most recent call last):
[rank0]:   File "/workspace/scripts/megatron_lm_ckpt_to_nemo.py", line 571, in <module>
[rank0]:     convert(local_rank, rank, world_size, args)
[rank0]:   File "/workspace/scripts/megatron_lm_ckpt_to_nemo.py", line 485, in convert
[rank0]:     checkpoint, consumed, steps, version = load_from_checkpoint(
[rank0]:   File "/workspace/scripts/megatron_lm_ckpt_to_nemo.py", line 308, in load_from_checkpoint
[rank0]:     old_checkpoint = pl_load(checkpoint_path, map_location=lambda storage, loc: storage)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/lightning_fabric/utilities/cloud_io.py", line 56, in _load
[rank0]:     with fs.open(path_or_url, "rb") as f:
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/fsspec/spec.py", line 1293, in open
[rank0]:     f = self._open(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/fsspec/implementations/local.py", line 184, in _open
[rank0]:     return LocalFileOpener(path, mode, fs=self, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/fsspec/implementations/local.py", line 306, in __init__
[rank0]:     self._open()
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/fsspec/implementations/local.py", line 311, in _open
[rank0]:     self.f = open(self.path, mode=self.mode)
[rank0]: IsADirectoryError: [Errno 21] Is a directory: '/.../checkpoints/megatron_llama3.1_8b--val_loss=1.13-step=7500-consumed_samples=1920000.0'
```

Full traceback of the second `KeyError`:

```[rank0]: Traceback (most recent call last):
[rank0]:   File "/workspace/scripts/megatron_lm_ckpt_to_nemo.py", line 575, in <module>
[rank0]:     convert(local_rank, rank, world_size, args)
[rank0]:   File "/workspace/scripts/megatron_lm_ckpt_to_nemo.py", line 489, in convert
[rank0]:     checkpoint, consumed, steps, version = load_from_checkpoint(
[rank0]:   File "/workspace/scripts/megatron_lm_ckpt_to_nemo.py", line 315, in load_from_checkpoint
[rank0]:     old_checkpoint['model'],
[rank0]: KeyError: 'model'
```



_* Note: The lines in the traceback may be different from the official `megatron_lm_ckpt_to_nemo.py` as I have added some debugging code._



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.