feat(megatron): report peak memory for log_probs and actor_train

[EazyReal]

Summary

• print_memory samples current allocation only at fixed lifecycle points, so it does not expose a phase's high-water mark and an OOM can happen before the next report
• add report_peak_memory(phase) to miles/utils/memory_utils.py: torch.cuda.reset_peak_memory_stats() on entry scopes the peak to exactly one phase; max_memory_allocated / max_memory_reserved are logged in a finally block, so an OOMing phase still reports its peak
• wrap the {store_prefix}log_probs forward and actor_train in the Megatron actor with it, composed with the existing timer(...) scopes; each executed ref/teacher/actor pass reports under its existing prefix, while a skipped actor forward emits no report (--use-rollout-logprobs unless mismatch metrics require it, or --skip-actor-forward-only)
• does not change model execution or allocations; each scope intentionally rebases the current CUDA device's allocator peak counters. The two call sites are strictly sequential because nested scopes would discard the outer peak, and the load-time CI peak check in ci_utils.py fires before any rollout, so the scoped resets cannot affect it. The critic phases are left for a follow-up.

Measured on a 128x H100 DeepSeek-V4-Flash run at 262,144-token context: the finally-path report is what attributed a 132k-context log_probs OOM to 25 GiB of retained activations — the OOM traceback alone never shows the phase's high-water mark.

Testing

• pytest tests/fast/utils/test_memory_utils.py — scoped reset, GB formatting, and the report-on-raise path (torch.cuda mocked, CPU-only); 3/3 pass locally on CPU-only torch
• black, ruff, and isort checks pass on the touched files

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.