github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2026-01-05 15:44:49 +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

12
tests/components_to_test/repeated_computed_layer.py
View File

@@ -36,9 +36,15 @@ class DummyDataLoader(DummyDataGenerator):
@non_distributed_component_funcs.register(name='repeated_computed_layers')
def get_training_components():
model = NetWithRepeatedlyComputedLayers(checkpoint=True)
def model_builder(checkpoint=True):
return NetWithRepeatedlyComputedLayers(checkpoint)
trainloader = DummyDataLoader()
testloader = DummyDataLoader()
optim = torch.optim.Adam(model.parameters(), lr=0.001)
def optim_builder(model):
return torch.optim.Adam(model.parameters(), lr=0.001)
criterion = torch.nn.CrossEntropyLoss()
return model, trainloader, testloader, optim, criterion
return model_builder, trainloader, testloader, optim_builder, criterion