pull/306/head
mwiacx 2023-09-27 09:36:50 +08:00
parent b2134f4dd4
commit 7d52e223a8
1 changed files with 2 additions and 0 deletions

View File

@ -817,6 +817,8 @@ class HybridZeroOptimizer(BaseOptimizer):
self.optim.step()
# release the fp32 grad
for group_id in range(self.num_param_groups):
if not self.param_group_has_params[group_id]:
continue
release_param_grad(self._fp32_orig_param_groups_of_current_rank[group_id])
if self._enable_memory_balance and self._memory_balance_role == 1: