pytorch

8085e08a - [TP] Add prepareInput and output for input/output DTensor layout annotation in the parent module in TP API (#111166)

[TP] Add prepareInput and output for input/output DTensor layout annotation in the parent module in TP API (#111166) In some use cases, we found that users might want to annote the input/output DTensor layout for the parent module rather than the submodule whose parameters are to be distributed so that we want to have these two class for users to annote input/output DTensor layouts so that we register pre-FWD/FWD hook for the TP-lized module. Pull Request resolved: https://github.com/pytorch/pytorch/pull/111166 Approved by: https://github.com/wanchaol ghstack dependencies: #111160