nemo_rl.models.automodel.convert#

Offline conversion of Automodel DCP checkpoints to Hugging Face format.

Module Contents#

Functions#

convert_dcp_to_hf

Convert a Torch DCP checkpoint to a Hugging Face checkpoint.

API#

nemo_rl.models.automodel.convert.convert_dcp_to_hf(
dcp_ckpt_path: str,
hf_ckpt_path: str,
model_name_or_path: str,
tokenizer_name_or_path: str,
overwrite: bool = False,
hf_overrides: Optional[dict[str, Any]] = {},
) → str#

Convert a Torch DCP checkpoint to a Hugging Face checkpoint.

This is not an optimized utility. If checkpoint is too large, consider saving DCP during training and using this utility to convert to HF format.

Parameters:

  • dcp_ckpt_path – Checkpoint dir that contains model/, for example step_N/policy/weights. The DCP shards are read from <dcp_ckpt_path>/model.

  • hf_ckpt_path – Path to save HF checkpoint.

  • model_name_or_path – Model name or path for config.

  • tokenizer_name_or_path – Tokenizer name or path.

  • overwrite – Whether to overwrite existing checkpoint. Defaults to False.

  • hf_overrides – Extra keyword arguments forwarded to AutoConfig.from_pretrained.

Returns:

Path to the saved HF checkpoint.

Raises:

  • FileExistsError – If HF checkpoint already exists and overwrite is False.

  • FileNotFoundError – If <dcp_ckpt_path>/model/.metadata does not exist.