Support MX4 E3M0 format and add stochastic rounding #477
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/477
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
| @@ -21,6 +21,8 @@ | |||
|
|
|||
| from torchao.prototype.mx_formats.constants import ( | |||
| DTYPE_FP4, | |||
| DTYPE_FP4_E2M1, | |||
There was a problem hiding this comment.
for all the changes in this file, can you rebase past #363 ?
and in the code after that PR, perhaps we can add stochastic rounding option to _f32_to_fpx_unpacked?
for rounding mode, how about something like this instead of a boolean?
class RoundingMode(enum.Enum):
TIE_TO_EVEN = auto() # default
STOCHASTIC = auto() # added in this PR
def foo(..., rounding_mode=RoundingMode.TIE_TO_EVEN, ...): ...
| @pytest.mark.parametrize("device", ["cuda", "cpu"]) | ||
| @pytest.mark.parametrize("sign", [1, -1]) | ||
| @pytest.mark.parametrize("use_stochastic_rounding", [False, True]) | ||
| def test_overflow_cast(hp_dtype, device, sign, use_stochastic_rounding): |
There was a problem hiding this comment.
can we add these tests to test/prototype/mx_formats/test_custom_cast.py to keep the testing of MX numerics in one place?
|
thanks for adding this! left some comments, mostly on rebasing past https://github.com/pytorch/ao/pull/363/files and code style |
| denormal_x = denormal_x.view(torch.float) | ||
|
|
||
| # adjust the denormal values back | ||
| denormal_x -= min_normal |
There was a problem hiding this comment.
The SR code up to this line looks good to me.
NicoleMayer
force-pushed
the
hanmei-e3m0-impl
branch
3 times, most recently
from
57165ff to
60fe4c3
Compare
NicoleMayer
force-pushed
the
hanmei-e3m0-impl
branch
from
60fe4c3 to
e202e6e
Compare
NicoleMayer
force-pushed
the
hanmei-e3m0-impl
branch
from
55dfe1f to
58a9f01
Compare
NicoleMayer
force-pushed
the
hanmei-e3m0-impl
branch
from
58a9f01 to
45520d2
Compare
No description provided.