Changes done internally at Facebook

d269be2fc7d84738a642d1d53eb44e6886a28d0c Alex Beloi <[email protected]> [fx] add deferred weights (xl_weight) and tracing for xl_embedding_bag
6f233bc9c72d90a908db0548c9d2dbe853895137 Alex Beloi <[email protected]> [fx] fix out of bounds indices/offsets for embedding_bag ops with xl_weight
3ca3b21c6a85ab9a6e9de503d0f13ee713a7b67c Janet Yang <[email protected]> Support div, torch.norm
52955d93d25e857510ed1b765220e8e5b0b0bb08 Janet Yang <[email protected]> Pass to replace sum(elmtwise(X))/numel(X) w/ mean(elmtwise(X))
89c56ef76a7a329f244a013ac5ccb099cb00c3c0 Janet Yang <[email protected]> Support scalar clamp, fixes for nan_to_num and benchmark
afdc533da031a64e162bb08c8629ff38739e24f8 Wei Wei <[email protected]> [fx2trt] disable dispatch trace leaf node test
d160a7a5e554d37c142e13f100bf4d8739ced232 Wei Wei <[email protected]> add option to remove passes
c22f691e6eae1b06ecd301eb6285b32d5dc9717c Mike Iovine <[email protected]> [fx2trt] Support dict inputs in acc tracer
8c05a3c57b1f5c63108b979ef8c61411525d0b1f Mike Iovine <[email protected]> [fx2trt] Support namedtuple access in acc tracer getattr
ff2000594e3f3ff75e0074edf9c38b5609128bbd Janet Yang <[email protected]> Generalize remove split ops more
1580805d827eb40c941e769b0b99e7c6a3ed6f89 Wei Wei <[email protected]> [fx2trt] add reshape unit test
d6a975462071a3747d18edcbe87a3b143b3ece88 Archie Sravankumar <[email protected]> Added FX tracing for `log_softmax`
6943ac0e322077b36a03c50c4c9065de6cd32837 Sungmin Cho <[email protected]> Add replace_mutable_op lower pass
baab27b81b1275de92fdaf760a158ce951564d33 Donglin Xia <[email protected]> Register avg_pool3d for acc_op in acc_op.py
ae4c4e2c3c18d78542140fcc30e1c24f7c647ef3 Wei Wei <[email protected]> [aten2trt] init check-in
87ef03338c9a25c5a610a2eb590345e8935f8d75 Wei Wei <[email protected]> [aten2trt] add binary ops
2bb168517ace7e638cffc7a241b1cbf528790b92 Mike Iovine <[email protected]> [fx2trt] Add acc normalization blocklist
8c912e085cf8722d572698286020ae1ce055023d Zhijing Li (Accelerator Enablement) <[email protected]> Skip unstable test_conv_add_standalone_module
b80dca9c9afa3b7d253e7806f48a890b9f83bf04 Jonathan Amazon <[email protected]> [PyTorch][FX][Compiler] Add acc_op tracing support for torch.baddbmm in FX
137a3977ffeb03d0387e8a95ff2f32f3d15b3de8 Wei Wei <[email protected]> [aten2trt] resnet support
fef54c237589a70c007c861e2d59c4052e3de054 Kefei Lu <[email protected]> [easy] fx2xxx: fix fuse_parallel_linear which changes getitem slices from tuple to list
4b062ef361cd7797e72c51bb4dc41766aca7b6db Kefei Lu <[email protected]> fx2trt: fix bad reshape pattern x.reshape(y.size(0), ...)
49573920892bb2fe75fe011a8cad9887bdc8bd04 Alex Beloi <[email protected]> [FX] add tracing for torch.detach
fe3cc75e775af53f603a83e8b4899b28f3cb6ddc Yinghai Lu <[email protected]> [fx2ait] add support to torch.clip
42c54d69c68dc58ac348647acada88b1e5634b40 Fei Kou <[email protected]> Fix clamping float32 boundary values
e013621dedf5960f81b915cef8d2ce19ca349a7a Kefei Lu <[email protected]> trt lower: change preset application logic to in-place instead of immutable update
adc9f8ff48c01a0ce70080c930221ac81f048563 Kefei Lu <[email protected]> [easy]: fix another instance of [slice(), ...] to (slice(), ...)
a22e9ff2cc55eb8669690eedd6971be93a2a356b Rui Zhu <[email protected]> Support NoneType in acc_tracing by setting its meta shape to be 1
4f54ce9283f02fe416ff3f502ef1a4e4f80c0f37 Mike Iovine <[email protected]> [fx2ait] Avoid extra copies from view ops
0baf42ebf6ce4146df1bee2d2e62fa2b77dbd7fb Mor Tzur <[email protected]> add torch.concat to acc_ops
9cd933707772b0f05b8aca62bcc813929bd52868 Shirong Wu <[email protected]> replace assert_allclose with assert_close
e418d0653752022ea4ee186036b79dc8ca0ae87b Valeriu Lacatusu <[email protected]> [PyTorch][FX][Compiler] Add acc_op tracing support for torch.nn.functional.softplus in FX
afb2f560b3995ea3a1cd440df3cdd66d92472e46 Wei Wei <[email protected]> [fx2trt] test fix to adopt new interface of dynamo
8ca2307c744f13ef15bad49f5030dddd2b787b9d Huamin Li <[email protected]> rename test_setitem to test_setitem_trt
e0b75bbfda8604d4b60599ddba4d4aa7023887a5 Valeriu Lacatusu <[email protected]> [FX] Replace deprecated torch.testing.assert_allclose with torch.testing.assert_close
4a233da979a755fa605e9750c6035ed885597afa Valeriu Lacatusu <[email protected]> [PyTorch][FX][Compiler] Add acc_op tracing support for torch.ops._caffe2.RoIAlign in FX

Author: Wei Wei

docs/_sources/tutorials/ptq.rst.txt

@@ -136,7 +136,7 @@ Then all thats required to setup the module for INT8 calibration is to set the f
 If you have an existing Calibrator implementation for TensorRT you may directly set the ``ptq_calibrator`` field with a pointer to your calibrator and it will work as well.
 From here not much changes in terms of how to execution works. You are still able to fully use LibTorch as the sole interface for inference. Data should remain
 in FP32 precision when it's passed into `trt_mod.forward`. There exists an example application in the Torch-TensorRT demo that takes you from training a VGG16 network on
-CIFAR10 to deploying in INT8 with Torch-TensorRT here: https://github.com/pytorch/TensorRT/tree/master/cpp/ptq
+CIFAR10 to deploying in INT8 with Torch-TensorRT here: https://github.com/pytorch/TensorRT/tree/master/examples/int8/ptq
 
 .. _writing_ptq_python:
 
@@ -194,8 +194,8 @@ to use ``CacheCalibrator`` to use in INT8 mode.
                                         calibrator=calibrator)
 
 If you already have an existing calibrator class (implemented directly using TensorRT API), you can directly set the calibrator field to your class which can be very convenient.
-For a demo on how PTQ can be performed on a VGG network using Torch-TensorRT API, you can refer to https://github.com/pytorch/TensorRT/blob/master/tests/py/test_ptq_dataloader_calibrator.py
-and https://github.com/pytorch/TensorRT/blob/master/tests/py/test_ptq_trt_calibrator.py
+For a demo on how PTQ can be performed on a VGG network using Torch-TensorRT API, you can refer to https://github.com/pytorch/TensorRT/blob/master/tests/py/ptq/test_ptq_dataloader_calibrator.py
+and https://github.com/pytorch/TensorRT/blob/master/tests/py/ptq/test_ptq_trt_calibrator.py
 
 Citations
 ^^^^^^^^^^^

examples/fx/hugging_face_torchdynamo_example.py

@@ -15,12 +15,12 @@
 )
 from transformers import BertConfig, ReformerConfig, XLNetModel, XLNetConfig
 
-import torchdynamo
-from torchdynamo.optimizations import backends
-from torchdynamo.optimizations.training import aot_autograd_debug_strategy1
-from torchdynamo.optimizations.training import aot_autograd_speedup_strategy
-from torchdynamo.testing import collect_results
-from torchdynamo.testing import same
+import torch._dynamo as torchdynamo
+from torch._dynamo.optimizations import backends
+from torch._dynamo.optimizations.training import aot_autograd_debug_strategy1
+from torch._dynamo.optimizations.training import aot_autograd_speedup_strategy
+from torch._dynamo.testing import collect_results
+from torch._dynamo.testing import same
 
 torch.backends.cuda.matmul.allow_tf32 = True
 

examples/fx/torchdynamo_example.py

@@ -3,11 +3,11 @@
 from dataclasses import dataclass, field, replace
 
 import torch
-import torchdynamo
+import torch._dynamo as torchdynamo
 import torchvision
 from torch_tensorrt.fx.lower import compile
 from torch_tensorrt.fx.utils import LowerPrecision
-from torchdynamo.optimizations import backends
+from torch._dynamo.optimizations import backends
 
 """
 The purpose of this example is to demostrate the lowering flow to TRT and Torchdynamo

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -25,6 +25,7 @@
     trt_transposed_linear,
     trt_transposed_matmul,
 )
+from torch_tensorrt.fx.tracer.acc_tracer.acc_ops import contiguous
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -3371,6 +3372,9 @@ def acc_ops_gelu(
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
     input_val = kwargs["input"]
+    approximate = kwargs["approximate"]
+    if approximate is not "none":
+        raise RuntimeError("GeLU converter currently doesn't support fast gelu compute")
     if not isinstance(input_val, TRTTensor):
         raise RuntimeError(
             f"GELU received input {input_val} that is not part "

py/torch_tensorrt/fx/lower.py

@@ -29,6 +29,7 @@
 def compile(
     module: nn.Module,
     input,
+    min_acc_module_size: int = 10,
     max_batch_size: int = 2048,
     max_workspace_size=1 << 25,
     explicit_batch_dimension=False,
@@ -48,6 +49,7 @@ def compile(
         module: Original module for lowering.
         input: Input for module.
         max_batch_size: Maximum batch size (must be >= 1 to be set, 0 means not set)
+        min_acc_module_size: Minimal number of nodes for an accelerated submodule
         max_workspace_size: Maximum size of workspace given to TensorRT.
         explicit_batch_dimension: Use explicit batch dimension in TensorRT if set True, otherwise use implicit batch dimension.
         lower_precision: lower_precision config given to TRTModule.
@@ -61,6 +63,7 @@ def compile(
     """
     lower_setting = LowerSetting(
         max_batch_size=max_batch_size,
+        min_acc_module_size=min_acc_module_size,
         max_workspace_size=max_workspace_size,
         explicit_batch_dimension=explicit_batch_dimension,
         lower_precision=lower_precision,
@@ -237,6 +240,7 @@ def __call__(
         module: nn.Module,
         inputs: Input,
         additional_inputs: Optional[Input] = None,
+        fp16_conversion_fn: Optional[Callable[[Input], Input]] = None,
     ) -> nn.Module:
         lower_setting = self.lower_pass_manager_builder.lower_setting
         atol = lower_setting.correctness_atol
@@ -253,9 +257,26 @@ def do_lower(module: nn.Module, inputs: Input) -> nn.Module:
                 == LowerPrecision.FP16
             ):
                 module.half()
-                inputs = tuple(
-                    x.half() if x is not None and x.dtype == torch.float32 else x
-                    for x in inputs
+                # A custom conversion function can be passed to the lowerer to
+                # handle inputs with custom types. By default, just handle
+                # tensors and NoneType.
+                if fp16_conversion_fn is None:
+                    conversion_fn = (
+                        lambda x: x.half()
+                        if x is not None and x.dtype == torch.float32
+                        else x
+                    )
+                else:
+                    conversion_fn = fp16_conversion_fn
+
+                inputs = tuple(conversion_fn(x) for x in inputs)
+            if lower_setting.is_aten:
+                pm = self.lower_pass_manager_builder.build_aten2trt_lower_pipeline(
+                    inputs, additional_inputs
+                )
+            else:
+                pm = self.lower_pass_manager_builder.build_trt_lower_pipeline(
+                    inputs, additional_inputs
                 )
             if lower_setting.is_aten:
                 pm = self.lower_pass_manager_builder.build_aten2trt_lower_pipeline(

py/torch_tensorrt/fx/test/converters/acc_op/test_cat.py

@@ -1,32 +1,51 @@
 import torch
 import torch.nn as nn
 import torch_tensorrt.fx.tracer.acc_tracer.acc_ops as acc_ops
+from parameterized import param, parameterized
 from torch.testing._internal.common_utils import run_tests
 from torch_tensorrt.fx.tools.common_fx2trt import AccTestCase, InputTensorSpec
 
 
 class TestCatConverter(AccTestCase):
-    def test_cat(self):
+    @parameterized.expand(
+        [
+            param("cat", torch.cat),
+            param("concat", torch.concat),
+        ]
+    )
+    def test_cat(self, _, op):
         class Cat(nn.Module):
             def forward(self, x, y, z):
-                return torch.cat((x, y, z), 1)
+                return op((x, y, z), 1)
 
         inputs = [torch.randn(1, 2, 3), torch.randn(1, 1, 3), torch.randn(1, 3, 3)]
         self.run_test(Cat(), inputs, expected_ops={acc_ops.cat})
 
-    def test_cat_neg(self):
+    @parameterized.expand(
+        [
+            param("cat", torch.cat),
+            param("concat", torch.concat),
+        ]
+    )
+    def test_cat_neg(self, _, op):
         class Cat(nn.Module):
             def forward(self, x, y, z):
-                return torch.cat((x, y, z), -1)
+                return op((x, y, z), -1)
 
         inputs = [torch.randn(1, 2, 3), torch.randn(1, 2, 3), torch.randn(1, 2, 2)]
         self.run_test(Cat(), inputs, expected_ops={acc_ops.cat})
 
-    def test_cat_with_dynamic_shape(self):
+    @parameterized.expand(
+        [
+            param("cat", torch.cat),
+            param("concat", torch.concat),
+        ]
+    )
+    def test_cat_with_dynamic_shape(self, _, op):
         class Cat(nn.Module):
             def forward(self, x, y):
                 x = x + y
-                return torch.cat((x, y), 0)
+                return op((x, y), 0)
 
         input_specs = [
             InputTensorSpec(
@@ -42,11 +61,17 @@ def forward(self, x, y):
         ]
         self.run_test_with_dynamic_shape(Cat(), input_specs, expected_ops={acc_ops.cat})
 
-    def test_cat_with_dynamic_shape_four_dimensions(self):
+    @parameterized.expand(
+        [
+            param("cat", torch.cat),
+            param("concat", torch.concat),
+        ]
+    )
+    def test_cat_with_dynamic_shape_four_dimensions(self, _, op):
         class Cat(nn.Module):
             def forward(self, x, y):
                 x = x + y
-                return torch.cat((x, y), 0)
+                return op((x, y), 0)
 
         input_specs = [
             InputTensorSpec(
@@ -63,6 +88,14 @@ def forward(self, x, y):
 
         self.run_test_with_dynamic_shape(Cat(), input_specs, expected_ops={acc_ops.cat})
 
+    def test_concat(self):
+        class Cat(nn.Module):
+            def forward(self, x, y, z):
+                return torch.concat((x, y, z), 1)
+
+        inputs = [torch.randn(1, 2, 3), torch.randn(1, 1, 3), torch.randn(1, 3, 3)]
+        self.run_test(Cat(), inputs, expected_ops={acc_ops.cat})
+
 
 if __name__ == "__main__":
     run_tests()

py/torch_tensorrt/fx/test/converters/acc_op/test_gelu.py

@@ -57,6 +57,39 @@ def forward(self, x):
             TestModule(), input_specs, expected_ops={acc_ops.gelu}
         )
 
+    def test_gelu_module(self):
+        class TestModule(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.gelu = torch.nn.GELU()
+
+            def forward(self, x):
+                return self.gelu(x)
+
+        inputs = [torch.randn(3, 10, 20)]
+        self.run_test(
+            TestModule(),
+            inputs,
+            expected_ops={acc_ops.gelu},
+            test_implicit_batch_dim=False,
+        )
+
+    def test_gelu_module_throw(self):
+        class TestModule(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.gelu = torch.nn.GELU(approximate="tanh")
+
+            def forward(self, x):
+                return self.gelu(x)
+
+        inputs = [torch.randn(3, 10, 20)]
+        self.run_test_with_assert_error(
+            TestModule(),
+            inputs,
+            expect_error=RuntimeError,
+        )
+
 
 if __name__ == "__main__":
     run_tests()

py/torch_tensorrt/fx/test/converters/acc_op/test_new_ones.py

@@ -6,19 +6,6 @@
 
 
 class TestNewOnesConverter(AccTestCase):
-    def test_newone(self):
-        class TestModule(nn.Module):
-            def forward(self, x):
-                return x.new_ones((3, 5), dtype=torch.float16)
-
-        inputs = [torch.randn(1, 10)]
-        self.run_test(
-            TestModule(),
-            inputs,
-            expected_ops={acc_ops.new_ones},
-            test_implicit_batch_dim=False,
-        )
-
     def test_newone_no_dtype(self):
         class TestModule(nn.Module):
             def forward(self, x):
@@ -47,23 +34,6 @@ def forward(self, x):
 
 
 class TestNewOnesConverterWithDynamicShape(AccTestCase):
-    def test_newone(self):
-        class TestModule(nn.Module):
-            def forward(self, x):
-                return x.new_ones((3, 5), dtype=torch.float16)
-
-        input_specs = [
-            InputTensorSpec(
-                shape=(-1, -1, -1, -1),
-                dtype=torch.float32,
-                shape_ranges=[((1, 1, 1, 1), (2, 3, 4, 5), (2, 3, 10, 10))],
-            ),
-        ]
-
-        self.run_test_with_dynamic_shape(
-            TestModule(), input_specs, expected_ops={acc_ops.new_ones}
-        )
-
     def test_newone_no_dtype(self):
         class TestModule(nn.Module):
             def forward(self, x):

py/torch_tensorrt/fx/test/converters/acc_op/test_unary_ops.py

@@ -9,29 +9,31 @@
 from torch_tensorrt.fx.tools.common_fx2trt import AccTestCase, InputTensorSpec
 
 unary_ops = [
-    (torch.sin, acc_ops.sin),
-    (torch.cos, acc_ops.cos),
-    (torch.tan, acc_ops.tan),
-    (torch.sinh, acc_ops.sinh),
-    (torch.cosh, acc_ops.cosh),
-    (torch.asin, acc_ops.asin),
-    (torch.acos, acc_ops.acos),
-    (torch.atan, acc_ops.atan),
-    (torch.abs, acc_ops.abs),
-    (torch.neg, acc_ops.neg),
-    (torch.reciprocal, acc_ops.reciprocal),
-    (torch.sqrt, acc_ops.sqrt),
-    (torch.log, acc_ops.log),
-    (torch.exp, acc_ops.exp),
-    (torch.floor, acc_ops.floor),
-    (torch.ceil, acc_ops.ceil),
-    (torch.sign, acc_ops.sign),
+    (torch.sin, acc_ops.sin, False),
+    (torch.cos, acc_ops.cos, False),
+    (torch.tan, acc_ops.tan, False),
+    (torch.sinh, acc_ops.sinh, False),
+    (torch.cosh, acc_ops.cosh, False),
+    (torch.asin, acc_ops.asin, True),
+    (torch.acos, acc_ops.acos, True),
+    (torch.atan, acc_ops.atan, True),
+    (torch.abs, acc_ops.abs, False),
+    (torch.neg, acc_ops.neg, False),
+    (torch.reciprocal, acc_ops.reciprocal, False),
+    (torch.sqrt, acc_ops.sqrt, False),
+    (torch.log, acc_ops.log, False),
+    (torch.exp, acc_ops.exp, False),
+    (torch.floor, acc_ops.floor, False),
+    (torch.ceil, acc_ops.ceil, False),
+    (torch.sign, acc_ops.sign, False),
 ]
 
 
 class TestUnaryOpConverters(AccTestCase):
-    @parameterized.expand([(op[1].__name__, op[0], op[1]) for op in unary_ops])
-    def test_unary_ops(self, name, orig_op: Callable, expected_op):
+    @parameterized.expand([(op[1].__name__, op[0], op[1], op[2]) for op in unary_ops])
+    def test_unary_ops(
+        self, name, orig_op: Callable, expected_op: Callable, range_req: bool
+    ):
         class TestModule(nn.Module):
             def __init__(self, orig_op):
                 super().__init__()
@@ -41,11 +43,15 @@ def forward(self, x):
                 return self.orig_op(x)
 
         m = TestModule(orig_op)
-        inputs = [torch.randn(2, 2, 3)]
+        inputs = (
+            [torch.distributions.uniform.Uniform(-1, 1).sample([2, 2, 3])]
+            if range_req
+            else [torch.randn(2, 2, 3)]
+        )
         self.run_test(m, inputs, expected_ops={expected_op})
 
 
-class TestUnaryOpConvertersWithDynamicShapeFourDimensions(AccTestCase):
+class TestUnaryVOpConvertersWithDynamicShapeFourDimensions(AccTestCase):
     @parameterized.expand([(op[1].__name__, op[0], op[1]) for op in unary_ops])
     def test_unary_ops(self, name, orig_op: Callable, expected_op):
         class TestModule(nn.Module):

py/torch_tensorrt/fx/test/core/test_trt_module.py

@@ -27,7 +27,7 @@ def forward(self, x):
         torch.save(trt_mod, "trt.pt")
         reload_trt_mod = torch.load("trt.pt")
 
-        torch.testing.assert_allclose(
+        torch.testing.assert_close(
             reload_trt_mod(inputs[0].cuda()).cpu(), ref_output, rtol=1e-04, atol=1e-04
         )
         os.remove(f"{os.getcwd()}/trt.pt")
@@ -49,7 +49,7 @@ def forward(self, x):
         new_trt_mod = TRTModule()
         new_trt_mod.load_state_dict(st)
 
-        torch.testing.assert_allclose(
+        torch.testing.assert_close(
             new_trt_mod(inputs[0].cuda()).cpu(), ref_output, rtol=1e-04, atol=1e-04
         )