polish utils docstring (#620)

colossalai/utils/memory_tracer/async_memtracer.py

@@ -11,19 +11,13 @@ from colossalai.utils import get_current_device
 class AsyncMemoryMonitor:
     """
     An Async Memory Monitor runing during computing. Sampling memory usage of the current GPU
-    at interval of 1/(10**power) sec.
+    at interval of `1/(10**power)` sec.
 
     The idea comes from Runtime Memory Tracer of PatrickStar
-    PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management
-    https://arxiv.org/abs/2108.05818
-    
-    :param power: the power of time interval, defaults to 10
-    :type power: int
+    `PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management`_
 
-    Usage:
-    ::
+    Usage::
 
-        ```python
         async_mem_monitor = AsyncMemoryMonitor()
         input = torch.randn(2, 20).cuda()
         OP1 = torch.nn.Linear(20, 30).cuda()
@@ -36,7 +30,12 @@ class AsyncMemoryMonitor:
         output = OP2(output)
         async_mem_monitor.finish()
         async_mem_monitor.save('log.pkl')
-        ```
+
+    Args:
+        power (int, optional): the power of time interva. Defaults to 10.
+
+    .. _PatrickStar\: Parallel Training of Pre-trained Models via Chunk-based Memory Management:
+        https://arxiv.org/abs/2108.05818
     """
 
     def __init__(self, power: int = 10):

colossalai/utils/profiler/mem_profiler.py

@@ -12,6 +12,8 @@ class MemProfiler(BaseProfiler):
     To use this profiler, you need to pass an `engine` instance. And the usage is same like
     CommProfiler.
 
+    Usage::
+
         mm_prof = MemProfiler(engine)
         with ProfilerContext([mm_prof]) as prof:
             writer = SummaryWriter("mem")
@@ -36,11 +38,7 @@ class MemProfiler(BaseProfiler):
     def to_tensorboard(self, writer: SummaryWriter) -> None:
         stats = self._mem_tracer.async_mem_monitor.state_dict['mem_stats']
         for info, i in enumerate(stats):
-            writer.add_scalar(
-                "memory_usage/GPU",
-                info,
-                i
-            )
+            writer.add_scalar("memory_usage/GPU", info, i)
 
     def to_file(self, data_file: Path) -> None:
         self._mem_tracer.save_results(data_file)

colossalai/utils/profiler/prof_utils.py

@@ -70,12 +70,10 @@ class BaseProfiler(ABC):
 
 
 class ProfilerContext(object):
-    """
-    Profiler context manager
-    Usage:
-    ::
+    """Profiler context manager
+
+    Usage::
 
-        ```python
         world_size = 4
         inputs = torch.randn(10, 10, dtype=torch.float32, device=get_current_device())
         outputs = torch.empty(world_size, 10, 10, dtype=torch.float32, device=get_current_device())
@@ -92,7 +90,6 @@ class ProfilerContext(object):
             dist.reduce(inputs, 0)
 
         prof.show()
-        ```
     """
 
     def __init__(self, profilers: List[BaseProfiler] = None, enable: bool = True):