unified_export_megatron

Code that export quantized Megatron Core models for deployment.

Functions

export_mcore_gpt_to_hf

Export Megatron Core GPTModel to unified checkpoint and save to export_dir.

import_mcore_gpt_from_hf

Import GPTModel state_dict from supported HuggingFace pretrained model path.
export_mcore_gpt_to_hf(model, pretrained_model_name_or_path=None, dtype=torch.float16, export_dir='/tmp')

Export Megatron Core GPTModel to unified checkpoint and save to export_dir.

Parameters:

  • model (Module) – The Megatron Core GPTModel instance.

  • pretrained_model_name_or_path (str | PathLike | None) – Can be either: the model id of a pretrained model hosted inside a model repo on huggingface.co; or a directory containing model weights saved using [~PreTrainedModel.save_pretrained], e.g., ./my_model_directory/.

  • dtype (dtype) – The weights data type to export the unquantized layers.

  • export_dir (Path | str) – The target export path.

import_mcore_gpt_from_hf(model, pretrained_model_path, dtype=torch.float16)

Import GPTModel state_dict from supported HuggingFace pretrained model path.

Parameters:

  • model (Module) – The Megatron Core GPTModel instance.

  • pretrained_model_path (str) – A path to a directory containing model weights saved using [~PreTrainedModel.save_pretrained], e.g., ./my_model_directory/.

  • dtype (dtype) – The weights data type to import.