This PR rewrites Float8DynamicLinear to use module hooks, as we think long term this is more composable with other PyTorch features. For now there is no plan to land this, this is just reproducing / sharing what breaks when we try this today.
Test Plan:
// note: all tests pass without this PR
// eager mode is fine
> pytest -s test/test_base.py | with-proxy gh gist create
https://gist.github.com/vkuzo/aded224af91092c8326becc855b125c9
// compile has some errors in aot_eager backend
> pytest -s test/test_compile.py | with-proxy gh gist create
https://gist.github.com/vkuzo/cab55b11a2c3cee0d1ff94169131b171
// dtensor + float8 has numeric issues
> ./test/test_dtensor.sh | with-proxy gh gist create
https://gist.github.com/vkuzo/d1035200db22f2e3357438824cd3594f
Summary:
This PR rewrites Float8DynamicLinear to use module hooks, as we think long term this is more composable with other PyTorch features. For now there is no plan to land this, this is just reproducing / sharing what breaks when we try this today.
Test Plan:
Reviewers:
Subscribers:
Tasks:
Tags: