github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-11-02 06:54:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

[zero] adapt for no-leaf module in zero (#535)

Browse Source 
only process module's own parameters in Zero context

add zero hooks for all modules that contrain parameters

gather parameters only belonging to module itself
This commit is contained in:
HELSON
2022-03-28 17:42:18 +08:00
committed by GitHub
parent 705f56107c
commit a30e2b4c24
7 changed files with 70 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_zero_data_parallel/test_shard_model_v2.py
View File

@@ -24,7 +24,7 @@ from common import CONFIG, check_grads_padding, run_fwd_bwd
@parameterize("enable_autocast", [True])
@parameterize("shard_strategy_class", [TensorShardStrategy, BucketTensorShardStrategy])
def run_model_test(enable_autocast, shard_strategy_class):
test_models = ['repeated_computed_layers', 'resnet18', 'bert']
test_models = ['repeated_computed_layers', 'resnet18', 'bert', 'no_leaf_module']
shard_strategy = shard_strategy_class()
for model_name in test_models:
get_components_func = non_distributed_component_funcs.get_callable(model_name)