[wip] add option to do activation/grad cast from hooks

[vkuzo]

Summary:

Testing moving activation casting logic into hooks, so we can start building towards composability of Float8 with DTensor

Note: needs pytorch/pytorch#116454 to land to enable backward pre hook in all cases

Current status:

1. test_base.py works for eager mode
2. test_compile.py fails with https://gist.github.com/vkuzo/3bfa1fcbb7b3c0ee0186eae2944eb75e , we need dynamo team help to debug this

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

[wanchaol]

I'm a bit surprised that torch.compile does not support this case with autograd function, as both DTensor pre-forward hook and forward hook are using autograd functions, probably because DTensor is using sth like allow_in_graph in dynamo. Maybe one way to workaround is to use a function that calls NoopFwToFloat8E5M2Bw.apply(output, module.emulate) inside, and then make it a allow_in_graph in dynamo?

[vkuzo]

note: backward pre hook works with pytorch/pytorch#116454, we still see the same dynamo error though (and not using torch.autograd.Function anymore)

[vkuzo]

@wanchaol , looks like the current dynamo issue is lack of support of Float8Tensor being in a subgraph boundary. Logs with evidence supporting this: https://gist.github.com/vkuzo/3cf42b6a54a68be1a45e8e1ca07eeb26. #166 has a more minimal repro.

Would you expect allow_in_graph to be related?