diff --git a/tests/unit_tests/test_tools_utils.py b/tests/unit_tests/test_tools_utils.py index 4965038763..8cc342baac 100644 --- a/tests/unit_tests/test_tools_utils.py +++ b/tests/unit_tests/test_tools_utils.py @@ -4,10 +4,16 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from unittest.mock import Mock + import pytest import torch -from torchtitan.tools.utils import get_cuda_flash_attention_impl, get_local_device +from torchtitan.tools.utils import ( + GarbageCollection, + get_cuda_flash_attention_impl, + get_local_device, +) class _FakeDeviceModule: @@ -18,6 +24,18 @@ def device_count(self) -> int: return self.num_devices +def test_gc_debug_collects_once_per_step(monkeypatch: pytest.MonkeyPatch) -> None: + gc_collect = Mock() + monkeypatch.setattr("torchtitan.tools.utils.gc.collect", gc_collect) + garbage_collection = GarbageCollection.__new__(GarbageCollection) + garbage_collection.debug = True + + for step in (1, 2): + assert garbage_collection.run(step) + gc_collect.assert_called_once_with(2) + gc_collect.reset_mock() + + def test_get_local_device_uses_local_rank_when_multiple_devices_visible( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/torchtitan/tools/utils.py b/torchtitan/tools/utils.py index 5aac0d35ad..8e71958c69 100644 --- a/torchtitan/tools/utils.py +++ b/torchtitan/tools/utils.py @@ -120,7 +120,6 @@ def run(self, step_count: int) -> bool: "Force GC to perform collection to obtain debug information", generation=2, ) - gc.collect() sl.add_step_tag("gc") return True if step_count > 1 and step_count % self.gc_freq == 0: