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

[shardformer] zero1+pp and the corresponding tests (#4517)

Browse Source 
* pause

* finish pp+zero1

* Update test_shard_vit.py
This commit is contained in:
Jianghai
2023-08-28 10:51:16 +08:00
committed by GitHub
parent 44eab2b27f
commit 376533a564
10 changed files with 109 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
tests/test_shardformer/test_model/test_shard_t5.py
View File

@@ -118,6 +118,15 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
'zero_stage': 2,
'precision': 'fp16',
'initial_scale': 1
}, {
'tp_size': 1,
'pp_size': 2,
'num_microbatches': 2,
'enable_all_optimization': True,
'use_lazy_init': True,
'zero_stage': 1,
'precision': 'fp16',
'initial_scale': 1
}])
@clear_cache_before_run()
def run_t5_test(test_config):