Backward register #423
Open
Backward register #423
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
StrongSpoon
force-pushed
the
bwd
branch
2 times, most recently
from
9f79739 to
01bee17
Compare
tongxin
reviewed
Feb 19, 2025
| save_invstd=None, | ||
| train=False, | ||
| eps=1e-05, | ||
| output_mask=None, |
There was a problem hiding this comment.
The last argument should be grad_input_mask.
| affine: tl.constexpr, | ||
| input_grad_mask: tl.constexpr, | ||
| weight_grad_mask: tl.constexpr, | ||
| bias_grad_mask: tl.constexpr, |
There was a problem hiding this comment.
The backward kernel may need is_train arg also, to distinguish between train and non-train cases.
There was a problem hiding this comment.
We can leave it for future work tho.
| running_var=None, | ||
| save_mean=None, | ||
| save_invstd=None, | ||
| train=False, |
There was a problem hiding this comment.
kernel should be able to handle train=True case.
|
|
||
| def native_dropout(x, p=0.5, train=True): | ||
| return NativeDropout.apply(x, p, train) | ||
| def dropout(input, p, train): |
| logging.debug("GEMS NATIVE DROPOUT FORWARD") | ||
| assert p > 0.0 and p < 1.0, "p must be in (0, 1)" | ||
| device = input.device | ||
| input = input.contiguous() |
There was a problem hiding this comment.
Add a note that we'll remove contiguous enforcement in the future.
Comment on lines
+119
to
+120
| indices = indices.contiguous() | ||
| weight = weight.contiguous() |
| mean = mean.contiguous() | ||
| rstd = rstd.contiguous() | ||
| weight = None if weight is None else weight.contiguous() | ||
| group_size = C // group |
| BLOCK_GROUP_SIZE=triton.next_power_of_2(C // num_groups), | ||
| BLOCK_HW_SIZE=triton.next_power_of_2(HW), | ||
| HxW, | ||
| BLOCK_GROUP_SIZE=triton.next_power_of_2(C // group), |
tongxin
reviewed
Feb 19, 2025
|
|
||
| def native_dropout(x, p=0.5, train=True): | ||
| return NativeDropout.apply(x, p, train) | ||
| def dropout(input, p, train): |
There was a problem hiding this comment.
I realized we didn't handle we train=False correctly in the previous version. Let's fix that.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
PR Category
Operator
Type of Change
New Feature
Description
register backward functions as aten interfaces
implement threshold operator incidentally
Issue
Progress
Performance