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

[waitong94]

Summary

• Add a configurable Megatron after-train-step hook path.
• Invoke the hook after each successful Megatron train step with train-step context, mutable loss metrics, and the per-step microbatch count.
• Document the hook alongside the existing Megatron customization hooks.

Why

This gives training integrations a generic extension point for post-step diagnostics or metric enrichment without changing the core train loop for each use case.

Validation

• source .venv/bin/activate && python -m pytest --confcutdir=tests/fast tests/fast/backends/megatron_utils/test_model_initialize.py
• source .venv/bin/activate && pre-commit run --files docs/user-guide/usage.md docs/user-guide/customization.md docs/user-guide/cli-reference.md miles/backends/megatron_utils/model.py miles/utils/arguments.py tests/fast/backends/megatron_utils/test_model_initialize.py tests/fast/utils/test_arguments.py

[gemini-code-assist]

Code Review

This pull request introduces a new command-line argument --custom-megatron-after-train-step-hook-path and integrates a corresponding post-train-step hook into the Megatron training loop. It also adds unit tests to verify that the hook is correctly parsed and executed before logging. The review feedback highlights a critical issue where the hook is executed unconditionally, even when a training step is discarded. This could pass an empty loss dictionary to the hook and cause crashes, so it is recommended to only run the hook when the step outcome is normal.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[waitong94]

/gemini rereview

[gemini-code-assist]

I have completed the re-review of your pull request. The implementation of the after-train-step hook correctly follows the existing patterns for Megatron customization hooks in the codebase. The logic in miles/backends/megatron_utils/model.py is appropriately guarded, and the provided tests verify both the successful invocation of the hook and the expected behavior when skipping it. The documentation updates in the user guide are also clear and consistent with the new functionality. No issues were identified.