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

[Feature] Auto scaling factor tuning for FP8 collective communication #140

Open
wants to merge 15 commits into
base: main
Choose a base branch
from
+111 −22
Open

[Feature] Auto scaling factor tuning for FP8 collective communication #140

Changes from 1 commit
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
6352d6f
wgrad auto scaling in optimizer
wkcn Dec 7, 2023
57d4a42
pre_scale support
wkcn Dec 7, 2023
69f73ce
fp8e4m3 for wgrad
wkcn Dec 7, 2023
55735b7
refine code
wkcn Dec 7, 2023
bcd0f74
comment
wkcn Dec 7, 2023
d248505
pre_scale support in add_to_fp8
wkcn Dec 7, 2023
4a7fd1f
pre_scale as Tensor
wkcn Dec 8, 2023
c850d63
pre_scale *fp32
wkcn Dec 8, 2023
617ee0b
pre_scale ut
wkcn Dec 8, 2023
199eb3a
test meta
wkcn Dec 8, 2023
1be092c
auto scaling freq
wkcn Dec 10, 2023
c08bc0d
Merge branch 'wgrad_auto_scaling' of github.com:wkcn/ms-amp into wgra…
wkcn Dec 10, 2023
6e58e1e
unittest
wkcn Dec 11, 2023
df68631
bug fixed when stacking pre_scales
wkcn Dec 11, 2023
664025d
Merge branch 'main' into wgrad_auto_scaling
wkcn Dec 14, 2023
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
Prev Previous commit
Next Next commit
fp8e4m3 for wgrad
@wkcn
wkcn committed Dec 7, 2023

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
GPG key ID: B5690EEEBB952194
Verified
Learn about vigilant mode
commit 69f73ce35b6a1d48e441643d2cefab2388b54423
13 changes: 6 additions & 7 deletions msamp/megatron/optimizer/distrib_optimizer.py
View file
Original file line number Diff line number Diff line change
@@ -547,16 +547,15 @@ def reduce_model_grads(self, args, timers): # noqa: C901
for p in model_group:
g = p.main_grad
if g is not None and not torch.is_tensor(g):
if g.qtype != WGRAD_QTYPE:
raise TypeError('g.qtype != WGRAD_QTYPE: {} != {}'.format(g.qtype, WGRAD_QTYPE))
if g.qtype != Dtypes.kfloat8_e4m3:
raise TypeError('g.qtype != Dtypes.kfloat8_e4m3: {}'.format(g.qtype))
# stat overflow ratio
num_infs = torch.count_nonzero((g.value & 0x7f) == 126)
overflow_ratio = num_infs / g.numel()
if args.wgrad_auto_scaling_ratio is not None:
if overflow_ratio > args.wgrad_auto_scaling_ratio:
g.meta.pre_scale /= 2.0
else:
g.meta.pre_scale *= 2.0**(1.0 / args.wgrad_auto_scaling_window)
if overflow_ratio > args.wgrad_auto_scaling_ratio:
g.meta.pre_scale /= 2.0
else:
g.meta.pre_scale *= 2.0**(1.0 / args.wgrad_auto_scaling_window)

# synchonize pre_scale
for model_id, model in enumerate(self.models):