github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-07 12:01:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

[zero] refactor memstats collector (#706)

Browse Source 
* refactor memstats collector

* fix disposable

* polish code
This commit is contained in:
ver217
2022-04-11 10:46:08 +08:00
committed by GitHub
parent 3fc8a204dc
commit ab8c6b4a0e
8 changed files with 44 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/utils/__init__.py
View File

@@ -5,7 +5,7 @@ from .common import (clip_grad_norm_fp32, conditional_context, copy_tensor_paral
ensure_path_exists, free_port, is_dp_rank_0, is_model_parallel_parameter, is_no_pp_or_last_stage,
is_tp_rank_0, is_using_ddp, is_using_pp, is_using_sequence, multi_tensor_applier,
param_is_not_tensor_parallel_duplicate, print_rank_0, switch_virtual_pipeline_parallel_rank,
sync_model_param)
sync_model_param, disposable)
from .data_sampler import DataParallelSampler, get_dataloader
from .gradient_accumulation import accumulate_gradient
from .memory_utils.memory_monitor import report_memory_usage
@@ -19,5 +19,5 @@ __all__ = [
'param_is_not_tensor_parallel_duplicate', 'get_current_device', 'synchronize', 'empty_cache', 'set_to_cuda',
'report_memory_usage', 'Timer', 'MultiTimer', 'multi_tensor_applier', 'accumulate_gradient', 'DataParallelSampler',
'get_dataloader', 'switch_virtual_pipeline_parallel_rank', 'TensorDetector', 'load_checkpoint', 'save_checkpoint',
'ensure_path_exists'
'ensure_path_exists', 'disposable'
]