From 8c4f9b5b8c5d5b258dc1bbafa3ca39f41d6363ce Mon Sep 17 00:00:00 2001 From: ver217 Date: Tue, 30 Jul 2024 12:33:07 +0800 Subject: [PATCH] [zero] hotfix update master params --- colossalai/zero/low_level/low_level_optim.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/colossalai/zero/low_level/low_level_optim.py b/colossalai/zero/low_level/low_level_optim.py index f1db7c540f51..01382cd8e865 100644 --- a/colossalai/zero/low_level/low_level_optim.py +++ b/colossalai/zero/low_level/low_level_optim.py @@ -807,8 +807,8 @@ def update_master_params(self, model: nn.Module) -> None: """ for p in model.parameters(): p_id = id(p) - pg = self.param_to_pg[p] if p_id in self.working_to_master_param: + pg = self.param_to_pg[p] master_param = self.working_to_master_param[p_id] padding_size = self.get_param_padding_size(p) working_param = p.data.view(-1)