feat(megatron): add after-train-step hook

docs/user-guide/cli-reference.md

@@ -157,6 +157,7 @@ Sections mirror the launch-script argument groups.
 | `--load` | path | – | Actor checkpoint to resume from. |
 | `--save` | path | – | Actor checkpoint write directory. |
 | `--save-interval` | int | – | Rollouts between saves. |
+| `--custom-megatron-after-train-step-hook-path` | `<module>.<fn>` | – | Callback after each successful Megatron training step. |
 | `--custom-megatron-post-save-hook-path` | `<module>.<fn>` | – | Rank-0 callback after each checkpoint save. |
 | `--model-name` | str | – | Set in multi-node to avoid `transformers` file-system race. |
 | `--spec` | `<module> <fn>` | – | Plugin spec for custom architectures (e.g. `miles_plugins.models.qwen3_5 get_qwen3_5_spec`). |

docs/user-guide/customization.md

@@ -28,6 +28,7 @@ and the default it replaces.
 | **Megatron hooks** | `--custom-megatron-init-path` | After Megatron init |
 | | `--custom-megatron-before-log-prob-hook-path` | Before logprob compute |
 | | `--custom-megatron-before-train-step-hook-path` | Before each train step |
+| | `--custom-megatron-after-train-step-hook-path` | After each successful train step |
 | | `--custom-megatron-post-save-hook-path` | After each checkpoint save |
 | **Logging** | `--custom-rollout-log-function-path` | Train-rollout logging |
 | | `--custom-eval-rollout-log-function-path` | Eval-rollout logging |
@@ -224,10 +225,13 @@ def convert_samples_to_train_data(args, samples) -> dict:
 | `--custom-megatron-init-path` | `def custom_init(args) -> None` |
 | `--custom-megatron-before-log-prob-hook-path` | `def custom_hook(args, model, store_prefix) -> None` |
 | `--custom-megatron-before-train-step-hook-path` | `def custom_hook(args, rollout_id, step_id, model, optimizer, opt_param_scheduler) -> None` |
+| `--custom-megatron-after-train-step-hook-path` | `def custom_hook(args, rollout_id, step_id, model, optimizer, opt_param_scheduler, loss_dict, num_microbatches) -> None` |
 | `--custom-megatron-post-save-hook-path` | `def hook(args, rollout_id: int, checkpoint_dir: str, hf_checkpoint_dir: str | None) -> None` |
 
 The Megatron init, log-prob, and train-step hooks give access to the live model
 and optimizer, useful for custom probes, weight clipping, or surgical interventions.
+The after-train-step hook runs only after a successful training step and can add
+metrics to `loss_dict` before the standard train-step logging path runs.
 The post-save hook runs on rank 0 after checkpoint save completion and receives
 the saved checkpoint paths instead of live model objects.
 

docs/user-guide/usage.md

@@ -102,13 +102,14 @@ command.
 
 ### Hooks
 
-Three extension points override Megatron behavior without forking:
+Four extension points override Megatron behavior without forking:
 
 | Flag | Runs |
 |---|---|
 | `--custom-megatron-init-path` | After Megatron initialization |
 | `--custom-megatron-before-log-prob-hook-path` | Before every log-probability computation |
 | `--custom-megatron-before-train-step-hook-path` | Before every training step |
+| `--custom-megatron-after-train-step-hook-path` | After every successful training step |
 
 Typical use cases: mixing in an auxiliary loss, instrumenting per-step metrics, or
 clipping weights surgically. See [Customization](/user-guide/customization#megatron-hooks).

miles/backends/megatron_utils/model.py

@@ -709,6 +709,21 @@ def train(
                 config.param_sync_func = param_sync_func
                 pre_hook_enabled = True
 
+        if train_step_outcome == TrainStepOutcome.NORMAL and args.custom_megatron_after_train_step_hook_path:
+            from miles.utils.misc import load_function
+
+            custom_after_train_step_hook = load_function(args.custom_megatron_after_train_step_hook_path)
+            custom_after_train_step_hook(
+                args,
+                rollout_id,
+                step_id,
+                model,
+                optimizer,
+                opt_param_scheduler,
+                loss_dict,
+                num_microbatches[step_id],
+            )
+
         if args.enable_mtp_training:
             from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper
 

miles/utils/arguments.py

@@ -1946,6 +1946,11 @@ def add_custom_megatron_plugins_arguments(parser):
                 type=str,
                 default=None,
             )
+            parser.add_argument(
+                "--custom-megatron-after-train-step-hook-path",
+                type=str,
+                default=None,
+            )
             return parser
 
         def add_mtp_training_arguments(parser):

tests/fast/backends/megatron_utils/test_model_initialize.py

@@ -2,10 +2,13 @@
 import types
 from argparse import Namespace
 from contextlib import ExitStack
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
 
+from miles.utils.misc import function_registry
+
 
 def _stub_module(name: str, attrs: dict[str, object] | None = None, is_package: bool = False) -> types.ModuleType:
     module = types.ModuleType(name)
@@ -50,6 +53,9 @@ def __init__(self, **kwargs):
 class _FakeModelChunk:
     role: str | None = None
 
+    def train(self) -> None:
+        pass
+
 
 @pytest.fixture(scope="module", autouse=True)
 def _mock_megatron_environment():
@@ -130,7 +136,70 @@ def _mock_megatron_environment():
                 "_setup_lora_model_via_bridge": MagicMock(),
             },
         )
+        _stub_module(
+            "miles.backends.megatron_utils.ft.indep_dp",
+            {"allreduce_grads_and_losses_across_replicas": MagicMock(return_value=(True, {}))},
+        )
+        _stub_module(
+            "miles.backends.megatron_utils.local_weight_checksum", {"dump_local_weight_checksums": MagicMock()}
+        )
+        _stub_module("miles.utils.audit_utils.witness.allocator", {"WitnessInfo": MagicMock()})
+        _stub_module("miles.utils.audit_utils.witness.module", {"witness_dump_and_clear_stale": MagicMock()})
+        _stub_module(
+            "miles.utils.dumper_utils",
+            {
+                "DumperMegatronUtil": MagicMock(),
+                "DumperPhase": types.SimpleNamespace(FWD_BWD="fwd_bwd", FWD_ONLY="fwd_only"),
+            },
+        )
+        _stub_module("miles.utils.memory_utils", {"clear_memory": MagicMock()})
+        _stub_module("miles.utils.test_utils.ft_test_actions", {"FTTestActionActorExecutor": MagicMock()})
+        _stub_module("miles.utils.tracking_utils.structured_log", {"log_structured": MagicMock()})
+        _stub_module(
+            "miles.backends.training_utils.ci_utils", {"check_grad_norm": MagicMock(), "check_kl": MagicMock()}
+        )
+        _stub_module("miles.backends.training_utils.data", {"DataIterator": MagicMock(), "get_batch": MagicMock()})
+        _stub_module(
+            "miles.backends.training_utils.log_utils",
+            {
+                "aggregate_forward_results": MagicMock(),
+                "aggregate_train_losses": MagicMock(return_value={}),
+                "log_train_step": MagicMock(),
+            },
+        )
+        _stub_module("miles.backends.training_utils.loss", {"loss_function": MagicMock()})
+        _stub_module("miles.backends.training_utils.parallel", {"get_parallel_state": MagicMock()})
+        _stub_module(
+            "miles.backends.megatron_utils.checkpoint",
+            {
+                "load_checkpoint": MagicMock(),
+                "save_checkpoint": MagicMock(),
+                "save_checkpoint_with_lora": MagicMock(),
+            },
+        )
+        _stub_module(
+            "miles.backends.megatron_utils.ci_utils",
+            {
+                "check_model_hashes": MagicMock(),
+                "check_peak_gpu_memory_after_load": MagicMock(),
+                "compute_model_hashes_by_layer": MagicMock(),
+                "save_model_hashes": MagicMock(),
+            },
+        )
+        _stub_module(
+            "miles.backends.megatron_utils.initialize",
+            {"is_first_replica_megatron_main_rank": MagicMock(return_value=True)},
+        )
+        _stub_module(
+            "miles.backends.megatron_utils.lora_utils",
+            {
+                "is_lora_enabled": MagicMock(return_value=False),
+                "is_lora_model": MagicMock(return_value=False),
+                "save_lora_checkpoint": MagicMock(),
+            },
+        )
         _stub_module("miles.backends.megatron_utils.model_provider", {"get_model_provider_func": MagicMock()})
+        _stub_module("miles.backends.megatron_utils.parallel", {"get_packed_seq_params": MagicMock()})
         yield
     finally:
         sys.modules.clear()
@@ -187,3 +256,154 @@ def test_initialize_steps_scheduler_when_checkpoint_did_not_restore_it():
 
     assert result == (model, optimizer, opt_param_scheduler, 100)
     opt_param_scheduler.step.assert_called_once_with(increment=800)
+
+
+def test_train_invokes_after_train_step_hook_before_logging():
+    from miles.backends.megatron_utils.ft.types import TrainStepOutcome
+    from miles.backends.megatron_utils.model import train
+
+    args = Namespace(
+        debug_disable_optimizer=True,
+        custom_megatron_after_train_step_hook_path="test:after_train_step_hook",
+        overlap_grad_reduce=False,
+        overlap_param_gather=False,
+        align_param_gather=False,
+        reset_optimizer_states=False,
+        manual_gc=False,
+        enable_mtp_training=False,
+        ci_test=False,
+    )
+    model = [_FakeModelChunk()]
+    data_iterator = [MagicMock()]
+    config = SimpleNamespace(no_sync_func=None, param_sync_func=None)
+    parallel_state = SimpleNamespace(indep_dp=SimpleNamespace(size=1))
+    hook_calls = []
+
+    def after_train_step_hook(
+        hook_args,
+        rollout_id,
+        step_id,
+        hook_model,
+        optimizer,
+        opt_param_scheduler,
+        loss_dict,
+        num_microbatches,
+    ):
+        loss_dict["custom_metric"] = 2.5
+        hook_calls.append(
+            (
+                hook_args,
+                rollout_id,
+                step_id,
+                hook_model,
+                optimizer,
+                opt_param_scheduler,
+                loss_dict,
+                num_microbatches,
+            )
+        )
+
+    with function_registry.temporary("test:after_train_step_hook", after_train_step_hook), ExitStack() as stack:
+        stack.enter_context(patch("miles.backends.megatron_utils.model.get_args", return_value=args))
+        stack.enter_context(
+            patch("miles.backends.megatron_utils.model.get_parallel_state", return_value=parallel_state)
+        )
+        stack.enter_context(patch("miles.backends.megatron_utils.model.get_model_config", return_value=config))
+        stack.enter_context(
+            patch("miles.backends.megatron_utils.model.should_disable_forward_pre_hook", return_value=False)
+        )
+        stack.enter_context(
+            patch(
+                "miles.backends.megatron_utils.model.train_one_step",
+                return_value=({"loss": 1.0}, 0.5, TrainStepOutcome.NORMAL),
+            )
+        )
+        stack.enter_context(
+            patch("miles.backends.megatron_utils.model.is_first_replica_megatron_main_rank", return_value=True)
+        )
+        log_train_step = stack.enter_context(
+            patch("miles.backends.megatron_utils.model.log_train_step", return_value={"loss": 1.0})
+        )
+
+        result = train(
+            rollout_id=3,
+            model=model,
+            optimizer=None,
+            opt_param_scheduler=None,
+            data_iterator=data_iterator,
+            num_microbatches=[7],
+            witness_info=None,
+            attempt=0,
+        )
+
+    assert result == TrainStepOutcome.NORMAL
+    assert hook_calls == [
+        (
+            args,
+            3,
+            0,
+            model,
+            None,
+            None,
+            {"loss": 1.0, "custom_metric": 2.5},
+            7,
+        )
+    ]
+    log_train_step.assert_called_once()
+    assert log_train_step.call_args.kwargs["loss_dict"] == {"loss": 1.0, "custom_metric": 2.5}
+
+
+def test_train_skips_after_train_step_hook_for_discarded_step():
+    from miles.backends.megatron_utils.ft.types import TrainStepOutcome
+    from miles.backends.megatron_utils.model import train
+
+    args = Namespace(
+        debug_disable_optimizer=True,
+        custom_megatron_after_train_step_hook_path="test:after_train_step_hook_discarded",
+        overlap_grad_reduce=False,
+        overlap_param_gather=False,
+        align_param_gather=False,
+        reset_optimizer_states=False,
+        manual_gc=False,
+        enable_mtp_training=False,
+        ci_test=False,
+    )
+    model = [_FakeModelChunk()]
+    data_iterator = [MagicMock()]
+    config = SimpleNamespace(no_sync_func=None, param_sync_func=None)
+    parallel_state = SimpleNamespace(indep_dp=SimpleNamespace(size=1))
+    after_train_step_hook = MagicMock()
+
+    with function_registry.temporary(
+        "test:after_train_step_hook_discarded", after_train_step_hook
+    ), ExitStack() as stack:
+        stack.enter_context(patch("miles.backends.megatron_utils.model.get_args", return_value=args))
+        stack.enter_context(
+            patch("miles.backends.megatron_utils.model.get_parallel_state", return_value=parallel_state)
+        )
+        stack.enter_context(patch("miles.backends.megatron_utils.model.get_model_config", return_value=config))
+        stack.enter_context(
+            patch("miles.backends.megatron_utils.model.should_disable_forward_pre_hook", return_value=False)
+        )
+        stack.enter_context(
+            patch(
+                "miles.backends.megatron_utils.model.train_one_step",
+                return_value=({}, 0.0, TrainStepOutcome.DISCARDED_SHOULD_RETRY),
+            )
+        )
+        log_train_step = stack.enter_context(patch("miles.backends.megatron_utils.model.log_train_step"))
+
+        result = train(
+            rollout_id=3,
+            model=model,
+            optimizer=None,
+            opt_param_scheduler=None,
+            data_iterator=data_iterator,
+            num_microbatches=[7],
+            witness_info=None,
+            attempt=0,
+        )
+
+    assert result == TrainStepOutcome.DISCARDED_SHOULD_RETRY
+    after_train_step_hook.assert_not_called()
+    log_train_step.assert_not_called()

tests/fast/utils/test_arguments.py

@@ -65,6 +65,20 @@ def test_skips_function_without_add_arguments(self, path_arg):
             get_miles_extra_args_provider()(parser)
 
 
+def test_custom_megatron_after_train_step_hook_path_is_parsed():
+    parser = argparse.ArgumentParser()
+    get_miles_extra_args_provider()(parser)
+    args, _ = parser.parse_known_args(
+        [
+            "--custom-megatron-after-train-step-hook-path",
+            "custom_hooks.after_train_step",
+        ]
+        + REQUIRED_ARGS
+    )
+
+    assert args.custom_megatron_after_train_step_hook_path == "custom_hooks.after_train_step"
+
+
 class TestMaybeApplyDumperOverrides:
     def _make_args(
         self,