Is this page helpful?

core.post_training.modelopt.checkpointing#

Dist checkpointing modules needed for ModelOpt.

Module Contents#

Functions#

remove_per_module_state

Remove metadata from the modelopt_state.

save_modelopt_state

Save modelopt_state as a part of the per rank state_dict.

save_sharded_modelopt_state

Save modelopt_state in the sharded state_dict format.

_load_extra_state_from_sharded_checkpoint

Load extra state from sharded checkpoint.

restore_sharded_modelopt_state

Restore modelopt_state from the sharded state_dict format.

Data#

logger

API#

core.post_training.modelopt.checkpointing.logger#

‘getLogger(…)’

core.post_training.modelopt.checkpointing.remove_per_module_state(modelopt_state: dict[str, Any]) → None#

Remove metadata from the modelopt_state.

The metadata of the modelopt_state contains keys which may change with different pipeline and expert parallelism. As a result, the metadata must be stored as several ShardedObject with global and local layer offset mapping.

Parameters:

modelopt_state – the state_dict that contains all algorithms that have been applied to the given model.

core.post_training.modelopt.checkpointing.save_modelopt_state(
model: list[torch.nn.Module],
state_dict: dict[str, Any],
) → None#

Save modelopt_state as a part of the per rank state_dict.

NOTE: Only used for Megatron-LM.

Parameters:

  • model – the modelopt optimized model

  • state_dict – the current modelopt optimized model state_dict to store

core.post_training.modelopt.checkpointing.save_sharded_modelopt_state(
model: list[torch.nn.Module],
checkpoint_name: str | pathlib.Path,
sharded_strategy: tuple[str, int] | None = None,
prefix: str = '',
) → None#

Save modelopt_state in the sharded state_dict format.

Parameters:

  • model – the model to restore the modelopt optimization

  • checkpoint_name – the checkpoint folder path

  • sharded_strategy – configures sharded tensors saving behavior and backend

  • prefix – the prefix to add to the modelopt_state keys (“model.” for NeMo)

core.post_training.modelopt.checkpointing._load_extra_state_from_sharded_checkpoint(
model: torch.nn.Module,
checkpoint_name: str | pathlib.Path,
prefix: str,
) → None#

Load extra state from sharded checkpoint.

Note: since extra_state is a subset of full the sharded_state_dict, we use strict=StrictHandling.LOG_UNEXPECTED instead of LOG_ALL.

Parameters:

  • model – the model to load extra state into

  • checkpoint_name – the checkpoint folder path

  • prefix – the prefix to add to the modelopt_state keys

core.post_training.modelopt.checkpointing.restore_sharded_modelopt_state(
model: list[torch.nn.Module],
checkpoint_name: str | pathlib.Path,
prefix: str = '',
) → None#

Restore modelopt_state from the sharded state_dict format.

Parameters:

  • model – the model to restore the modelopt optimization

  • checkpoint_name – the checkpoint folder path

  • prefix – the prefix to add to the modelopt_state keys (“model.” for NeMo)