github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[Gemini] polish runtime tracer tests (#2077)

pull/2080/head
Jiarui Fang 2022-12-05 16:22:49 +08:00 committed by GitHub
parent 677e1e20d4
commit 4f21c9e8d9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
tests/test_gemini/test_runtime_mem_tracer.py
View File

@ -10,17 +10,6 @@ from tests.components_to_test import run_fwd_bwd
from tests.components_to_test.registry import non_distributed_component_funcs
def run_fwd_bwd(model, data, label, criterion, enable_autocast=False, dtype=torch.half):
with torch.cuda.amp.autocast(enabled=enable_autocast):
if criterion:
y = model(data)
loss = criterion(y, label)
else:
loss = model(data, label)
loss = loss.to(dtype)
model.backward(loss)
def test_runtime_mem_tracer():
test_models = ['gpt2', 'bert', 'simple_net', 'repeated_computed_layers', 'nested_model', 'albert']
@ -28,7 +17,7 @@ def test_runtime_mem_tracer():
get_components_func = non_distributed_component_funcs.get_callable(model_name)
model_builder, train_dataloader, _, _, criterion = get_components_func()
with ColoInitContext(device=torch.device('cpu')):
with ColoInitContext(device='cpu'):
model = model_builder(checkpoint=False)
model_bk = deepcopy(model)
@ -40,7 +29,7 @@ def test_runtime_mem_tracer():
data = data.cuda()
label = label.cuda()
run_fwd_bwd(runtime_mem_tracer, data, label, criterion, False)
run_fwd_bwd(runtime_mem_tracer, data, label, criterion, optimizer=runtime_mem_tracer)
for p1, p2 in zip(model_bk.parameters(), model.parameters()):
torch.allclose(p1.to(torch.half), p2)