github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-06 19:40:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

[zero] update zero context init with the updated test utils (#327)

Browse Source
This commit is contained in:
Jiarui Fang
2022-03-08 14:45:01 +08:00
committed by Frank Lee
parent 6268446b81
commit 11bddb6e55
10 changed files with 96 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
tests/test_engine/test_engine.py
View File

@@ -16,10 +16,11 @@ CONFIG = dict(parallel=dict(pipeline=dict(size=1), tensor=dict(size=1, mode=None
def run_train():
for get_components_func in non_distributed_component_funcs:
model, train_dataloader, _, optimizer, criterion = get_components_func()
model_builder, train_dataloader, _, optimizer_builder, criterion = get_components_func()
model = model_builder(checkpoint=False)
engine, train_dataloader, *args = colossalai.initialize(model=model,
optimizer=optimizer,
optimizer=optimizer_builder(model),
criterion=criterion,
train_dataloader=train_dataloader)