[Bug] Avoid duplicate full GC in debug mode

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:

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: