github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2026-05-12 17:53:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixed docstring in colossalai (#171)

Browse Source
This commit is contained in:
HELSON
2022-01-21 10:44:30 +08:00
committed by GitHub
parent e2089c5c15
commit 0f8c7f9804
77 changed files with 983 additions and 603 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
colossalai/utils/memory.py
View File

@@ -12,34 +12,36 @@ from colossalai.logging import get_dist_logger
def bytes_to_GB(val, decimal=2):
'''A byte-to-Gigabyte converter, defaultly using binary notation.
"""A byte-to-Gigabyte converter, defaultly using binary notation.
:param val: X bytes to convert
:return: X' GB
'''
"""
return round(val / (1024 * 1024 * 1024), decimal)
def bytes_to_MB(val, decimal=2):
'''A byte-to-Megabyte converter, defaultly using binary notation.
"""A byte-to-Megabyte converter, defaultly using binary notation.
:param val: X bytes to convert
:return: X' MB
'''
"""
return round(val / (1024 * 1024), decimal)
def report_memory_usage(message, logger=None, report_cpu=False):
'''Calculate and print RAM usage (in GB)
"""Calculate and print RAM usage (in GB)
:param message: a prefix message to add in the log
:param message: A prefix message to add in the log
:type message: str
:param logger: an instance of :class:`colossalai.logging.DistributedLogger`
:type logger: :class:`colossalai.logging.DistributedLogger`
:param report_cpu: whether to report CPU memory
:type report_cpu: bool
:raises EnvironmentError: raise error if no distributed environment has been initialized
'''
:param logger: An instance of :class:`colossalai.logging.DistributedLogger`
:type logger: :class:`colossalai.logging.DistributedLogger`, optional
:param report_cpu: Whether to report CPU memory
:type report_cpu: bool, optional
:raises EnvironmentError: Raise error if no distributed environment has been initialized
"""
if not gpc.is_initialized(ParallelMode.GLOBAL):
raise EnvironmentError("No distributed environment is initialized")
gpu_allocated = bytes_to_MB(torch.cuda.memory_allocated())
gpu_max_allocated = bytes_to_MB(torch.cuda.max_memory_allocated())