Use torch.repeat instead of expand on key & value in Triton MQA to prevent NaNs with certain h_dims #442
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.
For h_dim=8, we see NaNs due to a
.expand
of the key value tensors, which can be resolved with.repeat
. While h_dim=8 is an edge case, we are not sure if there are other cases of h_dims for which this might be problematic, or if there might be silent failures.Note that this does come at a performance hit, see MFU for 7b, so it may be desirable to revert this change in the future.
(blue curve:
![image](https://private-user-images.githubusercontent.com/4717972/251878086-385f222f-5ac4-430d-b10c-67c084b97955.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MjMwNzU4MjUsIm5iZiI6MTcyMzA3NTUyNSwicGF0aCI6Ii80NzE3OTcyLzI1MTg3ODA4Ni0zODVmMjIyZi01YWM0LTQzMGQtYjEwYy02N2MwODRiOTc5NTUucG5nP1gtQW16LUFsZ29yaXRobT1BV1M0LUhNQUMtU0hBMjU2JlgtQW16LUNyZWRlbnRpYWw9QUtJQVZDT0RZTFNBNTNQUUs0WkElMkYyMDI0MDgwOCUyRnVzLWVhc3QtMSUyRnMzJTJGYXdzNF9yZXF1ZXN0JlgtQW16LURhdGU9MjAyNDA4MDhUMDAwNTI1WiZYLUFtei1FeHBpcmVzPTMwMCZYLUFtei1TaWduYXR1cmU9ZTc4YTE4MmJlOTkyMTM4NzM4YWQ4OGFhZTQ2MjE1YTRmMWRjMWRjNGJkMWNlOTVhNDg1NTg0YTgzMmZlYWNkOSZYLUFtei1TaWduZWRIZWFkZXJzPWhvc3QmYWN0b3JfaWQ9MCZrZXlfaWQ9MCZyZXBvX2lkPTAifQ.K7g8O3ZvGq2bsQy2PxDq_UWT-u0Z3xDys_OEC8HeDjc)
.expand()
green curve.repeat()
red curve.expand().clone()
)