github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

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

pull/2686/head
HELSON 2023-02-13 18:00:16 +08:00 committed by GitHub
parent f0aa191f51
commit df4f020ee3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
colossalai/zero/sharded_optim/low_level_optim.py
View File

@ -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