NVIDIA-NeMo / NVIDIA-NeMo/RL

Direct support for megatron bridge checkpoints

Open
#2,054 0 comments 0 reactions 1 assignee Claimed by @ananthsub View on GitHub
enhancement
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

Support Pretrained Megatron Checkpoints for Policy Worker Initialization

Motivation:

Today, the Megatron policy worker always requires a HuggingFace model ID or path as its starting point. On first launch, the worker converts the HF checkpoint into Megatron format using Megatron Bridge. This conversion is cached so subsequent launches skip it, but it still requires an initial on-the-fly conversion. If users already have a Megatron-Bridge format checkpoint (either from a training run or offline conversion), this one-time HF -> Megatron conversion is inefficient. Currently there is no way to skip this. The only entry point is a HuggingFace path, and the system derives all Megatron checkpoint paths internally.

Goals:
1. Allow users to point directly at an existing Megatron checkpoint iteration directory (e.g., `/path/to/checkpoints/iter_0001000/`) and skip the HF-to-Megatron conversion entirely
2. Retain model_name (HF model ID) in the config since it is still needed for the tokenizer and for the AutoBridge HF refit/export
3. Keep the default behavior unchanged: if no Megatron checkpoint is provided, the existing HF conversion flow runs as before
4. Cleanly separate from the training-checkpoint-resume path (weights_path), which is an orthogonal mechanism

Existing flow:
1. Derives pretrained_path (a parent directory like ~/.cache/huggingface/nemo_rl/Qwen/Qwen2.5-1.5B) from config["model_name"]. Checks if iter_0000000/ exists there.
2. handle_model_import() -- If checkpoint doesn't exist, runs import_model_from_hf_name() to convert HF -> Megatron and save to pretrained_path/iter_0000000/
3. setup_model_config() -- Reads pretrained_path/iter_0000000/run_config.yaml (hardcoded) to get the TransformerConfig for model construction.
4. _create_checkpoint_config(pretrained_path, weights_path) -- Creates Bridge's CheckpointConfig with pretrained_checkpoint=pretrained_path and load=weights_path.

Proposal:

0. the current megatron bridge `load_checkpoint` API currently expects `pretrained_checkpoint` to be a parent directory containing tracker files like `latest_train_state.pt` that tell bridge which checkpoint to load from. This means its not yet supported for users to point to a specific iteration to load from without re-arranging the directory structure. We will refactor megatron bridge to also support loading from a specific checkpoint iteration directory: https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/2624

1. Add `pretrained_megatron_checkpoint` to the MegatronConfig as an optional field
Semantics: path to a specific Megatron checkpoint iteration directory (e.g., /mnt/checkpoints/iter_1234567/). The directory must contain `run_config.yaml` for Bridge's serialized model config

2. Update [`validate_model_paths`](https://github.com/NVIDIA-NeMo/RL/blob/1a32e9e2576da6ba913e3d9b30c266040b6f9d95/nemo_rl/models/megatron/setup.py#L249-L265) to read the user-provided pretrained checkpoint, if available

3. Update [`setup_model_config`](https://github.com/NVIDIA-NeMo/RL/blob/1a32e9e2576da6ba913e3d9b30c266040b6f9d95/nemo_rl/models/megatron/setup.py#L267-L279) to not use the hardcoded iteration path, and instead pass along the user-provided pretrained checkpoint

Related: https://github.com/NVIDIA-NeMo/RL/issues/671

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.