From 26f22ed283ae4ac3a24b756809e5961efe6f9da8 Mon Sep 17 00:00:00 2001 From: Zhihao Lin <36994684+LZHgrla@users.noreply.github.com> Date: Wed, 8 Nov 2023 23:35:12 +0800 Subject: [PATCH] [Fix] Support exclude_frozen_parameters for DeepSpeedStrategy's resume (#1424) --- mmengine/_strategy/deepspeed.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/mmengine/_strategy/deepspeed.py b/mmengine/_strategy/deepspeed.py index a439a4a9..378616db 100644 --- a/mmengine/_strategy/deepspeed.py +++ b/mmengine/_strategy/deepspeed.py @@ -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'))