bridge.training.pretrain#

Module Contents#

Functions#

pretrain

Main function to run the training pipeline.

_pretrain

Internal function containing the actual pretrain logic.

_abort_async_checkpoint_worker

Abort async checkpoint state before distributed teardown.

_safe_distributed_rank

Return a rank identifier without obscuring an active training failure.

_cleanup_after_pretrain_failure

Clean up framework-owned state after ordinary pretrain execution fails.

_maybe_destroy_process_group

Destroy or abort process groups created by this training session.

Data#

logger

API#

bridge.training.pretrain.logger#

‘getLogger(…)’

bridge.training.pretrain.pretrain(
config: megatron.bridge.training.config.ConfigContainer,
forward_step_func: megatron.bridge.training.forward_step_func_types.ForwardStepCallable,
callbacks: list[megatron.bridge.training.callbacks.Callback] | megatron.bridge.training.callbacks.CallbackManager | None = None,
) → None#

Main function to run the training pipeline.

Sets up the environment, model, optimizer, scheduler, and data iterators. Performs training, validation, and optionally testing based on the provided configuration.

Parameters:

  • config – The main configuration container holding all necessary parameters.

  • forward_step_func –

    A callable (function or functor) that performs a single forward and backward step, returning the loss and any computed metrics. Supports the following signatures:

    • 2 args: (data_iterator, model)

    • 3 args: (data_iterator, model, return_schedule_plan=False) OR (state: GlobalState, data_iterator, model)

    • 4 args: (state: GlobalState, data_iterator, model, return_schedule_plan=False)

  • callbacks –

    Optional callbacks for custom logic injection. Can be:

    • list[Callback]: List of Callback subclass instances

    • CallbackManager: Pre-configured manager with registered callbacks

    • None: No callbacks (default)

.. note::

Use the signature with GlobalState type hint for full access to configuration, timers, and training state. State injection is automatic based on type hints or parameter names. Functors (classes with call) are fully supported.

.. warning::

This is an experimental API and is subject to change in backwards incompatible ways without notice.

bridge.training.pretrain._pretrain(
state: megatron.bridge.training.state.GlobalState,
forward_step_func: megatron.bridge.training.forward_step_func_types.ForwardStepCallable,
callback_manager: megatron.bridge.training.callbacks.CallbackManager | None = None,
store: torch.distributed.Store | None = None,
inprocess_call_wrapper: nvidia_resiliency_ext.inprocess.CallWrapper | None = None,
) → None#

Internal function containing the actual pretrain logic.

Parameters:

  • state – Global training state containing the validated configuration and runtime objects

  • forward_step_func – Function or functor that performs a single forward/backward step

  • callback_manager – Optional CallbackManager for custom callback execution

  • store – Optional distributed Store used by in-process restart for coordination

  • inprocess_call_wrapper – Optional wrapper injected by nvrx to expose restart iteration

bridge.training.pretrain._abort_async_checkpoint_worker(
state: megatron.bridge.training.state.GlobalState,
) → None#

Abort async checkpoint state before distributed teardown.

bridge.training.pretrain._safe_distributed_rank() → str#

Return a rank identifier without obscuring an active training failure.

bridge.training.pretrain._cleanup_after_pretrain_failure(
state: megatron.bridge.training.state.GlobalState,
should_destroy_process_group: bool,
) → None#

Clean up framework-owned state after ordinary pretrain execution fails.

bridge.training.pretrain._maybe_destroy_process_group(
should_destroy: bool,
*,
synchronize: bool = True,
abort: bool = False,
) → None#

Destroy or abort process groups created by this training session.

Parameters:

  • should_destroy – Whether the process group should be destroyed

  • synchronize – Whether to synchronize ranks before destruction

  • abort – Whether to abort all process groups instead of waiting for their outstanding work to finish