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

Add main_grad #1140

Open
wants to merge 4 commits into
base: ngoyal_changes_for_pp_fp8
Choose a base branch
from
Open

Add main_grad #1140

wants to merge 4 commits into from

Conversation

jianyuh
Copy link
Member

@jianyuh jianyuh commented Oct 2, 2023

What does this PR do?

Fixes main_grad following up #1139 (comment)

Before submitting

  • Did you have fun?
    • Make sure you had fun coding 🙃
  • Did you read the contributor guideline?
  • Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
    • N/A
  • Did you make sure to update the docs?
    • N/A
  • Did you write any new necessary tests?
    • N/A
  • Did you update the changelog? (if needed)
    • N/A

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Oct 2, 2023
@jianyuh jianyuh marked this pull request as ready for review October 2, 2023 01:09
@jianyuh jianyuh mentioned this pull request Oct 2, 2023
Merged
10 tasks
@jianyuh jianyuh changed the base branch from ngoyal_changes_for_pp_fp8_fix_handle to ngoyal_changes_for_pp_fp8 October 2, 2023 03:03
@jianyuh jianyuh force-pushed the ngoyal_changes_for_pp_fp8_fix_handle_grad_main branch from 3f34441 to 239ed36 Compare October 2, 2023 03:11
vedanuj
vedanuj reviewed Oct 2, 2023
View reviewed changes
fairscale/nn/data_parallel/fully_sharded_data_parallel.py Outdated
@@ -1713,6 +1713,13 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:

# Switch to FP32 shard after backward.
self._use_fp32_param_shard([param])
if self.mixed_precision and self.fp32_reduce_scatter:
Copy link

Choose a reason for hiding this comment

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

Currently for fp8, we do not use mixed_precision, so we should remove this.

Only check

if self.fp32_reduce_scatter:

Copy link
Member Author

Choose a reason for hiding this comment

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

Addressed the comment.

@jianyuh jianyuh force-pushed the ngoyal_changes_for_pp_fp8_fix_handle_grad_main branch from 239ed36 to ad54660 Compare October 2, 2023 21:48
jianyuh
jianyuh commented Oct 4, 2023
View reviewed changes
fairscale/nn/data_parallel/fully_sharded_data_parallel.py

if self.fp32_reduce_scatter:
# Cast grad to FP32.
param.grad.data = param.grad.data.float()

orig_grad_data = param.grad.data
Copy link
Member Author

Choose a reason for hiding this comment

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

Move here to make orig_grad_data FP32. This was from #1139 (comment)

jspark1105
jspark1105 suggested changes Oct 4, 2023
View reviewed changes
fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -1721,23 +1728,31 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
# reductions in post_backward stream.
self._streams["post_backward"].wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(self._streams["post_backward"]):
orig_grad_data = param.grad.data

if self.fp32_reduce_scatter:
# Cast grad to FP32.
param.grad.data = param.grad.data.float()
Copy link

Choose a reason for hiding this comment

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

I don't feel this is right since param.grad will be None from L1722.

Overall, this PR creates main_grad for flat parameters while what we need to do is main_grad visible to TE modules. So probably we need to change FlatParameter as well?

Is this based on one of Naman's branches?

Copy link

Choose a reason for hiding this comment

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

I have a branch where i am adding param.main_grad to FlatParams to enable fuse wgrad accumulation. here is the PR : #1142

Copy link
Member Author

Choose a reason for hiding this comment

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

Thanks! Feel free to ignore the changes in this PR. Still learning about FlatParams etc.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@vedanuj vedanuj vedanuj left review comments

@jspark1105 jspark1105 jspark1105 requested changes

@jiecaoyu jiecaoyu Awaiting requested review from jiecaoyu

Assignees
No one assigned
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@jianyuh @jspark1105 @vedanuj @facebook-github-bot