Feat(PythonGC): Do garbage collection manually (#326)

* feat:add gc control

* feat:add gc control

* feat:add gc control

* feat:add gc

* re-lint

internlm/initialize/launch.py

@@ -2,6 +2,7 @@
 # -*- encoding: utf-8 -*-
 
 import argparse
+import gc
 import os
 from pathlib import Path
 from typing import Dict, Union
@@ -446,6 +447,8 @@ def initialize_distributed_env(
         master_port (str): The master port for distributed training. 8888 by default.
         seed (int, optional): Specified random seed for every process. 1024 by default.
     """
+    # close automatic garbage collection
+    gc.disable()
 
     torch.cuda.empty_cache()
 

internlm/utils/gputest.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 # -*- encoding: utf-8 -*-
 
+import gc
 import math
 import socket
 
@@ -41,6 +42,8 @@ def empty_cache_and_diag(batch_count, interval=50):
                 bench_net()
         # do empty_cache after the bench
         torch.cuda.empty_cache()
+        # do garbage collection
+        gc.collect()
 
 
 def benchmark_forward(