Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

set the default to use set_to_none for clearing gradients in BF16 optimizer. #5434

Merged
merged 8 commits into from
Apr 23, 2024
Merged

set the default to use set_to_none for clearing gradients in BF16 optimizer. #5434

Changes from 5 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
678a6e9
set clear_grad default to None, use foreach_zero
inkcherry Apr 18, 2024
6c568ac
fix typo
inkcherry Apr 18, 2024
d1ec544
Merge branch 'master' into fix_5175_
inkcherry Apr 18, 2024
6b77946
fix condition&add detach
inkcherry Apr 20, 2024
7c9c8a9
Merge branch 'fix_5175_' of https://github.com/inkcherry/DeepSpeed in…
inkcherry Apr 20, 2024
85f295d
Merge branch 'master' into fix_5175_
loadams Apr 22, 2024
32fc1ce
Merge branch 'master' into fix_5175_
loadams Apr 22, 2024
617e28a
Merge branch 'master' into fix_5175_
loadams Apr 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 13 additions & 4 deletions deepspeed/runtime/bf16_optimizer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -341,7 +341,7 @@ def _update_hp_grad(self, lp, group_idx, param_idx, clear_lp_grads):

# clear gradients
if clear_lp_grads:
lp.grad._zero()
lp.grad.zero_()
tjruwase marked this conversation as resolved.
Show resolved Hide resolved

@torch.no_grad()
def _update_hp_grads_func(self, clear_lp_grads=False):
Expand Down Expand Up @@ -441,11 +441,20 @@ def clear_hp_grads(self):
self.fp32_groups_has_gradients[i] = [False] * len(group)

def clear_lp_grads(self):

# using zero_() fixed memory address for graph replay
set_to_none = False if self.graph_harvesting else True
zero_grads_list = []
for group in self.bf16_groups:
for param in group:
if param.grad is not None:
# Using zero_() fixed memory address for graph replay
param.grad.zero_()
if set_to_none:
param.grad = None
elif param.grad is not None:
if param.grad.grad_fn is not None:
param.grad.detach_()
zero_grads_list.append(param.grad)
if not set_to_none and len(zero_grads_list) > 0:
torch._foreach_zero_(zero_grads_list)

def state_dict(self):
state_dict = {}
Expand Down