Llama3 conversion from Megatron DCP checkpoints to HF [NeMo 1.0]
State: closed | Pull request: yes
Labels: name: stale

# What does this PR do ?

Improves docs and fixes tool to convert from a NeMo/Megatron DCP checkpoint to a HF checkpoint

**Collection**: NLP

# Changelog 
- Added cmd launch to convert from NeMo/Megatron DCP checkpoint to .nemo 
- Fixed `scripts/checkpoint_converters/convert_llama_nemo_to_hf.py`

# Usage
Hi!

There has been some discussions regarding NeMo/Megatron checkpoint conversions in both projects [[1]](https://github.com/NVIDIA/NeMo/issues/10831), [[2]](https://github.com/NVIDIA/NeMo/issues/10480), [[3]](https://github.com/NVIDIA/Megatron-LM/issues/1206), [[4]](https://github.com/NVIDIA/NeMo/issues/10884). I would like to share some insights regarding this issue, as I've been able to successfully convert checkpoints to HF.

First, my dependencies:
- NeMo commit `fc7905d`, I know it's a month old but it's the version we are working with
- Megatron commit `bf74129`
- TransformerEngine commit `7d576ed`. It's a bit old, but it's the one used in the NeMo/Dockerfile.ci image
- PyTorch NGC Container `[nvcr.io/nvidia/pytorch:24.05-py3](http://nvcr.io/nvidia/pytorch:24.05-py3)`

Some comments after converting a NeMo/Megatron Llama3.1-8B checkpoint stored with the TORCH distributed backend:
- The [official documentation](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/checkpoints/convert_mlm.html#converting-checkpoints) points out the examples/nlp/language_modeling/megatron_lm_ckpt_to_nemo.py , but it doesn't works. Instead I could use a pretty similar script stored under the same folder examples/nlp/language_modeling/megatron_ckpt_to_nemo.py. If I'm not wrong, the docs state *You can convert your GPT-style model checkpoints trained with Megatron-LM into the NeMo Framework* but now when using NeMo with `model.dist_ckpt_format: torch_dist` we are indeed storing Megatron-LM checkpoints, right? So this step will be essential to convert to `.nemo` before converting to HF.

  - With this later script I could successfully convert a NeMo/Megatron TORCH distributed checkpoint to .nemo this way:
 ```python
python -m torch.distributed.launch --nproc_per_node=4 examples/nlp/language_modeling/megatron_ckpt_to_nemo.py --checkpoint_folder megatron_ckpts/checkpoints/ --checkpoint_name llama_3_1_70b-step=50-consumed_samples=750.0-last --nemo_file_path nemo_converted/converted.nemo --hparams_file megatron_ckpts/hparams.yaml --tensor_model_parallel_size 4 --model_type gpt --gpus_per_node 4 --pipeline_model_parallel_size 1
```
- Then, to convert from .nemo to HF I used this script: ‎scripts/checkpoint_converters/convert_llama_nemo_to_hf.py
  - Which I had to edit all the `model.state_dict()` calls with `model.model[0].state_dict()`. Also the key --> `[f'model. ...']`with `[f'...']`
  - Hardcode `model_config.sequence_parallel = 0` as it raised a error if the model was trained with TP (We are also overriding TP = 1)
  - To run the conversion:
```python
python3 -u scripts/checkpoint_converters/convert_llama_nemo_to_hf.py --input_name_or_path=nemo_converted/converted.nemo --output_path=hf_converted/pytorch_model.bin --cpu-only
```
  - I also experienced some issues loading the tokenizer, but I've seen some changes in the recent commits. Also I think it would be a great idea to convert to a HF .safetensors checkpoint automatically without the necessity of specifying a path to a hf model as we can build a model from the `hparams.yaml` file (I can push this changes too)
  - This changes might be also necessary for the other model conversion scripts, but I haven't tried

IMO I think it would be nice to clarify how to convert checkpoints to HF as it's the most popular platform to share checkpoints with the community. 

# GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

# Before your PR is "Ready for review"
**Pre checks**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [x] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [ ] New Feature
- [x] Bugfix
- [x] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Related to # (issue)

cc @MaximumEntropy, @ericharper, @ekmb, @yzhang123, @VahidooX, @vladgets, @okuchaiev


This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

This PR was closed because it has been inactive for 7 days since being marked as stale.