Refine GC logging (#1234)

Improve the code based on the
#1230 (comment)

Author: fegin

torchtitan/tools/utils.py

@@ -50,15 +50,18 @@ def __init__(self, gc_freq: int = 1000, debug: bool = False):
 
     def run(self, step_count: int):
         if self.debug:
-            logger.info("Force GC to perform collection to obtain debug information.")
+            self.collect(
+                "Force GC to perform collection to obtain debug information.",
+                generation=2,
+            )
             gc.collect()
         elif step_count > 1 and step_count % self.gc_freq == 0:
             self.collect("Peforming periodical GC collection.")
 
     @staticmethod
-    def collect(reason: str):
+    def collect(reason: str, generation: int = 1):
         begin = time.monotonic()
-        gc.collect(1)
+        gc.collect(generation)
         logger.info("[GC] %s %.2f seconds.", reason, time.monotonic() - begin)