[zero] hotfix update master params (#5951)

hpcaitech · Jul 30, 2024 · 0608921 · 0608921
1 parent bcf0181
commit 0608921
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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)