lightning-thunder
lightning-thunder copied to clipboard
Transformer Engine's recipes UX
🚀 Feature
Branching the conversation from #1908(comment) there are multiple possibilities for the user to specify a recipe for TE computation.
However on the high level, there are two approaches:
- Thunder takes care of deciding which TE recipe is best for your and uses that throughout the program and advanced users that want to customize the choice are given the opportunity to do so by adding an option to the instantiation of the executor, or
- Transformer Engine
fp8_autocastcontext manager is fully integrated with Thunder and the user has to create and decide which recipe and when to use.
In either case I believe that there should be complete commitment to the choice, for example: committing to point 2 means matching TE behavior, and allowing high precision computation outside the context manager.
cc. @kshitij12345 @IvanYashchuk
cc @borda
Is this still an open discussion, or do we have a preferred approach? Given #2437, maybe unifying everything under option 1 is clearer in terms of UX?