Skip to content

Sequence-parallel amax reduction fix - #74

Merged
ptrendx merged 3 commits into
NVIDIA:mainfrom
ksivaman:amax_reduce_fix
Feb 16, 2023
Merged

Sequence-parallel amax reduction fix#74
ptrendx merged 3 commits into
NVIDIA:mainfrom
ksivaman:amax_reduce_fix

Conversation

@ksivaman

@ksivamanksivaman commented Feb 14, 2023

Copy link
Copy Markdown
Member

Signed-off-by: Kirthi Shankar Sivamani ksivamani@nvidia.com

@ksivaman
Fix no reduce_amax option for SP case
265fade
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman

ksivaman commented Feb 14, 2023

Copy link
Copy Markdown
MemberAuthor

/te-ci

@ksivaman

ksivaman commented Feb 14, 2023

Copy link
Copy Markdown
MemberAuthor

@mikolajblaz has verified convergence with this fix

@ptrendx

ptrendx commented Feb 14, 2023

Copy link
Copy Markdown
Member

@ksivaman Could you give some background on the issue being fixed here?

@ksivaman

ksivaman commented Feb 14, 2023

Copy link
Copy Markdown
MemberAuthor

We have a reduce_amax flag in the FP8 recipe. If set to False, all of global amax reduction and the steps to set this up are skipped entirely. However, this is an issue because reduction of amaxes across the tensor parallel group is compulsory when sequence parallel is used because we gather the tensor in FP8 after doing the cast. This fix decouples the reduce_amax flag from the reduction across TP group. @ptrendx

@ksivaman

ksivaman commented Feb 15, 2023

Copy link
Copy Markdown
MemberAuthor

/te-ci

ptrendx
ptrendx reviewed Feb 16, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/module.py Outdated
# Previous iteration was grad_enabled
if self.fp8_meta.get("update_amax_and_scale_fwd", False):
if self.fp8_meta["recipe"].reduce_amax:
if self.fp8_meta["recipe"].reduce_amax or self.sequence_parallel:

@ptrendxptrendxFeb 16, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think there should be some sort of one-time warning if those two settings are set differently, since you are basically overriding the user here -> they should know that the amax reduction actually still happens.

@ksivaman
add warning about overriding reduce_amax
fa0f71c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman

ksivaman commented Feb 16, 2023

Copy link
Copy Markdown
MemberAuthor

/te-ci

ptrendx
ptrendx reviewed Feb 16, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/module.py
ptrendx
ptrendx approved these changes Feb 16, 2023
View reviewed changes

@ptrendxptrendx left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@ptrendx
ptrendx merged commit f06e2d8 into NVIDIA:mainFeb 16, 2023
@ksivaman
ksivaman deleted the amax_reduce_fix branch July 19, 2023 01:40
@XLzedXLzed mentioned this pull request Jul 3, 2024
Closed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ptrendxptrendxptrendx approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@ksivaman@ptrendx