From df4f020ee32c5b857d2a9806d5ec40d0b2064021 Mon Sep 17 00:00:00 2001 From: HELSON Date: Mon, 13 Feb 2023 18:00:16 +0800 Subject: [PATCH] [zero1&2] only append parameters with gradients (#2681) --- colossalai/zero/sharded_optim/low_level_optim.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/colossalai/zero/sharded_optim/low_level_optim.py b/colossalai/zero/sharded_optim/low_level_optim.py index d174fc6ac..89f5f9fad 100644 --- a/colossalai/zero/sharded_optim/low_level_optim.py +++ b/colossalai/zero/sharded_optim/low_level_optim.py @@ -131,7 +131,10 @@ class LowLevelZeroOptimizer(ColossalaiOptimizer): # partition these param groups for data parallel training # and add buffers to parameter store for future access for group_id, param_group in enumerate(self.optim.param_groups): - group_params = param_group['params'] + group_params = list() + for param in param_group['params']: + if param.requires_grad: + group_params.append(param) # add the fp16 params to fp16_param_groups for bookkeeping self._fp16_param_groups[group_id] = group_params