github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-01 09:07:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Tensor] simplify named param (#928)

Browse Source 
* simplify ColoModulize

* simplify ColoModulize

* polish

* polish
This commit is contained in:
Ziyue Jiang
2022-05-11 10:54:19 +08:00
committed by GitHub
parent 32a45cd7ef
commit dfc88b85ea
2 changed files with 14 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tests/test_tensor/test_model.py
View File

@@ -1,3 +1,4 @@
from colossalai.tensor.colo_parameter import ColoParameter
from tests.components_to_test.registry import non_distributed_component_funcs
import colossalai
@@ -371,7 +372,13 @@ def _run_pretrain_load():
dict_col = {}
for name, param in model_pretrained.named_parameters():
dict_pretrained[name] = param
for name, param in model.named_parameters():
c1 = 0
c2 = 0
for name, param in model.colo_named_parameters():
if isinstance(param, ColoParameter):
c1 = c1 + 1
else:
c2 = c2 + 1
dict_col[name] = param
for name, param in dict_pretrained.items():
@@ -416,4 +423,5 @@ if __name__ == '__main__':
# test_model_parameters()
# test_colo_optimizer()
# test_model()
_test_pretrain_load(4)
# _test_pretrain_load(4)
_run_pretrain_load()