From e2ff12b843afdf67b9a7132e9fbc9f7441d793e3 Mon Sep 17 00:00:00 2001 From: LZHgrla Date: Tue, 7 Nov 2023 18:38:13 +0800 Subject: [PATCH] fix --- 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 a439a4a952..378616db3d 100644 --- a/mmengine/_strategy/deepspeed.py +++ b/mmengine/_strategy/deepspeed.py @@ -463,8 +463,15 @@ def resume( 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'))