diff --git a/docs/user-guide/cli-reference.md b/docs/user-guide/cli-reference.md index 1ca81e8519..9d396e238a 100644 --- a/docs/user-guide/cli-reference.md +++ b/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` | `.` | – | Callback after each successful Megatron training step. | | `--custom-megatron-post-save-hook-path` | `.` | – | Rank-0 callback after each checkpoint save. | | `--model-name` | str | – | Set in multi-node to avoid `transformers` file-system race. | | `--spec` | ` ` | – | Plugin spec for custom architectures (e.g. `miles_plugins.models.qwen3_5 get_qwen3_5_spec`). | diff --git a/docs/user-guide/customization.md b/docs/user-guide/customization.md index 8067cb51bd..7c2128583a 100644 --- a/docs/user-guide/customization.md +++ b/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. diff --git a/docs/user-guide/usage.md b/docs/user-guide/usage.md index 3a2b0b0c34..e000002cfe 100644 --- a/docs/user-guide/usage.md +++ b/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). diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 61a49fe395..3e58f54242 100644 --- a/miles/backends/megatron_utils/model.py +++ b/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 diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 858b5c0efd..070e6eb993 100644 --- a/miles/utils/arguments.py +++ b/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): diff --git a/tests/fast/backends/megatron_utils/test_model_initialize.py b/tests/fast/backends/megatron_utils/test_model_initialize.py index d3e4ee9f11..668e75244c 100644 --- a/tests/fast/backends/megatron_utils/test_model_initialize.py +++ b/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() diff --git a/tests/fast/utils/test_arguments.py b/tests/fast/utils/test_arguments.py index 20a7bfb758..434fc27b1f 100644 --- a/tests/fast/utils/test_arguments.py +++ b/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,