[Gemini] param_tracer_wrapper and test case (#2009)

colossalai/gemini/memory_tracer/__init__.py

@@ -4,8 +4,9 @@ from .model_data_memtracer import GLOBAL_MODEL_DATA_TRACER    # isort:skip
 from .chunk_memstats_collector import ChunkMemStatsCollector    # isort:skip
 from .static_memstats_collector import StaticMemStatsCollector    # isort:skip
 from .module_tracer_wrapper import MemtracerWrapper    # isort:skip
+from .param_tracer_wrapper import ParamWrapper    # isort:skip
 
 __all__ = [
     'AsyncMemoryMonitor', 'SyncCudaMemoryMonitor', 'MemStatsCollector', 'ChunkMemStatsCollector',
-    'StaticMemStatsCollector', 'GLOBAL_MODEL_DATA_TRACER', 'MemtracerWrapper'
+    'StaticMemStatsCollector', 'GLOBAL_MODEL_DATA_TRACER', 'MemtracerWrapper', 'ParamWrapper'
 ]

colossalai/gemini/memory_tracer/param_tracer_wrapper.py

@@ -0,0 +1,51 @@
+import torch.nn
+
+from colossalai.tensor.colo_parameter import ColoParameter
+from colossalai.tensor.param_op_hook import ParamOpHookManager
+from colossalai.gemini.ophooks import ParamMemHook
+from colossalai.nn.parallel.data_parallel import _cast_float
+
+
+class ParamWrapper():
+
+    def __init__(self, module: torch.nn.Module, dtype: torch.dtype = torch.half):
+        super().__init__()
+        self.module = module
+        self.dtype = dtype
+        self.param_op_hook = ParamMemHook()
+
+        for p in module.parameters():
+            assert isinstance(p, ColoParameter)
+            p.data = p.data.to(dtype)
+
+        self._cast_buffers_to_cuda_dtype()
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
+
+    def _pre_forward(self):
+        self.param_op_hook.mem_monitor.start()
+
+    def forward(self, *args, **kwargs):
+        args, kwargs = _cast_float(args, self.dtype), _cast_float(kwargs, self.dtype)
+        self.module.zero_grad(set_to_none=True)
+        self._pre_forward()
+        with ParamOpHookManager.use_hooks(self.param_op_hook):
+            outputs = self.module(*args, **kwargs)
+        return outputs
+
+    def backward(self, loss):
+        with self.param_op_hook.switch_to_backward(), ParamOpHookManager.use_hooks(self.param_op_hook):
+            loss.backward()
+        self._post_backward()
+
+    def _post_backward(self):
+        cuda_volume = self.param_op_hook.mem_monitor.finish()
+        last_model_data = self.param_op_hook._model_data_list[-1]
+        self.param_op_hook._non_model_data_list.append(cuda_volume - last_model_data)
+
+    def _cast_buffers_to_cuda_dtype(self):
+        for buffer in self.module.buffers():
+            buffer.data = buffer.cuda()
+            if torch.is_floating_point(buffer):
+                buffer.data = buffer.data.to(self.dtype)