[Distributed] add pack-check method for float8_e4m3fn #135961
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/135961
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (2 Unrelated Failures)As of commit 183947e with merge base 0216936 ( FLAKY - The following jobs failed but were likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
oncall: distributed
| // x[i] = 0x7f - (x[i] + 0x01) | ||
| // (3) Only leave the sign bit by masking with 0x80. | ||
| // (4) If any x[i] is negative, then the whole x != 0. | ||
| #define CHECK_8_FLOAT8_E4M3FN(x) \ |
There was a problem hiding this comment.
Should we do this for the other float8 dtypes?
There was a problem hiding this comment.
Yeah, I will add support for e5m2 in another PR :)
There was a problem hiding this comment.
You should consider doing it for the float16 types as well. You'll will be reducing the number of SASS instructions considerably.
| // x[i] = 0x7f - (x[i] + 0x01) | ||
| // (3) Only leave the sign bit by masking with 0x80. | ||
| // (4) If any x[i] is negative, then the whole x != 0. | ||
| #define CHECK_8_FLOAT8_E4M3FN(x) \ |
There was a problem hiding this comment.
I suggest making this a function/template specialization (e.g. has_nan_vectorized) for type safety. Also it would help organizing similar optimizations for other scalar types.
There was a problem hiding this comment.
Yep, changed it to inline function now.
| // (2) If the result is 0x7f (is nan), then the following would be a negative number: | ||
| // x[i] = 0x7f - (x[i] + 0x01) |
There was a problem hiding this comment.
I feel doing x[i] = x[i] + 0x01 and checking the most significant bit is a bit easier to understand.
There was a problem hiding this comment.
Yeah, changed now :)
We check 8 x FP8 simultaneously, at size of 8 bytes. cc XilunWu H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Add support for Float8_e5m2, following similar algorithm used for Float8_e4m3fn (i.e. overflow check). Made `HasNanFP8x8` a template so that it is extendable based on dtype. Pull Request resolved: #136115 Approved by: https://github.com/Skylion007 ghstack dependencies: #135891, #135961
We check 8 x FP8 simultaneously, at size of 8 bytes. Pull Request resolved: pytorch#135961 Approved by: https://github.com/yifuwang, https://github.com/Skylion007 ghstack dependencies: pytorch#135891
Add support for Float8_e5m2, following similar algorithm used for Float8_e4m3fn (i.e. overflow check). Made `HasNanFP8x8` a template so that it is extendable based on dtype. Pull Request resolved: pytorch#136115 Approved by: https://github.com/Skylion007 ghstack dependencies: pytorch#135891, pytorch#135961
5 participants
Stack from ghstack (oldest at bottom):
We check 8 x FP8 simultaneously, at size of 8 bytes.
cc @XilunWu @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o