[Gemini] add bert for MemtracerWrapper unintests (#1982)

colossalai/gemini/memory_tracer/module_tracer_wrapper.py

@@ -28,6 +28,9 @@ class _Wrapper():
     def show_mem_stats(self):
         self._ophook_list[0].show_mem_stats()
 
+    def named_buffers(self):
+        return self._model.named_buffers()
+
 
 def MemtracerWrapper(model):
     ophook_list = [MemTracerOpHook()]

colossalai/gemini/ophooks/mem_trace_hook.py

@@ -7,6 +7,7 @@ from colossalai.gemini.ophooks import BaseOpHook
 class MemTracerOpHook(BaseOpHook):
     """
     TODO() what if parameters are sharded by multiple submodules.
+    register buff on its father node
     """
 
     def __init__(self):