Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[not for land] example of simple FP8 UEX with stateful scaling #44

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[not for land] example of simple FP8 UEX with stateful scaling #44

wants to merge 1 commit into from

Commits on Apr 28, 2023

  1. [not for land] example of simple FP8 UEX with stateful scaling 

    This is an example of implementing basic fp8 support with a Python tensor subclass.

tl;dr;
1. FP8Tensor is the Python object which contains raw fp8 data (as torch.bits8), a scale, and a flavor (e4m3/e5m2)
2. FP8Tensor.__torch__dispatch knows how to add gradients, but converts to fp32 for everything else
3. FP8Linear is a module which can do stateful delayed scaling. User is expected to manually swap their linears to something like this.

Note: E4M3 support has not been numerically validated, and E5M2 support is not there at all
Note: No testing other than the bare bones at the bottom of the PR has been done.
Note: scaling is not implemented, currently it's just scales of 1.0 everywhere
    @vkuzo
    vkuzo committed Apr 28, 2023
    Configuration menu
    Copy the full SHA
    d47c566 View commit details
    Browse the repository at this point in the history