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

[test] make zero engine test really work (#447)

Browse Source
This commit is contained in:
Jiarui Fang
2022-03-17 17:24:25 +08:00
committed by GitHub
parent bb2790cf0b
commit 0fcfb1e00d
7 changed files with 39 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
tests/test_zero_data_parallel/common.py
View File

@@ -6,6 +6,7 @@ import torch.distributed as dist
from colossalai.logging import get_dist_logger
from colossalai.utils import checkpoint
from colossalai.zero.sharded_model import ShardedModelV2
from colossalai.nn.optimizer import CPUAdam
LOGGER = get_dist_logger('zero_test')
@@ -19,16 +20,16 @@ _ZERO_MODEL_CONFIG = dict(reduce_scatter_bucket_size_mb=25,
use_memory_tracer=False)
_ZERO_OPTIMIZER_CONFIG = dict(
optimizer_class=torch.optim.Adam,
optimizer_class=torch.optim.Adam, #CPUAdam
cpu_offload=False,
initial_scale=2**32,
initial_scale=2**5,
min_scale=1,
growth_factor=2,
backoff_factor=0.5,
growth_interval=1000,
hysteresis=2,
max_scale=2**32,
)
lr=1e-3)
ZERO_PARALLEL_CONFIG = dict(fp16=dict(mode=None,),
zero=dict(