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

[booster] add low level zero plugin (#3594)

Browse Source 
* [booster] add low level zero plugin

* [booster] fix gemini plugin test

* [booster] fix precision

* [booster] add low level zero plugin test

* [test] fix booster plugin test oom

* [test] fix booster plugin test oom

* [test] fix googlenet and inception output trans

* [test] fix diffuser clip vision model

* [test] fix torchaudio_wav2vec2_base

* [test] fix low level zero plugin test
This commit is contained in:
Hongxin Liu
2023-04-26 14:37:25 +08:00
committed by GitHub
parent b9a8dff7e5
commit 4b3240cb59
9 changed files with 476 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
colossalai/zero/low_level/low_level_optim.py
View File

@@ -55,6 +55,7 @@ class LowLevelZeroOptimizer(ColossalaiOptimizer):
# 2. contiguous gradients
# 3. cpu offload
# 4. support when some parameters requires_grad = False
# 5. support layer drop
super(LowLevelZeroOptimizer, self).__init__(optim=optimizer)
self._dtype = self.optim.param_groups[0]['params'][0].dtype
self._logger = get_dist_logger()