nemo_rl.models.megatron.community_import#

Module Contents#

Functions#

import_model_from_hf_name

Import a Hugging Face model into Megatron checkpoint format and save the Megatron checkpoint to the output path.

export_model_from_megatron

API#

nemo_rl.models.megatron.community_import.import_model_from_hf_name(
hf_model_name: str,
output_path: str,
megatron_config: Optional[nemo_rl.models.policy.MegatronConfig] = None,
**config_overrides: Any,
)#

Import a Hugging Face model into Megatron checkpoint format and save the Megatron checkpoint to the output path.

Parameters:

  • hf_model_name – Hugging Face model ID or local path (e.g., ‘meta-llama/Llama-3.1-8B-Instruct’).

  • output_path – Directory to write the Megatron checkpoint (e.g., /tmp/megatron_ckpt).

  • megatron_config – Optional megatron config with paralellism settings for distributed megatron model import.

nemo_rl.models.megatron.community_import.export_model_from_megatron(
hf_model_name: str,
input_path: str,
output_path: str,
hf_tokenizer_path: str,
overwrite: bool = False,
hf_overrides: Optional[dict[str, Any]] = {},
)#