Update on "[ONNX] Support None in fx.args as torchlib inputs"

titaiwangms · titaiwangms · commit c961589165c2 · 2023-09-12T23:13:36.000Z
Prior to this PR, if None is returned from intermediate nodes, it will crashes the export because None is not expected to be passed into `_fill_tensor_shape_type`, and raise beartype roar. The function fills in shape and type to TorchScriptTensor according to its info from FX graph. This is discovered after microsoft/onnxscript#1043 is supported. The op specifically generates None in one of its inputs, but the only output from it being consumed is the first one (not None). Reference test from a TorchBench model: ```python def test_nanogpt(self): import sys sys.path.append("/home/titaiwang") from nanoGPT.model import GPT, GPTConfig # Load the model kwargs = { "block_size": 256, "vocab_size": 8096, # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency "n_layer": 2, "n_head": 2, "n_embd": 128, "dropout": 0.0, "bias": False, # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster } config = GPTConfig(**kwargs) with torch.backends.cuda.sdp_kernel( enable_flash=True, enable_mem_efficient=True ): model = GPT(config) print("Done loading model") inputs = torch.arange(128).view(2, 64) targets = torch.arange(128).view(2, 64) self.run_test_with_fx_to_onnx_exporter_and_onnx_runtime( model, (inputs,), input_kwargs={ "targets": targets, }, verbose=True, ) ``` [ghstack-poisoned]
diff --git a/torch/onnx/_internal/fx/fx_onnx_interpreter.py b/torch/onnx/_internal/fx/fx_onnx_interpreter.py
@@ -236,6 +236,9 @@ def _fill_tensor_shape_type(
     ):
         if expected_value is None:
             # There is no shape/type from None.
+            # NOTE: according to https://github.com/pytorch/pytorch/blob/main/torch/_meta_registrations.py,
+            # None could be a valid value for return type, so we need to handle it.
+            # e.g. the function: meta__scaled_dot_product_flash() in cpu mode.
             continue
         elif isinstance(expected_value, (torch.SymInt, torch.SymFloat, torch.SymBool)):
             # aten::sym_size output is a int, not a tensor, which stands
