bridge.training.post_training.checkpointing#

Input/output checkpointing for ModelOpt.

Module Contents#

Functions#

restore_sharded_modelopt_state

Restore ModelOpt state saved in legacy or current MCore checkpoint layouts.

_validate_modelopt_checkpointing

Reject checkpoint sources whose ModelOpt state cannot be selected consistently.

_get_modelopt_checkpoint_path

Get the checkpoint iteration path to use for ModelOpt operations.

has_modelopt_state

Check if ModelOpt state exists inside the checkpoint path.

load_modelopt_state

Load modelopt_state from a checkpoint.

Data#

logger

API#

bridge.training.post_training.checkpointing.logger#

‘getLogger(…)’

bridge.training.post_training.checkpointing.restore_sharded_modelopt_state(
model: list[megatron.core.transformer.module.MegatronModule],
checkpoint_path: str,
) → None#

Restore ModelOpt state saved in legacy or current MCore checkpoint layouts.

Parameters:

  • model – Unwrapped model chunks to restore.

  • checkpoint_path – Iteration checkpoint directory containing modelopt_state.

Raises:

ValueError – If virtual pipeline parallelism produced multiple model chunks.

bridge.training.post_training.checkpointing._validate_modelopt_checkpointing(
non_persistent_ckpt_type: str | None,
) → None#

Reject checkpoint sources whose ModelOpt state cannot be selected consistently.

bridge.training.post_training.checkpointing._get_modelopt_checkpoint_path(
checkpoint_path: str,
ckpt_step: int | None = None,
) → str#

Get the checkpoint iteration path to use for ModelOpt operations.

bridge.training.post_training.checkpointing.has_modelopt_state(
checkpoint_path: str,
ckpt_step: int | None = None,
) → bool#

Check if ModelOpt state exists inside the checkpoint path.

Checks for modelopt_state in iteration directories (iter_*) or root directory. NOTE: Ignores distillation state which is deprecated and unused.

Parameters:

  • checkpoint_path – Path to the checkpoint directory.

  • ckpt_step – Specific checkpoint iteration selected for native loading.

Returns:

True if modelopt_state folder exists and contains nontrivial state, else False.

bridge.training.post_training.checkpointing.load_modelopt_state(
model: list[megatron.core.transformer.module.MegatronModule],
checkpoint_path: str,
ckpt_step: int | None = None,
) → None#

Load modelopt_state from a checkpoint.

Parameters:

  • model – The model to load the modelopt_state into

  • checkpoint_path – Path to the checkpoint directory.

  • ckpt_step – Specific checkpoint iteration selected for native loading.