[zero1&2] only append parameters with gradients (#2681)

pull/2686/head
HELSON 2 years ago committed by GitHub
parent f0aa191f51
commit df4f020ee3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -131,7 +131,10 @@ class LowLevelZeroOptimizer(ColossalaiOptimizer):
# partition these param groups for data parallel training # partition these param groups for data parallel training
# and add buffers to parameter store for future access # and add buffers to parameter store for future access
for group_id, param_group in enumerate(self.optim.param_groups): 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 # add the fp16 params to fp16_param_groups for bookkeeping
self._fp16_param_groups[group_id] = group_params self._fp16_param_groups[group_id] = group_params

Loading…
Cancel
Save