This repository has been archived by the owner on Aug 7, 2024. It is now read-only.
make the backward of differentiable float8 casts pass gradient as is #255
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Summary: Behavior before: * high precision to float8 in fw, float8 to high precision in bw * float8 to high precision in fw, high precision to float8 in bw if grad is a Float8Tensor, pass gradient unchanged otherwise Behavior after: * high precision to float8 in fw, pass gradient unchanged in bw * float8 to high precision in fw, pass gradient unchanged in bw Motivation for the new state: 1. we want gradients to be in high precision unless specified otherwise by the float8 recipe, and the logic to specify grad casting to float8 before the matmul is better implemented elsewhere 2. there is actually no logic change in this diff as the backward casts were not getting hit from existing code, this diff just makes the intended behavior clearer Test Plan: ``` ./test/test_everything.sh ``` Reviewers: Subscribers: Tasks: Tags:
facebook-github-bot
added
the
CLA Signed
|
I noticed this logic was funky when copy-pastaing some of this code for MX prototyping, fixing here to clarify intent. |
wanchaol
reviewed
May 3, 2024
| else: | ||
| return grad | ||
|
|
||
| if isinstance(x, DTensor): |
There was a problem hiding this comment.
if I understand the old state correctly:
gradis never aFloat8Tensorin practice, because we always have the matmul output high precision- things already work without nested subclasses, which is why removing this function is fine
But, would be great to clarify ^
There was a problem hiding this comment.
Oh I checked the current workflow, we call:
- cast_to_float8_e4m3fn
- cast_to_float8_e5m2_bw
and it seems none of these calls would call into this from_fp8_no_autograd
so this should be fine
There was a problem hiding this comment.
actually it looks like I'm wrong, the cast_to_float8_e4m3fn's backward would call this, but maybe it's a no-op even as of current state, as the output of the torch.scaled_mm in backward is fp32?
wanchaol
approved these changes
May 3, 2024
|
@vkuzo has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
drisspg
approved these changes
May 6, 2024
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary:
Behavior before:
Behavior after:
Motivation for the new state:
Test Plan:
Reviewers:
Subscribers:
Tasks:
Tags: