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

[misc] update pre-commit and run all files (#4752)

Browse Source 
* [misc] update pre-commit

* [misc] run pre-commit

* [misc] remove useless configuration files

* [misc] ignore cuda for clang-format
This commit is contained in:
Hongxin Liu
2023-09-19 14:20:26 +08:00
committed by GitHub
parent 3c6b831c26
commit 079bf3cb26
1268 changed files with 50037 additions and 38444 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
tests/test_zero/test_gemini/test_chunk_mgrv2.py
View File

@@ -12,54 +12,53 @@ CUDA_MEM_1 = {False: 0, True: 1024}
CPU_MEM = {True: {True: 0, False: 0}, False: {True: 512, False: 0}}
@parameterize('keep_gathered', [True, False])
@parameterize('pin_memory', [True, False])
@parameterize("keep_gathered", [True, False])
@parameterize("pin_memory", [True, False])
def exam_chunk_memory(keep_gathered, pin_memory):
params = [ColoTensor(torch.rand(8, 8)) for _ in range(3)]
config = {2: dict(chunk_size=128, keep_gathered=keep_gathered)}
chunk_manager = ChunkManager(config)
assert chunk_manager.total_mem['cpu'] == 0
assert chunk_manager.total_mem['cuda'] == 0
assert chunk_manager.total_mem["cpu"] == 0
assert chunk_manager.total_mem["cuda"] == 0
process_group = _get_default_group()
for p in params:
chunk_manager.register_tensor(p, 'param', 2, process_group, pin_memory=pin_memory)
chunk_manager.register_tensor(p, "param", 2, process_group, pin_memory=pin_memory)
chunk_manager.close_all_groups()
assert chunk_manager.total_mem['cpu'] == CPU_MEM[keep_gathered][pin_memory]
assert chunk_manager.total_mem['cuda'] == CUDA_MEM_0[keep_gathered]
assert chunk_manager.total_mem["cpu"] == CPU_MEM[keep_gathered][pin_memory]
assert chunk_manager.total_mem["cuda"] == CUDA_MEM_0[keep_gathered]
chunks = chunk_manager.get_chunks(params)
for chunk in chunks:
chunk_manager.access_chunk(chunk)
assert chunk_manager.total_mem['cpu'] == CPU_MEM[keep_gathered][pin_memory]
assert chunk_manager.total_mem['cuda'] == CUDA_MEM_0[True]
assert chunk_manager.total_mem["cpu"] == CPU_MEM[keep_gathered][pin_memory]
assert chunk_manager.total_mem["cuda"] == CUDA_MEM_0[True]
for chunk in chunks:
chunk_manager.release_chunk(chunk)
assert chunk_manager.total_mem['cpu'] == CPU_MEM[keep_gathered][pin_memory]
assert chunk_manager.total_mem['cuda'] == CUDA_MEM_0[keep_gathered]
assert chunk_manager.total_mem["cpu"] == CPU_MEM[keep_gathered][pin_memory]
assert chunk_manager.total_mem["cuda"] == CUDA_MEM_0[keep_gathered]
for chunk in chunks:
chunk_manager.move_chunk(chunk, torch.device('cpu'))
assert chunk_manager.total_mem['cpu'] == CPU_MEM[keep_gathered][True]
assert chunk_manager.total_mem['cuda'] == CUDA_MEM_1[keep_gathered]
chunk_manager.move_chunk(chunk, torch.device("cpu"))
assert chunk_manager.total_mem["cpu"] == CPU_MEM[keep_gathered][True]
assert chunk_manager.total_mem["cuda"] == CUDA_MEM_1[keep_gathered]
def run_dist(rank, world_size, port):
colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
colossalai.launch(config={}, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
exam_chunk_memory()
@pytest.mark.dist
@pytest.mark.parametrize('world_size', [2])
@pytest.mark.parametrize("world_size", [2])
@rerun_if_address_is_in_use()
def test_chunk_manager(world_size):
spawn(run_dist, world_size)
if __name__ == '__main__':
if __name__ == "__main__":
test_chunk_manager(2)

42
tests/test_zero/test_gemini/test_chunkv2.py
View File

@@ -31,26 +31,28 @@ def check_equal(param, param_cp):
return torch.equal(temp, param_cp.data)
@parameterize('init_device', [None, torch.device('cpu')])
@parameterize('keep_gathered', [True, False])
@parameterize('pin_memory', [True, False])
@parameterize("init_device", [None, torch.device("cpu")])
@parameterize("keep_gathered", [True, False])
@parameterize("pin_memory", [True, False])
def exam_chunk_basic(init_device, keep_gathered, pin_memory):
world_size = torch.distributed.get_world_size()
pg = _get_default_group()
my_chunk = Chunk(chunk_size=1024,
process_group=pg,
dtype=torch.float32,
init_device=init_device,
cpu_shard_init=True,
keep_gathered=keep_gathered,
pin_memory=pin_memory)
my_chunk = Chunk(
chunk_size=1024,
process_group=pg,
dtype=torch.float32,
init_device=init_device,
cpu_shard_init=True,
keep_gathered=keep_gathered,
pin_memory=pin_memory,
)
param_list = []
param_cp_list = []
add_param(param_list, param_cp_list, 8, 8, 8, device='cuda')
add_param(param_list, param_cp_list, 8, 8, 8, device="cuda")
add_param(param_list, param_cp_list, 4, 4)
add_param(param_list, param_cp_list, 4, 8, 2, device='cuda')
add_param(param_list, param_cp_list, 4, 8, 2, device="cuda")
add_param(param_list, param_cp_list, 1, 1, 5)
for param in param_list:
@@ -62,12 +64,12 @@ def exam_chunk_basic(init_device, keep_gathered, pin_memory):
if keep_gathered is False:
assert my_chunk.cpu_shard.size(0) == 1024 // world_size
assert my_chunk.device_type == 'cpu'
assert my_chunk.device_type == "cpu"
assert my_chunk.can_move
my_chunk.shard_move(get_current_device())
else:
assert my_chunk.cuda_global_chunk.size(0) == 1024
assert my_chunk.device_type == 'cuda'
assert my_chunk.device_type == "cuda"
assert not my_chunk.can_move
assert dist_sum(my_chunk.valid_end) == my_chunk.utilized_size
@@ -75,7 +77,7 @@ def exam_chunk_basic(init_device, keep_gathered, pin_memory):
assert not flag, "has_inf_or_nan is {}".format(flag)
my_chunk.access_chunk()
assert my_chunk.device_type == 'cuda'
assert my_chunk.device_type == "cuda"
for param, param_cp in zip(param_list, param_cp_list):
check_equal(param, param_cp)
@@ -97,25 +99,25 @@ def exam_chunk_basic(init_device, keep_gathered, pin_memory):
if keep_gathered is False:
assert my_chunk.cuda_shard.size(0) == 1024 // world_size
assert my_chunk.device_type == 'cuda'
assert my_chunk.device_type == "cuda"
assert my_chunk.can_move
else:
assert my_chunk.cuda_global_chunk.size(0) == 1024
assert my_chunk.device_type == 'cuda'
assert my_chunk.device_type == "cuda"
assert not my_chunk.can_move
def run_dist(rank, world_size, port):
colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
colossalai.launch(config={}, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
exam_chunk_basic()
@pytest.mark.dist
@pytest.mark.parametrize('world_size', [1, 2, 4])
@pytest.mark.parametrize("world_size", [1, 2, 4])
@rerun_if_address_is_in_use()
def test_chunk_function(world_size):
spawn(run_dist, world_size)
if __name__ == '__main__':
if __name__ == "__main__":
test_chunk_function(4)

41
tests/test_zero/test_gemini/test_fwd_bwd.py
View File

@@ -16,21 +16,10 @@ from tests.components_to_test import run_fwd_bwd
from tests.components_to_test.registry import non_distributed_component_funcs
PLACEMENT_CONFIGS = [
{
'placement_policy': 'static',
'shard_param_frac': 0.0
}, # zero2
{
'placement_policy': 'static',
'shard_param_frac': 1.0
}, # zero3
{
'placement_policy': 'static',
'shard_param_frac': 0.5
}, # zero3-half
{
'placement_policy': 'auto'
}
{"placement_policy": "static", "shard_param_frac": 0.0}, # zero2
{"placement_policy": "static", "shard_param_frac": 1.0}, # zero3
{"placement_policy": "static", "shard_param_frac": 0.5}, # zero3-half
{"placement_policy": "auto"},
]
@@ -41,14 +30,14 @@ def check_grad(model: GeminiDDP, torch_model: torch.nn.Module):
for chunk in chunk_list:
chunk_manager.access_chunk(chunk)
for (p0, p1) in zip(model.parameters(), torch_model.parameters()):
for p0, p1 in zip(model.parameters(), torch_model.parameters()):
assert_close(p0, p1.grad, rtol=1e-3, atol=5e-5)
@parameterize('placement_config', PLACEMENT_CONFIGS)
@parameterize('keep_gather', [False, True])
@parameterize('model_name', ['gpt2', 'bert', 'albert'])
@parameterize('use_grad_checkpoint', [False, True])
@parameterize("placement_config", PLACEMENT_CONFIGS)
@parameterize("keep_gather", [False, True])
@parameterize("model_name", ["gpt2", "bert", "albert"])
@parameterize("use_grad_checkpoint", [False, True])
def exam_gpt_fwd_bwd(
placement_config,
keep_gather,
@@ -69,14 +58,14 @@ def exam_gpt_fwd_bwd(
world_size = torch.distributed.get_world_size()
config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100)
config_dict[world_size]['chunk_size'] = 5000
config_dict[world_size]['keep_gathered'] = keep_gather
config_dict[world_size]["chunk_size"] = 5000
config_dict[world_size]["keep_gathered"] = keep_gather
model = GeminiDDP(model, config_dict, init_device, pin_memory=True, **placement_config)
optimizer = HybridAdam(model.parameters(), lr=1e-3)
zero_optim = GeminiOptimizer(optimizer, model, initial_scale=1)
rank = dist.get_rank()
amp_config = dict(opt_level='O2', keep_batchnorm_fp32=False, loss_scale=1)
amp_config = dict(opt_level="O2", keep_batchnorm_fp32=False, loss_scale=1)
torch_optim = torch.optim.Adam(torch_model.parameters(), lr=1e-3)
torch_model, torch_optim = convert_to_apex_amp(torch_model, torch_optim, amp_config)
torch_model = DDP(torch_model, device_ids=[rank])
@@ -105,16 +94,16 @@ def exam_gpt_fwd_bwd(
def run_dist(rank, world_size, port):
config = {}
colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
colossalai.launch(config=config, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
exam_gpt_fwd_bwd()
@pytest.mark.dist
@pytest.mark.parametrize('world_size', [1, 4])
@pytest.mark.parametrize("world_size", [1, 4])
@rerun_if_address_is_in_use()
def test_gpt(world_size):
spawn(run_dist, world_size)
if __name__ == '__main__':
if __name__ == "__main__":
test_gpt(4)

45
tests/test_zero/test_gemini/test_gemini_use_rmt.py
View File

@@ -14,10 +14,10 @@ from tests.components_to_test.registry import non_distributed_component_funcs
# run gemini use the runtime memory tracer
@parameterize('placement_policy', ['auto'])
@parameterize('keep_gather', [False])
@parameterize('model_name', ['repeated_computed_layers', 'bert', 'albert', 'gpt2'])
@parameterize('use_grad_checkpoint', [False, True])
@parameterize("placement_policy", ["auto"])
@parameterize("keep_gather", [False])
@parameterize("model_name", ["repeated_computed_layers", "bert", "albert", "gpt2"])
@parameterize("use_grad_checkpoint", [False, True])
def run_gemini_use_rmt(placement_policy, keep_gather, model_name: str, use_grad_checkpoint: bool = False):
set_seed(42)
get_components_func = non_distributed_component_funcs.get_callable(model_name)
@@ -25,7 +25,7 @@ def run_gemini_use_rmt(placement_policy, keep_gather, model_name: str, use_grad_
model = model_builder(use_grad_checkpoint).cuda()
print(f'model_name {model_name}')
print(f"model_name {model_name}")
runtime_mem_tracer = RuntimeMemTracer(model)
for i, (input_ids, label) in enumerate(train_dataloader):
if i > 0:
@@ -37,17 +37,17 @@ def run_gemini_use_rmt(placement_policy, keep_gather, model_name: str, use_grad_
run_fwd_bwd(runtime_mem_tracer, input_ids, label, criterion, runtime_mem_tracer)
memstats = runtime_mem_tracer.memstats()
runtime_tracer_non_model_data = runtime_mem_tracer._memstats._non_model_data_cuda_list
print('runtime tracer non model data points: ', len(runtime_tracer_non_model_data))
print('runtime tracer: ', runtime_tracer_non_model_data)
print("runtime tracer non model data points: ", len(runtime_tracer_non_model_data))
print("runtime tracer: ", runtime_tracer_non_model_data)
print([memstats.param_used_step(p) for p in model.parameters()])
if model_name == 'repeated_computed_layers':
if model_name == "repeated_computed_layers":
for idx, p in enumerate(model.parameters()):
step_list = memstats.param_used_step(p)
if idx < 4:
assert len(step_list) == 4
if model_name == 'repeated_computed_layers':
if model_name == "repeated_computed_layers":
for idx, p in enumerate(model.parameters()):
step_list = memstats.param_used_step(p)
if idx < 4:
@@ -55,13 +55,11 @@ def run_gemini_use_rmt(placement_policy, keep_gather, model_name: str, use_grad_
world_size = torch.distributed.get_world_size()
config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100)
config_dict[world_size]['chunk_size'] = 5000
config_dict[world_size]['keep_gathered'] = keep_gather
model = GeminiDDP(model,
chunk_config_dict=config_dict,
placement_policy=placement_policy,
pin_memory=True,
memstats=memstats)
config_dict[world_size]["chunk_size"] = 5000
config_dict[world_size]["keep_gathered"] = keep_gather
model = GeminiDDP(
model, chunk_config_dict=config_dict, placement_policy=placement_policy, pin_memory=True, memstats=memstats
)
set_seed(dist.get_rank())
for i, (input_ids, label) in enumerate(train_dataloader):
@@ -73,29 +71,30 @@ def run_gemini_use_rmt(placement_policy, keep_gather, model_name: str, use_grad_
input_ids, label = input_ids.cuda(), label.cuda()
set_seed(42)
loss = run_fwd_bwd(model, input_ids, label, criterion, model)
run_fwd_bwd(model, input_ids, label, criterion, model)
gemini_non_model_data = model.gemini_manager._mem_stats_collector._memstats.non_model_data_list('cuda')
gemini_non_model_data = model.gemini_manager._mem_stats_collector._memstats.non_model_data_list("cuda")
# print('gemini non model data:', gemini_non_model_data)
assert len(gemini_non_model_data) == len(runtime_tracer_non_model_data), \
f'model_name {model_name} {len(gemini_non_model_data)} vs {len(runtime_tracer_non_model_data)}'
assert len(gemini_non_model_data) == len(
runtime_tracer_non_model_data
), f"model_name {model_name} {len(gemini_non_model_data)} vs {len(runtime_tracer_non_model_data)}"
def run_dist(rank, world_size, port):
config = {}
colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
colossalai.launch(config=config, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
run_gemini_use_rmt()
@pytest.mark.skip("this is not used")
@pytest.mark.dist
@pytest.mark.parametrize('world_size', [1, 4])
@pytest.mark.parametrize("world_size", [1, 4])
@rerun_if_address_is_in_use()
def test_gemini_use_rmt(world_size):
spawn(run_dist, world_size)
if __name__ == '__main__':
if __name__ == "__main__":
test_gemini_use_rmt(1)

63
tests/test_zero/test_gemini/test_grad_clip.py
View File

@@ -16,26 +16,24 @@ from tests.components_to_test.registry import non_distributed_component_funcs
PLACEMENT_CONFIGS = [
{
'placement_policy': 'static',
'shard_param_frac': 0.0,
'offload_optim_frac': 0.0,
'offload_param_frac': 0.0
}, # zero2
"placement_policy": "static",
"shard_param_frac": 0.0,
"offload_optim_frac": 0.0,
"offload_param_frac": 0.0,
}, # zero2
{
'placement_policy': 'static',
'shard_param_frac': 0.0,
'offload_optim_frac': 1.0,
'offload_param_frac': 0.0
}, # zero2-offload
"placement_policy": "static",
"shard_param_frac": 0.0,
"offload_optim_frac": 1.0,
"offload_param_frac": 0.0,
}, # zero2-offload
{
'placement_policy': 'static',
'shard_param_frac': 0.0,
'offload_optim_frac': 0.5,
'offload_param_frac': 0.0
}, # zero2-offload-half
{
'placement_policy': 'auto'
}
"placement_policy": "static",
"shard_param_frac": 0.0,
"offload_optim_frac": 0.5,
"offload_param_frac": 0.0,
}, # zero2-offload-half
{"placement_policy": "auto"},
]
@@ -52,15 +50,15 @@ def check_param(model: GeminiDDP, torch_model: torch.nn.Module):
assert_close(value, temp_zero_value, rtol=1e-3, atol=4e-3)
@parameterize('placement_config', PLACEMENT_CONFIGS)
@parameterize('model_name', ['gpt2'])
@parameterize("placement_config", PLACEMENT_CONFIGS)
@parameterize("model_name", ["gpt2"])
def exam_grad_clipping(placement_config, model_name: str):
set_seed(1912)
get_components_func = non_distributed_component_funcs.get_callable(model_name)
model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func()
torch_model = model_builder().cuda()
amp_config = dict(opt_level='O2', keep_batchnorm_fp32=False, loss_scale=32)
amp_config = dict(opt_level="O2", keep_batchnorm_fp32=False, loss_scale=32)
torch_optim = torch.optim.Adam(torch_model.parameters(), lr=1e-3)
torch_model, torch_optim = convert_to_apex_amp(torch_model, torch_optim, amp_config)
torch_model = DDP(torch_model, device_ids=[dist.get_rank()])
@@ -72,18 +70,16 @@ def exam_grad_clipping(placement_config, model_name: str):
world_size = torch.distributed.get_world_size()
config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100)
config_dict[world_size]['chunk_size'] = 5000
config_dict[world_size]['keep_gathered'] = False
if placement_config['placement_policy'] != 'cuda':
init_device = torch.device('cpu')
config_dict[world_size]["chunk_size"] = 5000
config_dict[world_size]["keep_gathered"] = False
if placement_config["placement_policy"] != "cuda":
init_device = torch.device("cpu")
else:
init_device = None
model = GeminiDDP(model,
chunk_config_dict=config_dict,
chunk_init_device=init_device,
pin_memory=True,
**placement_config)
model = GeminiDDP(
model, chunk_config_dict=config_dict, chunk_init_device=init_device, pin_memory=True, **placement_config
)
optimizer = HybridAdam(model.parameters(), lr=1e-3)
zero_optim = GeminiOptimizer(optimizer, model, initial_scale=32, clipping_norm=1.0)
@@ -106,6 +102,7 @@ def exam_grad_clipping(placement_config, model_name: str):
assert_close(torch_loss, loss)
import apex.amp as apex_amp
torch.nn.utils.clip_grad_norm_(apex_amp.master_params(torch_optim), 1.0)
torch_optim.step()
zero_optim.step()
@@ -115,16 +112,16 @@ def exam_grad_clipping(placement_config, model_name: str):
def run_dist(rank, world_size, port):
config = {}
colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
colossalai.launch(config=config, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
exam_grad_clipping()
@pytest.mark.dist
@pytest.mark.parametrize('world_size', [1, 2])
@pytest.mark.parametrize("world_size", [1, 2])
@rerun_if_address_is_in_use()
def test_grad_clip(world_size):
spawn(run_dist, world_size)
if __name__ == '__main__':
if __name__ == "__main__":
test_grad_clip(2)

37
tests/test_zero/test_gemini/test_inference.py
View File

@@ -18,21 +18,10 @@ from tests.components_to_test import run_fwd_bwd
from tests.components_to_test.registry import non_distributed_component_funcs
PLACEMENT_CONFIGS = [
{
'placement_policy': 'static',
'shard_param_frac': 0.0
}, # zero2
{
'placement_policy': 'static',
'shard_param_frac': 1.0
}, # zero3
{
'placement_policy': 'static',
'shard_param_frac': 0.5
}, # zero3-half
{
'placement_policy': 'auto'
}
{"placement_policy": "static", "shard_param_frac": 0.0}, # zero2
{"placement_policy": "static", "shard_param_frac": 1.0}, # zero3
{"placement_policy": "static", "shard_param_frac": 0.5}, # zero3-half
{"placement_policy": "auto"},
]
@@ -52,8 +41,8 @@ def check_param(model: GeminiDDP, torch_model: torch.nn.Module):
def multi_chunk_init(model: torch.nn.Module, placement_config: dict):
world_size = dist.get_world_size()
config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100)
config_dict[world_size]['chunk_size'] = 5000
config_dict[world_size]['keep_gathered'] = False
config_dict[world_size]["chunk_size"] = 5000
config_dict[world_size]["keep_gathered"] = False
model = GeminiDDP(model, config_dict, pin_memory=True, **placement_config)
return model
@@ -63,16 +52,16 @@ def single_chunk_init(model: torch.nn.Module, placement_config: dict):
return model
@parameterize('placement_config', PLACEMENT_CONFIGS)
@parameterize('model_name', ['gpt2'])
@parameterize('model_init_func', [single_chunk_init, multi_chunk_init])
@parameterize("placement_config", PLACEMENT_CONFIGS)
@parameterize("model_name", ["gpt2"])
@parameterize("model_init_func", [single_chunk_init, multi_chunk_init])
def exam_inference(placement_config: dict, model_name: str, model_init_func: Callable):
set_seed(19360226)
get_components_func = non_distributed_component_funcs.get_callable(model_name)
model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func()
torch_model = model_builder().cuda()
amp_config = dict(opt_level='O2', keep_batchnorm_fp32=False, loss_scale=128)
amp_config = dict(opt_level="O2", keep_batchnorm_fp32=False, loss_scale=128)
torch_optim = torch.optim.Adam(torch_model.parameters(), lr=1e-3)
torch_model, torch_optim = convert_to_apex_amp(torch_model, torch_optim, amp_config)
torch_model = DDP(torch_model, device_ids=[dist.get_rank()])
@@ -121,16 +110,16 @@ def exam_inference(placement_config: dict, model_name: str, model_init_func: Cal
def run_dist(rank, world_size, port):
config = {}
colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
colossalai.launch(config=config, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
exam_inference()
@pytest.mark.dist
@pytest.mark.parametrize('world_size', [1, 4])
@pytest.mark.parametrize("world_size", [1, 4])
@rerun_if_address_is_in_use()
def test_inference(world_size):
spawn(run_dist, world_size)
if __name__ == '__main__':
if __name__ == "__main__":
test_inference(1)

106
tests/test_zero/test_gemini/test_optim.py
View File

@@ -16,50 +16,30 @@ from tests.components_to_test import run_fwd_bwd
from tests.components_to_test.registry import non_distributed_component_funcs
PLACEMENT_CONFIGS = [
{"placement_policy": "static", "shard_param_frac": 0.0, "offload_optim_frac": 0.0}, # zero2
{"placement_policy": "static", "shard_param_frac": 0.0, "offload_optim_frac": 1.0}, # zero2-offload
{"placement_policy": "static", "shard_param_frac": 0.0, "offload_optim_frac": 0.5}, # zero2-offload-half
{"placement_policy": "static", "shard_param_frac": 1.0}, # zero3
{"placement_policy": "static", "shard_param_frac": 0.5}, # zero3-half
{
'placement_policy': 'static',
'shard_param_frac': 0.0,
'offload_optim_frac': 0.0
}, # zero2
{
'placement_policy': 'static',
'shard_param_frac': 0.0,
'offload_optim_frac': 1.0
}, # zero2-offload
{
'placement_policy': 'static',
'shard_param_frac': 0.0,
'offload_optim_frac': 0.5
}, # zero2-offload-half
{
'placement_policy': 'static',
'shard_param_frac': 1.0
}, # zero3
{
'placement_policy': 'static',
'shard_param_frac': 0.5
}, # zero3-half
{
'placement_policy': 'static',
'shard_param_frac': 1.0,
'offload_optim_frac': 1.0,
'offload_param_frac': 1.0
}, # zero3-offload-all
{
'placement_policy': 'auto'
}
"placement_policy": "static",
"shard_param_frac": 1.0,
"offload_optim_frac": 1.0,
"offload_param_frac": 1.0,
}, # zero3-offload-all
{"placement_policy": "auto"},
]
# this model is large enough to slice to chunks
TEST_MODELS = ['gpt2']
TEST_MODELS = ["gpt2"]
# these models are too small, all parameters in these models are compacted into one chunk
EXAMPLE_MODELS = ['albert', 'beit', 'bert', 'hanging_param_model', 'nested_model', 'repeated_computed_layers']
EXAMPLE_MODELS = ["albert", "beit", "bert", "hanging_param_model", "nested_model", "repeated_computed_layers"]
# bfloat16 cannot represent them exactly
BF16_IGNORED_KEYS = [
'albert.embeddings.word_embeddings.weight',
'albert.embeddings.position_embeddings.weight',
'masked_bias',
"albert.embeddings.word_embeddings.weight",
"albert.embeddings.position_embeddings.weight",
"masked_bias",
]
@@ -78,23 +58,25 @@ def check_param(model: GeminiDDP, torch_model: torch.nn.Module, dtype: torch.dty
if dtype is torch.bfloat16:
rtol, atol = 4e-3, 8e-3
# debug_print([0], "max range: ", key, torch.max(torch.abs(value - temp_zero_value)))
assert_close(value.float(),
temp_zero_value.float(),
rtol=rtol,
atol=atol,
msg=lambda s: s + f'\n{key}\n{temp_zero_value.dtype}')
assert_close(
value.float(),
temp_zero_value.float(),
rtol=rtol,
atol=atol,
msg=lambda s: s + f"\n{key}\n{temp_zero_value.dtype}",
)
@parameterize('placement_config', PLACEMENT_CONFIGS)
@parameterize('model_name', TEST_MODELS)
@parameterize('mixed_precision', [torch.half, torch.bfloat16])
@parameterize("placement_config", PLACEMENT_CONFIGS)
@parameterize("model_name", TEST_MODELS)
@parameterize("mixed_precision", [torch.half, torch.bfloat16])
def exam_model_step(placement_config, model_name: str, mixed_precision: torch.dtype):
set_seed(42)
get_components_func = non_distributed_component_funcs.get_callable(model_name)
model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func()
torch_model = model_builder().cuda()
amp_config = dict(opt_level='O2', keep_batchnorm_fp32=False, loss_scale=128)
amp_config = dict(opt_level="O2", keep_batchnorm_fp32=False, loss_scale=128)
torch_optim = torch.optim.Adam(torch_model.parameters(), lr=1e-3)
torch_model, torch_optim = convert_to_apex_amp(torch_model, torch_optim, amp_config)
torch_model = DDP(torch_model, device_ids=[dist.get_rank()])
@@ -106,8 +88,8 @@ def exam_model_step(placement_config, model_name: str, mixed_precision: torch.dt
world_size = torch.distributed.get_world_size()
config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100)
config_dict[world_size]['chunk_size'] = 5000
config_dict[world_size]['keep_gathered'] = False
config_dict[world_size]["chunk_size"] = 5000
config_dict[world_size]["keep_gathered"] = False
model = GeminiDDP(model, config_dict, **placement_config, mixed_precision=mixed_precision)
optimizer = HybridAdam(model.parameters(), lr=1e-3)
@@ -135,16 +117,16 @@ def exam_model_step(placement_config, model_name: str, mixed_precision: torch.dt
check_param(model, torch_model, mixed_precision)
@parameterize('placement_config', PLACEMENT_CONFIGS)
@parameterize('model_name', EXAMPLE_MODELS)
@parameterize('mixed_precision', [torch.half, torch.bfloat16])
@parameterize("placement_config", PLACEMENT_CONFIGS)
@parameterize("model_name", EXAMPLE_MODELS)
@parameterize("mixed_precision", [torch.half, torch.bfloat16])
def exam_tiny_example(placement_config, model_name: str, mixed_precision: torch.dtype):
set_seed(2008)
get_components_func = non_distributed_component_funcs.get_callable(model_name)
model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func()
torch_model = model_builder().cuda()
amp_config = dict(opt_level='O2', keep_batchnorm_fp32=False, loss_scale=2)
amp_config = dict(opt_level="O2", keep_batchnorm_fp32=False, loss_scale=2)
torch_optim = torch.optim.Adam(torch_model.parameters(), lr=1e-3)
torch_model, torch_optim = convert_to_apex_amp(torch_model, torch_optim, amp_config)
torch_model = DDP(torch_model, device_ids=[dist.get_rank()])
@@ -154,12 +136,14 @@ def exam_tiny_example(placement_config, model_name: str, mixed_precision: torch.
for torch_p, p in zip(torch_model.parameters(), model.parameters()):
p.data.copy_(torch_p.data)
model = GeminiDDP(model,
chunk_init_device=get_current_device(),
search_range_m=1,
pin_memory=True,
mixed_precision=mixed_precision,
**placement_config)
model = GeminiDDP(
model,
chunk_init_device=get_current_device(),
search_range_m=1,
pin_memory=True,
mixed_precision=mixed_precision,
**placement_config,
)
optimizer = HybridAdam(model.parameters(), lr=1e-3)
zero_optim = GeminiOptimizer(optimizer, model, initial_scale=2)
@@ -182,7 +166,7 @@ def exam_tiny_example(placement_config, model_name: str, mixed_precision: torch.
torch_loss = run_fwd_bwd(torch_model, input_ids, label, criterion, torch_optim)
loss = run_fwd_bwd(model, input_ids, label, criterion, zero_optim)
assert_close(torch_loss, loss, rtol=rtol, atol=atol) # atol should be 2e-5 for torch lower than 1.12
assert_close(torch_loss, loss, rtol=rtol, atol=atol) # atol should be 2e-5 for torch lower than 1.12
zero_optim.step()
torch_optim.step()
@@ -192,17 +176,17 @@ def exam_tiny_example(placement_config, model_name: str, mixed_precision: torch.
def run_dist(rank, world_size, port):
config = {}
colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
colossalai.launch(config=config, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
exam_model_step()
exam_tiny_example()
@pytest.mark.dist
@pytest.mark.parametrize('world_size', [1, 4])
@pytest.mark.parametrize("world_size", [1, 4])
@rerun_if_address_is_in_use()
def test_optim(world_size):
spawn(run_dist, world_size)
if __name__ == '__main__':
if __name__ == "__main__":
test_optim(1)

8
tests/test_zero/test_gemini/test_runtime_mem_tracer.py
View File

@@ -13,7 +13,7 @@ from tests.components_to_test.registry import non_distributed_component_funcs
@pytest.mark.skip("this is not used")
@clear_cache_before_run()
def test_runtime_mem_tracer():
test_models = ['gpt2', 'bert', 'simple_net', 'repeated_computed_layers', 'nested_model', 'albert']
test_models = ["gpt2", "bert", "simple_net", "repeated_computed_layers", "nested_model", "albert"]
for model_name in test_models:
get_components_func = non_distributed_component_funcs.get_callable(model_name)
@@ -35,7 +35,7 @@ def test_runtime_mem_tracer():
for p1, p2 in zip(model_bk.parameters(), model.parameters()):
torch.allclose(p1.to(torch.half), p2)
non_model_data_list = runtime_mem_tracer._memstats.non_model_data_list('cuda')
non_model_data_list = runtime_mem_tracer._memstats.non_model_data_list("cuda")
cuda_non_model_data_list = np.array(non_model_data_list) / 1024**2
print("cuda_non_model_data_list", len(cuda_non_model_data_list))
print(non_model_data_list)
@@ -46,9 +46,9 @@ def test_runtime_mem_tracer():
cnt2 = 0
for p in model.parameters():
cnt2 += 1
assert cnt2 == cnt1, f'visited param number {cnt1} vs real param number {cnt2}'
assert cnt2 == cnt1, f"visited param number {cnt1} vs real param number {cnt2}"
del model
if __name__ == '__main__':
if __name__ == "__main__":
test_runtime_mem_tracer()

36
tests/test_zero/test_gemini/test_search.py
View File

@@ -11,19 +11,17 @@ from tests.components_to_test.registry import non_distributed_component_funcs
def exam_search_chunk_size():
world_size = torch.distributed.get_world_size()
get_components_func = non_distributed_component_funcs.get_callable('gpt2')
get_components_func = non_distributed_component_funcs.get_callable("gpt2")
model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func()
# make sure torch_model and model has the same parameter values
model = model_builder()
config_dict, *_ = search_chunk_configuration(model,
search_range_m=1,
search_interval=16,
min_chunk_size_m=0,
filter_exlarge_params=True)
config_dict, *_ = search_chunk_configuration(
model, search_range_m=1, search_interval=16, min_chunk_size_m=0, filter_exlarge_params=True
)
for key in config_dict:
chunk_size = config_dict[key]['chunk_size']
chunk_size = config_dict[key]["chunk_size"]
if world_size == 1 or True:
assert chunk_size == 31616
else:
@@ -33,34 +31,36 @@ def exam_search_chunk_size():
def exam_chunk_manager():
world_size = torch.distributed.get_world_size()
get_components_func = non_distributed_component_funcs.get_callable('gpt2')
get_components_func = non_distributed_component_funcs.get_callable("gpt2")
model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func()
sharded_ddp_model = model_builder()
chunk_manager = init_chunk_manager(sharded_ddp_model,
get_current_device(),
hidden_dim=16,
search_range_m=1,
min_chunk_size_m=0,
filter_exlarge_params=True,
strict_ddp_flag=True)
chunk_manager = init_chunk_manager(
sharded_ddp_model,
get_current_device(),
hidden_dim=16,
search_range_m=1,
min_chunk_size_m=0,
filter_exlarge_params=True,
strict_ddp_flag=True,
)
config_dict = chunk_manager.dp_degree_chunk_size_dict
assert len(config_dict) == 1
assert config_dict[world_size] == 31616
def run_dist(rank, world_size, port):
colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
colossalai.launch(config={}, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
exam_search_chunk_size()
exam_chunk_manager()
@pytest.mark.dist
@pytest.mark.parametrize('world_size', [1, 4])
@pytest.mark.parametrize("world_size", [1, 4])
@rerun_if_address_is_in_use()
def test_search(world_size):
spawn(run_dist, world_size)
if __name__ == '__main__':
if __name__ == "__main__":
test_search(4)

51
tests/test_zero/test_gemini/test_zeroddp_state_dict.py
View File

@@ -10,21 +10,10 @@ from colossalai.zero.gemini.chunk import search_chunk_configuration
from tests.components_to_test.registry import non_distributed_component_funcs
PLACEMENT_CONFIGS = [
{
'placement_policy': 'static',
'shard_param_frac': 0.0
}, # zero2
{
'placement_policy': 'static',
'shard_param_frac': 1.0
}, # zero3
{
'placement_policy': 'static',
'shard_param_frac': 0.5
}, # zero3-half
{
'placement_policy': 'auto'
}
{"placement_policy": "static", "shard_param_frac": 0.0}, # zero2
{"placement_policy": "static", "shard_param_frac": 1.0}, # zero3
{"placement_policy": "static", "shard_param_frac": 0.5}, # zero3-half
{"placement_policy": "auto"},
]
@@ -35,9 +24,9 @@ def ignore_the_first_parameter(model: torch.nn.Module):
return
@parameterize('placement_config', PLACEMENT_CONFIGS)
@parameterize('keep_gathered', [True, False])
@parameterize('model_name', ['gpt2', 'bert'])
@parameterize("placement_config", PLACEMENT_CONFIGS)
@parameterize("keep_gathered", [True, False])
@parameterize("model_name", ["gpt2", "bert"])
def exam_state_dict(placement_config, keep_gathered, model_name: str):
set_seed(431)
get_components_func = non_distributed_component_funcs.get_callable(model_name)
@@ -51,8 +40,8 @@ def exam_state_dict(placement_config, keep_gathered, model_name: str):
world_size = torch.distributed.get_world_size()
config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100)
config_dict[world_size]['chunk_size'] = 5000
config_dict[world_size]['keep_gathered'] = keep_gathered
config_dict[world_size]["chunk_size"] = 5000
config_dict[world_size]["keep_gathered"] = keep_gathered
model = GeminiDDP(model, config_dict, **placement_config, pin_memory=True)
model.train()
@@ -65,9 +54,9 @@ def exam_state_dict(placement_config, keep_gathered, model_name: str):
assert_close(value, temp_zero_value, rtol=1e-3, atol=1e-5)
@parameterize('placement_config', PLACEMENT_CONFIGS)
@parameterize('keep_gathered', [True, False])
@parameterize('model_name', ['gpt2', 'bert'])
@parameterize("placement_config", PLACEMENT_CONFIGS)
@parameterize("keep_gathered", [True, False])
@parameterize("model_name", ["gpt2", "bert"])
def exam_load_state_dict(placement_config, keep_gathered, model_name: str):
set_seed(431)
get_components_func = non_distributed_component_funcs.get_callable(model_name)
@@ -76,12 +65,12 @@ def exam_load_state_dict(placement_config, keep_gathered, model_name: str):
model = model_builder()
set_seed(451)
torch_model = model_builder() # get a different model
torch_model = model_builder() # get a different model
world_size = torch.distributed.get_world_size()
config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100)
config_dict[world_size]['chunk_size'] = 5000
config_dict[world_size]['keep_gathered'] = keep_gathered
config_dict[world_size]["chunk_size"] = 5000
config_dict[world_size]["keep_gathered"] = keep_gathered
model = GeminiDDP(model, config_dict, **placement_config, pin_memory=True)
@@ -95,8 +84,8 @@ def exam_load_state_dict(placement_config, keep_gathered, model_name: str):
assert_close(value, temp_zero_value, rtol=1e-3, atol=1e-5)
@parameterize('placement_config', PLACEMENT_CONFIGS)
@parameterize('model_name', ['gpt2', 'bert'])
@parameterize("placement_config", PLACEMENT_CONFIGS)
@parameterize("model_name", ["gpt2", "bert"])
def exam_state_dict_shard(placement_config, model_name: str):
get_components_func = non_distributed_component_funcs.get_callable(model_name)
model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func()
@@ -122,18 +111,18 @@ def exam_state_dict_shard(placement_config, model_name: str):
def run_dist(rank, world_size, port):
config = {}
colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
colossalai.launch(config=config, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
exam_state_dict()
exam_load_state_dict()
exam_state_dict_shard()
@pytest.mark.dist
@pytest.mark.parametrize('world_size', [1, 4])
@pytest.mark.parametrize("world_size", [1, 4])
@rerun_if_address_is_in_use()
def test_zero_ddp(world_size):
spawn(run_dist, world_size)
if __name__ == '__main__':
if __name__ == "__main__":
test_zero_ddp(1)

44
tests/test_zero/test_gemini/test_zerooptim_state_dict.py
View File

@@ -11,32 +11,18 @@ from colossalai.zero.gemini.chunk import search_chunk_configuration
from tests.components_to_test.registry import non_distributed_component_funcs
PLACEMENT_CONFIGS = [
{
'placement_policy': 'static',
'shard_param_frac': 0.0,
'offload_optim_frac': 0.0
}, # zero2
{
'placement_policy': 'static',
'shard_param_frac': 0.0,
'offload_optim_frac': 1.0
}, # zero2-offload
{
'placement_policy': 'static',
'shard_param_frac': 0.0,
'offload_optim_frac': 0.5
}, # zero2-offload-half
{
'placement_policy': 'auto'
}
{"placement_policy": "static", "shard_param_frac": 0.0, "offload_optim_frac": 0.0}, # zero2
{"placement_policy": "static", "shard_param_frac": 0.0, "offload_optim_frac": 1.0}, # zero2-offload
{"placement_policy": "static", "shard_param_frac": 0.0, "offload_optim_frac": 0.5}, # zero2-offload-half
{"placement_policy": "auto"},
]
@parameterize('placement_config', PLACEMENT_CONFIGS)
@parameterize('keep_gathered', [True, False])
@parameterize("placement_config", PLACEMENT_CONFIGS)
@parameterize("keep_gathered", [True, False])
def exam_zero_optim_state_dict(placement_config, keep_gathered):
set_seed(431)
get_components_func = non_distributed_component_funcs.get_callable('gpt2')
get_components_func = non_distributed_component_funcs.get_callable("gpt2")
model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func()
model = model_builder()
@@ -45,13 +31,13 @@ def exam_zero_optim_state_dict(placement_config, keep_gathered):
world_size = torch.distributed.get_world_size()
config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100)
config_dict[world_size]['chunk_size'] = 5000
config_dict[world_size]['keep_gathered'] = keep_gathered
config_dict[world_size]["chunk_size"] = 5000
config_dict[world_size]["keep_gathered"] = keep_gathered
model = GeminiDDP(model, config_dict, **placement_config, pin_memory=True)
optimizer = HybridAdam(model.parameters())
optim = GeminiOptimizer(optimizer, model, initial_scale=32) # initialize the link between chunk16 and chunk32
optim = GeminiOptimizer(optimizer, model, initial_scale=32) # initialize the link between chunk16 and chunk32
set_seed(dist.get_rank() * 3 + 128)
model.train()
@@ -67,8 +53,8 @@ def exam_zero_optim_state_dict(placement_config, keep_gathered):
optim_state_dict = optim.state_dict()
optim.load_state_dict(optim_state_dict)
new_state = optim.state_dict()['state']
org_state = optim_state_dict['state']
new_state = optim.state_dict()["state"]
org_state = optim_state_dict["state"]
for k, v in org_state.items():
w = new_state[k]
@@ -82,16 +68,16 @@ def exam_zero_optim_state_dict(placement_config, keep_gathered):
def run_dist(rank, world_size, port):
config = {}
colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
colossalai.launch(config=config, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
exam_zero_optim_state_dict()
@pytest.mark.dist
@pytest.mark.parametrize('world_size', [1, 4])
@pytest.mark.parametrize("world_size", [1, 4])
@rerun_if_address_is_in_use()
def test_zero_optim(world_size):
spawn(run_dist, world_size)
if __name__ == '__main__':
if __name__ == "__main__":
test_zero_optim(1)