github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2026-01-29 21:49:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Gemini] add albert in test models. (#2075)

Browse Source
This commit is contained in:
Jiarui Fang
2022-12-05 14:09:34 +08:00
committed by GitHub
parent 616ed91ecd
commit 40b7d55bf3
5 changed files with 64 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_gemini/update/test_optim.py
View File

@@ -27,7 +27,7 @@ from tests.test_tensor.common_utils import debug_print, set_seed
# this model is large enough to slice to chunks
TEST_MODELS = ['gpt2']
# these models are too small, all parameters in these models are compacted into one chunk
EXAMPLE_MODELS = ['hanging_param_model', 'bert', 'simple_net', 'nested_model', 'repeated_computed_layers']
EXAMPLE_MODELS = ['albert', 'hanging_param_model', 'bert', 'simple_net', 'nested_model', 'repeated_computed_layers']
def check_param(model: ZeroDDP, torch_model: torch.nn.Module):