Backward register #423
Backward register #423
Conversation
StrongSpoon
force-pushed
the
bwd
branch
2 times, most recently
from
9f79739 to
01bee17
Compare
| 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.
There was a problem hiding this comment.
the definition of native_batch_norm_backward in aten lib is like below and we keep it the same:
native_batch_norm_backward(Tensor grad_out, Tensor input, Tensor? weight, Tensor? running_mean, Tensor? running_var, Tensor? save_mean, Tensor? save_invstd, bool train, float eps, bool[3] output_mask) -> (Tensor, Tensor, Tensor)
| 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.
There was a problem hiding this comment.
it's a bit complex. fix it later QAQ
| 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.
src/flag_gems/ops/dropout.py
Outdated
|
|
||
| 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.
| indices = indices.contiguous() | ||
| weight = weight.contiguous() |
src/flag_gems/ops/groupnorm.py
Outdated
| mean = mean.contiguous() | ||
| rstd = rstd.contiguous() | ||
| weight = None if weight is None else weight.contiguous() | ||
| group_size = C // group |
src/flag_gems/ops/groupnorm.py
Outdated
| 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), |
src/flag_gems/ops/dropout.py
Outdated
|
|
||
| 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.
PR Category
Operator
Type of Change
New Feature
Description
register backward functions as aten interfaces
implement threshold operator incidentally
Issue
Progress
Performance