mirrors/mmengine
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmengine.git synced 2025-06-03 21:54:44 +08:00
Code Issues Projects Releases Wiki Activity

[Fix] Support exclude_frozen_parameters for DeepSpeedStrategy's resume (#1424)

Browse Source
This commit is contained in:
Zhihao Lin 2023-11-08 23:35:12 +08:00 committed by GitHub
parent 46784185cf
commit 26f22ed283
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 9 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
mmengine/_strategy/deepspeed.py
View File

@ -463,8 +463,15 @@ class DeepSpeedStrategy(BaseStrategy):
self.logger.info(f'Resume checkpoint from {filename}')
dirname, basename = osp.split(filename)
_, extra_ckpt = self.model.load_checkpoint(
dirname, tag=basename, load_optimizer_states=resume_optimizer)
if digit_version(deepspeed.__version__) >= digit_version('0.10.1'):
_, extra_ckpt = self.model.load_checkpoint(
dirname,
tag=basename,
load_optimizer_states=resume_optimizer,
load_module_strict=not self.exclude_frozen_parameters)
else:
_, extra_ckpt = self.model.load_checkpoint(
dirname, tag=basename, load_optimizer_states=resume_optimizer)
if resume_optimizer:
self.load_optim_state_dict(extra_ckpt.pop('optim_wrapper'))