core.dist_checkpointing.strategies.common

Common strategies.

Module Contents

Functions

save_common	Save common part of the state dict.
load_common	Load common (non-sharded) objects state dict from the checkpoint.

Data

COMMON_STATE_FNAME
logger

API

core.dist_checkpointing.strategies.common.COMMON_STATE_FNAME

‘common.pt’

core.dist_checkpointing.strategies.common.logger

‘getLogger(…)’

core.dist_checkpointing.strategies.common.save_common(

common_state_dict: megatron.core.dist_checkpointing.mapping.StateDict,

checkpoint_dir: str,

)

Save common part of the state dict.

core.dist_checkpointing.strategies.common.load_common(checkpoint_dir: str)

Load common (non-sharded) objects state dict from the checkpoint.

Parameters:

checkpoint_dir (Path) – checkpoint directory

Returns:

state dict with non-sharded objects from the checkpoint

Return type:

StateDict