github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-04 10:34:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

[plugin] a workaround for zero plugins' optimizer checkpoint (#3780)

Browse Source 
* [test] refactor torch ddp checkpoint test

* [plugin] update low level zero optim checkpoint

* [plugin] update gemini optim checkpoint
This commit is contained in:
Hongxin Liu
2023-05-19 19:42:31 +08:00
committed by GitHub
parent 60e6a154bc
commit 3c07a2846e
6 changed files with 128 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
colossalai/booster/plugin/gemini_plugin.py
View File

@@ -52,8 +52,16 @@ class GeminiCheckpointIO(GeneralCheckpointIO):
Save optimizer to checkpoint but only on master process.
"""
# TODO(ver217): optimizer state dict is sharded
warnings.warn('GeminiPlugin does not support save full optimizer checkpoint now. Save it on every process.')
checkpoint = f'{checkpoint}.rank{self.coordinator.rank}'
super().save_unsharded_optimizer(optimizer, checkpoint, gather_dtensor)
def load_optimizer(self, optimizer: Optimizer, checkpoint: str):
warnings.warn(
'GeminiPlugin can only load optimizer checkpoint saved by itself with the same number of processes.')
checkpoint = f'{checkpoint}.rank{self.coordinator.rank}'
super().load_optimizer(optimizer, checkpoint)
def save_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str):
"""
Save model to checkpoint but only on master process.