diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index 1d20abc96..e918c4298 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/slime/backends/megatron_utils/actor.py @@ -20,7 +20,7 @@ from slime.utils import accelerator from slime.utils.data import process_rollout_data from slime.utils.distributed_utils import get_gloo_group -from slime.utils.memory_utils import clear_memory, print_memory +from slime.utils.memory_utils import clear_memory, print_memory, report_peak_memory from slime.utils.misc import Box from slime.utils.reloadable_process_group import ( destroy_process_groups, @@ -360,7 +360,7 @@ def compute_log_prob( num_microbatches: list[int], store_prefix: str = "", ) -> dict[str, list[torch.Tensor]]: - with timer(f"{store_prefix}log_probs"): + with timer(f"{store_prefix}log_probs"), report_peak_memory(f"{store_prefix}log_probs"): return forward_only( get_log_probs_and_entropy, self.args, @@ -521,7 +521,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data capture_log_probs = self.args.save_debug_train_data is not None and "log_probs" not in rollout_data if capture_log_probs: enable_log_prob_capture() - with timer("actor_train"): + with timer("actor_train"), report_peak_memory("actor_train"): train( rollout_id, self.model, diff --git a/slime/utils/accelerator/torch_accelerator.py b/slime/utils/accelerator/torch_accelerator.py index d961ebe3d..cd3b964cb 100644 --- a/slime/utils/accelerator/torch_accelerator.py +++ b/slime/utils/accelerator/torch_accelerator.py @@ -154,6 +154,11 @@ def supports(self, capability: str) -> bool: module = self._module() if capability == "device_memory": return all(hasattr(module, name) for name in ("empty_cache", "mem_get_info", "memory_allocated")) + if capability == "peak_memory": + return all( + hasattr(module, name) + for name in ("reset_peak_memory_stats", "max_memory_allocated", "max_memory_reserved") + ) if capability == "events": return hasattr(module, "Event") if capability == "rng": diff --git a/slime/utils/memory_utils.py b/slime/utils/memory_utils.py index 4bad36b51..b72aa9e8e 100644 --- a/slime/utils/memory_utils.py +++ b/slime/utils/memory_utils.py @@ -1,5 +1,6 @@ import gc import logging +from contextlib import contextmanager import psutil import torch @@ -40,6 +41,29 @@ def _byte_to_gb(n: int): return round(n / (1024**3), 2) +@contextmanager +def report_peak_memory(phase: str): + """Log the phase's peak allocated/reserved memory when supported. + + Scopes must not nest: the reset on entry discards an outer scope's peak. + """ + backend = accelerator.get_accelerator() + if not backend.supports("peak_memory"): + yield + return + + device_module = backend.accelerator_module() + device_module.reset_peak_memory_stats() + try: + yield + finally: + logger.info( + f"[Rank {dist.get_rank()}] Peak-Memory {phase}: " + f"max_allocated_GB={_byte_to_gb(device_module.max_memory_allocated())}, " + f"max_reserved_GB={_byte_to_gb(device_module.max_memory_reserved())}" + ) + + def print_memory(msg, clear_before_print: bool = False): if clear_before_print: clear_memory() diff --git a/tests/test_accelerator.py b/tests/test_accelerator.py index 301737550..163ef5d0a 100644 --- a/tests/test_accelerator.py +++ b/tests/test_accelerator.py @@ -1,8 +1,10 @@ +import logging from types import SimpleNamespace +from unittest.mock import Mock import pytest -from slime.utils import accelerator +from slime.utils import accelerator, memory_utils NUM_GPUS = 0 @@ -69,6 +71,56 @@ def reset_accelerator_selection(monkeypatch): accelerator._MUSA_BOOTSTRAP_CHECKED = bootstrap_checked +@pytest.fixture +def peak_memory_module(monkeypatch): + module = SimpleNamespace( + reset_peak_memory_stats=Mock(), + max_memory_allocated=Mock(return_value=25 * 1024**3), + max_memory_reserved=Mock(return_value=30 * 1024**3), + ) + backend = SimpleNamespace( + supports=lambda capability: capability == "peak_memory", + accelerator_module=lambda: module, + ) + monkeypatch.setattr(accelerator, "get_accelerator", lambda: backend) + monkeypatch.setattr(memory_utils.dist, "get_rank", lambda: 3) + return module + + +@pytest.mark.unit +def test_peak_memory_report_uses_selected_accelerator(peak_memory_module, caplog): + with caplog.at_level(logging.INFO, logger="slime.utils.memory_utils"): + with memory_utils.report_peak_memory("actor_train"): + peak_memory_module.reset_peak_memory_stats.assert_called_once() + + assert caplog.messages[-1] == ("[Rank 3] Peak-Memory actor_train: max_allocated_GB=25.0, max_reserved_GB=30.0") + + +@pytest.mark.unit +def test_peak_memory_report_runs_when_body_raises(peak_memory_module, caplog): + error = RuntimeError("device out of memory") + + with caplog.at_level(logging.INFO, logger="slime.utils.memory_utils"): + with pytest.raises(RuntimeError) as exc_info: + with memory_utils.report_peak_memory("log_probs"): + raise error + + assert exc_info.value is error + assert caplog.messages[-1].startswith("[Rank 3] Peak-Memory log_probs:") + + +@pytest.mark.unit +def test_peak_memory_report_skips_unsupported_accelerator(monkeypatch, caplog): + backend = SimpleNamespace(supports=lambda _capability: False) + monkeypatch.setattr(accelerator, "get_accelerator", lambda: backend) + + with caplog.at_level(logging.INFO, logger="slime.utils.memory_utils"): + with memory_utils.report_peak_memory("actor_train"): + pass + + assert "Peak-Memory" not in caplog.text + + @pytest.mark.unit def test_cuda_selection_does_not_bootstrap_musa(monkeypatch): monkeypatch.setenv("SLIME_ACCELERATOR", "cuda") @@ -89,7 +141,12 @@ def test_cuda_selection_does_not_bootstrap_musa(monkeypatch): @pytest.mark.unit def test_selected_musa_bootstraps_patch_once(monkeypatch): imports = [] - fake_musa = SimpleNamespace(is_available=lambda: True) + fake_musa = SimpleNamespace( + is_available=lambda: True, + reset_peak_memory_stats=lambda: None, + max_memory_allocated=lambda: 0, + max_memory_reserved=lambda: 0, + ) def import_musa_patch(): imports.append("musa_patch") @@ -102,6 +159,7 @@ def import_musa_patch(): assert imports == [] assert accelerator.initialize_accelerator().name == "musa" assert accelerator.initialize_accelerator().name == "musa" + assert accelerator.initialize_accelerator().supports("peak_memory") assert imports == ["musa_patch"] @@ -163,6 +221,7 @@ def test_cuda_backend_uses_torch_cuda_namespace(monkeypatch): assert backend.is_available() assert backend.device_name() == "cuda:1" assert backend.memory_allocated() == 123 + assert backend.supports("peak_memory") @pytest.mark.unit