github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-08 04:24:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

add interleaved pipeline, fix naive amp and update pipeline model initializer (#80)

Browse Source
This commit is contained in:
ver217
2021-12-20 23:26:19 +08:00
committed by GitHub
parent 91c327cb44
commit 8f02a88db2
17 changed files with 544 additions and 170 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
colossalai/utils/__init__.py
View File

@@ -3,7 +3,7 @@ from .common import (print_rank_0, sync_model_param_in_dp, is_dp_rank_0,
is_tp_rank_0, is_no_pp_or_last_stage, is_using_ddp,
is_using_pp, conditional_context, is_model_parallel_parameter,
clip_grad_norm_fp32, count_zeros_fp32, copy_tensor_parallel_attributes,
param_is_not_tensor_parallel_duplicate)
param_is_not_tensor_parallel_duplicate, switch_virtual_pipeline_parallel_rank)
from .cuda import get_current_device, synchronize, empty_cache, set_to_cuda
from .memory import report_memory_usage
from .timer import MultiTimer, Timer
@@ -22,5 +22,6 @@ __all__ = ['checkpoint',
'Timer', 'MultiTimer',
'multi_tensor_applier',
'accumulate_gradient',
'DataParallelSampler', 'get_dataloader'
'DataParallelSampler', 'get_dataloader',
'switch_virtual_pipeline_parallel_rank'
]