bridge.training.pretrain#

Module Contents#

Functions#

pretrain

Main function to run the pretraining pipeline.

API#

bridge.training.pretrain.pretrain(
config: megatron.bridge.training.config.ConfigContainer,
forward_step_func: Callable,
) → None#

Main function to run the pretraining 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 that performs a single forward and backward step, returning the loss and any computed metrics.

.. warning::

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