From 060892162a592c1abf239ad11edd9a91b47188e0 Mon Sep 17 00:00:00 2001 From: Hongxin Liu Date: Tue, 30 Jul 2024 13:36:00 +0800 Subject: [PATCH] [zero] hotfix update master params (#5951) --- 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 f1db7c540..01382cd8e 100644 --- a/colossalai/zero/low_level/low_level_optim.py +++ b/colossalai/zero/low_level/low_level_optim.py @@ -807,8 +807,8 @@ class LowLevelZeroOptimizer(OptimizerWrapper): """ 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)