Skip to content

modify clip_grad with no to_global #10443

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

Open
hanwen-sun wants to merge 11 commits into
base: master
Choose a base branch
from
Open

modify clip_grad with no to_global #10443

Changes from 2 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
6560c00
modify clip_grad
hanwen-sun Mar 11, 2024
1e35800
auto format by CI
oneflow-ci-bot Mar 11, 2024
355dfe2
modify test_clip_grad
hanwen-sun Mar 27, 2024
dddf91e
test
hanwen-sun Apr 3, 2024
908f1c6
auto format by CI
oneflow-ci-bot Apr 3, 2024
766bf1e
Merge branch 'master' into modify_clip_grad_with_no_to_global
hanwen-sun Apr 7, 2024
ac47c0d
add comment
hanwen-sun Apr 8, 2024
a19966e
auto format by CI
oneflow-ci-bot Apr 8, 2024
381b985
recover threshold of 1e-4
Apr 16, 2024
e55a7be
auto format by CI
oneflow-ci-bot Apr 16, 2024
8eadd8e
Merge branch 'master' into modify_clip_grad_with_no_to_global
levi131 Apr 16, 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
18 changes: 5 additions & 13 deletions python/oneflow/nn/utils/clip_grad.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,31 +100,23 @@ def clip_grad_norm_(
param0_placement = parameters[0].placement
if norm_type == float("inf"):
norms = [
p.grad.detach()
.to_global(sbp=sbp_broadcast)
.abs()
.max()
.to_global(placement=param0_placement)
p.grad.detach().abs().max().to_global(placement=param0_placement)
for p in parameters
]
total_norm = norms[0] if len(norms) == 1 else flow.max(flow.stack(norms))
elif norm_type == float("-inf"):
norms = [
p.grad.detach()
.to_global(sbp=sbp_broadcast)
.abs()
.min()
.to_global(placement=param0_placement)
p.grad.detach().abs().min().to_global(placement=param0_placement)
for p in parameters
]
total_norm = norms[0] if len(norms) == 1 else flow.min(flow.stack(norms))
else:
total_norm = flow.linalg.vector_norm(
flow.stack(
[
flow.linalg.vector_norm(
p.grad.detach().to_global(sbp=sbp_broadcast), norm_type
).to_global(placement=param0_placement)
flow.linalg.vector_norm(p.grad.detach(), norm_type).to_global(
placement=param0_placement
)
for p in parameters
]
),
Expand Down