github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-07-05 11:37:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

Flake8 code restyle

Browse Source
This commit is contained in:
yuxuan-lou 2022-03-09 15:17:01 +08:00 committed by Frank Lee
parent af801cb4df
commit 3b88eb2259
2 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/engine/gradient_handler/_zero_gradient_handler.py
View File

@ -5,7 +5,7 @@ from ._base_gradient_handler import BaseGradientHandler
@GRADIENT_HANDLER.register_module @GRADIENT_HANDLER.register_module
class ZeROGradientHandler(BaseGradientHandler): class ZeROGradientHandler(BaseGradientHandler):
"""A helper class to handle all-reduce operations in a data parallel group. """A helper class to handle all-reduce operations in a data parallel group.
A all-reduce collective communication will be operated in A all-reduce collective communication will be operated in
:func:`handle_gradient` among a data parallel group. :func:`handle_gradient` among a data parallel group.
This class is specialized with ZeRO optimization. This class is specialized with ZeRO optimization.
""" """

3
colossalai/engine/ophooks/_memtracer_ophook.py
View File

@ -14,7 +14,6 @@ import math
class MemTracerOpHook(BaseOpHook): class MemTracerOpHook(BaseOpHook):
""" """
Collect GPU memory usage information Collect GPU memory usage information
:param warmup: This parameter indicates how many iterations to truncate before profiling, defaults to 50 :param warmup: This parameter indicates how many iterations to truncate before profiling, defaults to 50
:type warmup: int :type warmup: int
:param refreshrate: This parameter decides the frequency of write file, defaults to 10 :param refreshrate: This parameter decides the frequency of write file, defaults to 10
@ -112,4 +111,4 @@ class MemTracerOpHook(BaseOpHook):
def save_results(self): def save_results(self):
datafile = f"{self._data_prefix}-{self._rank}.pkl" datafile = f"{self._data_prefix}-{self._rank}.pkl"
self.async_mem_monitor.save(datafile) self.async_mem_monitor.save(datafile)