github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-10-25 01:40:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

[hotfix] fix initialize bug with zero (#442)

Browse Source
This commit is contained in:
Jiarui Fang
2022-03-17 13:16:22 +08:00
committed by GitHub
parent 725a39f4bd
commit 496cbb0760
12 changed files with 87 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
tests/test_zero_data_parallel/common.py
View File

@@ -2,23 +2,38 @@ from functools import partial
import torch
import torch.distributed as dist
import torch.nn as nn
from colossalai.logging import get_dist_logger
from colossalai.utils import checkpoint
from colossalai.zero.sharded_model import ShardedModelV2
LOGGER = get_dist_logger()
LOGGER = get_dist_logger('zero_test')
_ZERO_OPTIMIZER_CONFIG = dict(optimizer_type=torch.optim.Adam, optimizer_config=dict(lr=1e-3))
_ZERO_OFFLOAD_OPTIMIZER_CONFIG = dict(device='cpu', pin_memory=True, buffer_count=5, fast_init=False)
_ZERO_OFFLOAD_PARAM_CONFIG = dict(device='cpu', pin_memory=True, buffer_count=5, buffer_size=1e8, max_in_cpu=1e9)
MP_PARALLEL_CONFIG = dict(fp16=dict(mode=None,), parallel=dict(pipeline=dict(size=1), tensor=dict(size=2, mode=None)))
_ZERO_MODEL_CONFIG = dict(reduce_scatter_bucket_size_mb=25,
fp32_reduce_scatter=False,
offload_config=None,
gradient_predivide_factor=1.0,
shard_param=True,
use_memory_tracer=False)
_ZERO_OPTIMIZER_CONFIG = dict(
optimizer_class=torch.optim.Adam,
cpu_offload=False,
initial_scale=2**32,
min_scale=1,
growth_factor=2,
backoff_factor=0.5,
growth_interval=1000,
hysteresis=2,
max_scale=2**32,
)
ZERO_PARALLEL_CONFIG = dict(fp16=dict(mode=None,),
zero=dict(
optimzer=_ZERO_OPTIMIZER_CONFIG,
offload_optimizer_config=_ZERO_OFFLOAD_OPTIMIZER_CONFIG,
offload_param_config=_ZERO_OFFLOAD_PARAM_CONFIG,
model_config=_ZERO_MODEL_CONFIG,
optimizer_config=_ZERO_OPTIMIZER_CONFIG,
),
parallel=dict(pipeline=dict(size=1), tensor=dict(size=1, mode=None)))
@@ -72,8 +87,8 @@ def check_grads(model, zero_model, loose=False):
def check_params(model, zero_model, loose=False):
for p, zero_p in zip(model.parameters(), zero_model.parameters()):
zero_p = zero_p.clone().to(p.device)
assert p.dtype == zero_p.dtype
assert allclose(p, zero_p, loose=loose)
# assert p.dtype == zero_p.dtype
assert allclose(p.float(), zero_p.float(), loose=loose), f"diff {p.float() - zero_p.float()}"
def check_grads_padding(model, zero_model, loose=False):