Add callbacks to finetune function

src/megatron/bridge/training/finetune.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from megatron.bridge.training.callbacks import Callback, CallbackManager
 from megatron.bridge.training.config import ConfigContainer
 from megatron.bridge.training.forward_step_func_types import ForwardStepCallable
 from megatron.bridge.training.pretrain import pretrain
@@ -22,6 +23,7 @@
 def finetune(
     config: ConfigContainer,
     forward_step_func: ForwardStepCallable,
+    callbacks: list[Callback] | CallbackManager | None = None,
 ) -> None:
     """Main function to run the finetuning.
 
@@ -34,6 +36,7 @@ def finetune(
                           - 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 list of Callback instances, a CallbackManager, or None.
 
     Note:
         Use the signature with GlobalState type hint for full access to configuration, timers, and training state.
@@ -47,4 +50,4 @@ def finetune(
     assert config.checkpoint.pretrained_checkpoint is not None or config.checkpoint.load is not None, (
         "Finetuning requires a loading from a pretrained checkpoint or resuming from a checkpoint"
     )
-    return pretrain(config, forward_step_func)
+    return pretrain(config, forward_step_func, callbacks=callbacks)

tests/unit_tests/training/test_finetune.py

@@ -80,7 +80,7 @@ def test_finetune_succeeds_with_pretrained_checkpoint(self):
                 # This should not raise an AssertionError
                 finetune(container, mock_forward_step_func)
                 # Verify that pretrain was called with the correct arguments
-                mock_pretrain.assert_called_once_with(container, mock_forward_step_func)
+                mock_pretrain.assert_called_once_with(container, mock_forward_step_func, callbacks=None)
             finally:
                 restore_get_world_size_safe(og_ws, cfg_mod)
 
@@ -108,6 +108,6 @@ def test_finetune_succeeds_with_load_checkpoint(self):
                 # This should not raise an AssertionError
                 finetune(container, mock_forward_step_func)
                 # Verify that pretrain was called with the correct arguments
-                mock_pretrain.assert_called_once_with(container, mock_forward_step_func)
+                mock_pretrain.assert_called_once_with(container, mock_forward_step_func, callbacks=None)
             finally:
                 restore_get_world_size_safe(og_ws, cfg_mod)