Llama3 conversion from Megatron DCP checkpoints to HF [NeMo 1.0]

[TJ-Solergibert]

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], [2], [3], [4]. 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 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 -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:

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:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• 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
• Bugfix
• 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 contains specific people who can review PRs to various areas.

Additional Information

• Related to # (issue)

cc @MaximumEntropy, @ericharper, @ekmb, @yzhang123, @VahidooX, @vladgets, @okuchaiev