Type Conversions in Layer Norm Fusion #17281
Closed
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
reedwm
requested changes
Sep 27, 2024
Comment on lines
959
to
963
| // shapes of scale and bias must match. If a conversion to the type of the | ||
| // input is the only user of the output, set the output to the conversion. | ||
| // Similarly, if one of the users of the scale/bias is a conversion to the | ||
| // type of the bias/scale, set the scale/bias to the conversion. | ||
| if (instr->user_count() == 1 && |
There was a problem hiding this comment.
I am nervous about doing an incorrect rewrite based on these conversions you allow. What if the user, e.g., converts the input of layer norm from f32 to s2 and does all the layer norm logic in s2, before casting back to f32? It would be incorrect i think to rewrite this to a cudnn layer norm of full precision. Can we check at least all the types are floats of at least precision bf16?
Comment on lines
969
to
975
| for (HloInstruction* scale_user : scale->users()) { | ||
| if (scale_user->opcode() == HloOpcode::kConvert && | ||
| ShapeUtil::SameElementType(scale_user->shape(), bias->shape())) { | ||
| scale = scale_user; | ||
| break; | ||
| } | ||
| } |
There was a problem hiding this comment.
Why check if there is a convert user? This user might not even be part of the layer norm.
|
This looks good to me. As we discussed offline, you want to do some more testing on models first, so I'll wait until you do that before approving, otherwise this may get prematurely merged. |
reedwm
approved these changes
Oct 8, 2024
|
Approving as you confirmed the testing did not find issues. |
copybara-service bot
pushed a commit
that referenced
this pull request
Oct 8, 2024
Imported from GitHub PR #17281 Enables the fusion of layer norm graphs with type conversions of input, scale and bias. Copybara import of the project: -- 2880fde by Philipp Hack <[email protected]>: Layer norm fusion with type conversions of input, scale and bias. -- 898f002 by Philipp Hack <[email protected]>: Layer norm fusion with type conversions of input, scale and bias. Merging this change closes #17281 FUTURE_COPYBARA_INTEGRATE_REVIEW=#17281 from philipphack:u_layer_convert_xla 898f002 PiperOrigin-RevId: 683664734
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Oct 8, 2024
Imported from GitHub PR openxla/xla#17281 Enables the fusion of layer norm graphs with type conversions of input, scale and bias. Copybara import of the project: -- 2880fde5f71ad1aba23651ff866fd00becc706e5 by Philipp Hack <[email protected]>: Layer norm fusion with type conversions of input, scale and bias. -- 898f002f419909f367e6f9eedc72c61f4c73d201 by Philipp Hack <[email protected]>: Layer norm fusion with type conversions of input, scale and bias. Merging this change closes #17281 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#17281 from philipphack:u_layer_convert_xla 898f002f419909f367e6f9eedc72c61f4c73d201 PiperOrigin-RevId: 683664734
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Oct 8, 2024
Imported from GitHub PR openxla/xla#17281 Enables the fusion of layer norm graphs with type conversions of input, scale and bias. Copybara import of the project: -- 2880fde5f71ad1aba23651ff866fd00becc706e5 by Philipp Hack <[email protected]>: Layer norm fusion with type conversions of input, scale and bias. -- 898f002f419909f367e6f9eedc72c61f4c73d201 by Philipp Hack <[email protected]>: Layer norm fusion with type conversions of input, scale and bias. Merging this change closes #17281 PiperOrigin-RevId: 683682296
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Oct 8, 2024
FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#17281 from philipphack:u_layer_convert_xla 898f002f419909f367e6f9eedc72c61f4c73d201 PiperOrigin-RevId: 671073597
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.
Enables the fusion of layer norm graphs with type conversions of input, scale and bias.