ColossalAI/tests/test_zero/test_legacy/test_param_op.py

83 lines
2.5 KiB
Python

import copy
import torch
from colossalai.testing import clear_cache_before_run
from colossalai.zero.legacy.gemini.paramhooks import BaseParamHookMgr
from tests.components_to_test.registry import non_distributed_component_funcs
def allclose(tensor_a: torch.Tensor, tensor_b: torch.Tensor, loose=False) -> bool:
if loose:
return torch.allclose(tensor_a, tensor_b, atol=1e-3, rtol=1e-3)
return torch.allclose(tensor_a, tensor_b)
def run_model(model, inputs, label, criterion, use_param_hook=False):
if use_param_hook:
class HooKWrapper:
def __init__(self) -> None:
self.hook_triggered_times = 0
def wrapper_func(self):
def hook(param, grad) -> torch.Tensor or None:
self.hook_triggered_times += 1
return grad
return hook
hookwrapper = HooKWrapper()
param_list = [p for p in model.parameters()]
hook_mgr = BaseParamHookMgr(param_list)
hook_mgr.register_backward_hooks(hookwrapper.wrapper_func())
model.zero_grad(set_to_none=True)
with torch.cuda.amp.autocast():
if criterion:
y = model(inputs)
loss = criterion(y, label)
else:
loss = model(inputs, label)
loss = loss.float()
loss.backward()
if use_param_hook:
hook_mgr.remove_hooks()
return hookwrapper.hook_triggered_times
@clear_cache_before_run()
def test_base_param_hook():
test_models = ['repeated_computed_layers', 'resnet18', 'hanging_param_model', 'inline_op_model']
# test_models = ['bert']
for model_name in test_models:
get_components_func = non_distributed_component_funcs.get_callable(model_name)
model_builder, train_dataloader, _, _, criterion = get_components_func()
torch.manual_seed(0)
model = model_builder(checkpoint=True).cuda()
model.train()
for i, (inputs, label) in enumerate(train_dataloader):
if i > 0:
break
model_copy = copy.deepcopy(model)
run_model(model, inputs.cuda(), label.cuda(), criterion, False)
ret2 = run_model(model_copy, inputs.cuda(), label.cuda(), criterion, True)
# Make sure param hook has only be fired once in case of parameter sharing
assert ret2 == len(list(model.parameters()))
for p, p_copy in zip(model.parameters(), model_copy.parameters()):
assert allclose(p.grad, p_copy.grad), f"{p.grad} vs {p_copy.grad}"
if __name__ == '__main__':
test_base_param_hook()