Added fuse_rms_norm lowering #4017
Conversation
github-actions
bot
added
component: lowering
cehongwang
force-pushed
the
fused_rms_norm_lowering
branch
from
2bd6654 to
55010b6
Compare
|
|
||
| gm.graph.erase_node(node) | ||
|
|
||
| return x_normalized |
There was a problem hiding this comment.
In the replacement with the normalized, in line 80, should we also check for the second return value of the op? The rstd as we see in the test cases. Though yeah I think application wise, its mainly used for the gradient, but to be consistent with the op signature
There was a problem hiding this comment.
yeah I think it should have two Tensor outputs per https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/native_functions.yaml
zewenli98
left a comment
zewenli98
left a comment
There was a problem hiding this comment.
The PR looks good overall. But what I'm thinking is that, comparing with the converter-style implementation, the ops are almost same except impl.slice.expand and a few tensor casting. Do you know 1) which op/layer causes the perf discrepancy (60% as you mentioned in the last meeting)? 2) do these two approaches build the same size engines? If we can identify the issue, we probably can optimize other ops as well.
| # If the getitem is extracting the first element (the output tensor) | ||
| if not x_normalized.meta: | ||
| x_normalized.meta = copy.copy(node.meta) | ||
| user.replace_all_uses_with(x_normalized) |
There was a problem hiding this comment.
As Naren previously mentioned, can you add a log here when each node is changed?
|
|
||
| gm.graph.erase_node(node) | ||
|
|
||
| return x_normalized |
There was a problem hiding this comment.
yeah I think it should have two Tensor outputs per https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/native_functions.yaml
cehongwang
force-pushed
the
fused_rms_norm_lowering
branch
from
595aa5e to
6d99cd6
Compare
Co-authored-by: cehongwang <wangcehong@gmail.com>
cehongwang
force-pushed
the
fused_rms_norm_lowering
branch
from
121d636 to
a3ed926
Compare
Description
Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.
Fixes # (issue)
Type of change
Please delete options that are not relevant and/or add your own.
Checklist: