pytorch
757e3cbf - ns for fx: add support for shadowing linear fp16 patterns (#54275)

Commit
3 years ago
ns for fx: add support for shadowing linear fp16 patterns (#54275) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/54275 Adds support for NS shadow activations path for the fp16 emulation pattern such as ``` ... -> dequantize -> linear -> relu -> to(torch.float16) -> ... ``` There are a couple of changes necessary here: 1. removing the restriction on the shadowing graph pass that the B subgraph is a single node (since this subgraph is four nodes), and modifying the code to correctly add the relevant inputs versus output loggers (input loggers and subgraph copy if we are at start_node, and output logger if we are at end_node) 2. modifying the logic for calculating node input and output type to work correcty for the `to` and `dequantize` nodes: 2a. make the function return the first input and output, instead of just the first input 2b. make the function handle `dequantize` correctly by recursively using the output if its input 2c. make the function handle `to` correctyl by recursively using the output of its input and the target dtype 3. a bug fix to handle observers in kwargs, while copying subgraphs Note: input logging for these patterns is not tested yet, this will be in the next PR. Test Plan: ``` python test/test_quantization.py TestFXNumericSuiteCoreAPIs.test_linear_fp16 ``` Imported from OSS Reviewed By: jerryzh168 Differential Revision: D27172655 fbshipit-source-id: 3bdc86618b2a5782627fcf303d58af7f47fbc30d
Author
Parents
Loading