bridge.training.finetune#

Module Contents#

Functions#

finetune

Main function to run the finetuning.

API#

bridge.training.finetune.finetune(
config: megatron.bridge.training.config.ConfigContainer,
forward_step_func: megatron.bridge.training.forward_step_func_types.ForwardStepCallable,
) → None#

Main function to run the finetuning.

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)

.. 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.