[Gemini] update the non model data record method in runtime memory tracer (#2128)

colossalai/zero/utils/gemini_hook.py

@@ -32,7 +32,7 @@ class GeminiZeROHook(ColoParamOpHook):
         self._gemini_manager.adjust_layout(chunks)
         for chunk in chunks:
             self._chunk_manager.access_chunk(chunk)
-        self._gemini_manager.sample_model_data()
+        self._gemini_manager.record_model_data_volume()
 
     def post_op(self, params):
         params = [p for p in params if not getattr(p, '_ddp_to_ignore', False)]

colossalai/zero/utils/zero_hook.py

@@ -67,7 +67,7 @@ class ZeroHook(BaseOpHook):
 
         # record model data statistics
         if self._memstarts_collector:
-            self._memstarts_collector.sample_model_data()
+            self._memstarts_collector.record_model_data_volume()
 
     def pre_fwd_exec(self, module: torch.nn.Module, *args):
         self.adjust_module_data(module)