ColossalAI/tests/test_booster/test_plugin
Baizhou Zhang 21ba89cab6
[gemini] support gradient accumulation (#4869)
* add test

* fix no_sync bug in low level zero plugin

* fix test

* add argument for grad accum

* add grad accum in backward hook for gemini

* finish implementation, rewrite tests

* fix test

* skip stuck model in low level zero test

* update doc

* optimize communication & fix gradient checkpoint

* modify doc

* cleaning codes

* update cpu adam fp16 case
2023-10-17 14:07:21 +08:00
..
test_3d_plugin.py [misc] update pre-commit and run all files (#4752) 2023-09-19 14:20:26 +08:00
test_dp_plugin_base.py [misc] update pre-commit and run all files (#4752) 2023-09-19 14:20:26 +08:00
test_gemini_plugin.py [misc] update pre-commit and run all files (#4752) 2023-09-19 14:20:26 +08:00
test_low_level_zero_plugin.py [gemini] support gradient accumulation (#4869) 2023-10-17 14:07:21 +08:00
test_torch_ddp_plugin.py [misc] update pre-commit and run all files (#4752) 2023-09-19 14:20:26 +08:00
test_torch_fsdp_plugin.py [misc] update pre-commit and run all files (#4752) 2023-09-19 14:20:26 +08:00