test

Author: lanluo-nvidia

examples/dynamo/vgg16_ptq.py

@@ -200,8 +200,6 @@ def calibrate_loop(model):
     quant_cfg = mtq.INT8_DEFAULT_CFG
 elif args.quantize_type == "fp8":
     quant_cfg = mtq.FP8_DEFAULT_CFG
-elif args.quantize_type == "fp4":
-    quant_cfg = mtq.NVFP4_DEFAULT_CFG
 # PTQ with in-place replacement to quantized modules
 mtq.quantize(model, quant_cfg, forward_loop=calibrate_loop)
 # model has FP8 qdq nodes at this point
@@ -241,8 +239,6 @@ def calibrate_loop(model):
             enabled_precisions = {torch.int8}
         elif args.quantize_type == "fp8":
             enabled_precisions = {torch.float8_e4m3fn}
-        elif args.quantize_type == "fp4":
-            enabled_precisions = {torch.float4_e2m1fn_x2}
         trt_model = torchtrt.dynamo.compile(
             exp_program,
             inputs=[input_tensor],

py/torch_tensorrt/dynamo/conversion/impl/quantize.py

@@ -88,7 +88,9 @@ def dynamic_block_quantize(
     Adds quantize and dequantize ops (QDQ) which quantize to FP4 based
     on the output_type set and dequantizes them back.
     """
-
+    print(
+        f"dynamic_block_quantize entered: {target=} {source_ir=} {name=} {input_tensor.shape=} {input_tensor.dtype=} {block_size=} {amax=} {num_bits=} {exponent_bits=} {scale_num_bits=} {scale_exponent_bits=}"
+    )
     with unset_fake_temporarily():
         if not isinstance(input_tensor, TRTTensor):
             input_tensor = get_trt_tensor(
@@ -114,18 +116,28 @@ def dynamic_block_quantize(
         # Add Q node
         dynamic_quantize_layer = ctx.net.add_dynamic_quantize(
             input_tensor,
-            -1,
-            16,
-            trt.DataType.FP4,
-            trt.DataType.FP8,
+            axis=1,
+            block_size=16,
+            output_type=trt.DataType.FP4,
+            scale_type=trt.DataType.FP8,
         )
-
         set_layer_name(
             dynamic_quantize_layer, target, name + "_dynamic_quantize", source_ir
         )
         q_output = dynamic_quantize_layer.get_output(0)
-        # Add DQ node
-        dequantize_layer = ctx.net.add_dequantize(q_output, scale)
+        q_scale = dynamic_quantize_layer.get_output(1)
+
+        # Add double DQ node
+        scale_dequantize_layer = ctx.net.add_dequantize(q_scale, scale)
+        scale_dequantize_layer.axis = 0
+        set_layer_name(
+            scale_dequantize_layer, target, name + "_scale_dequantize", source_ir
+        )
+        scale_dequantize_layer.precision = trt.DataType.FP8
+        scale_dq_output = scale_dequantize_layer.get_output(0)
+
+        dequantize_layer = ctx.net.add_dequantize(q_output, scale_dq_output)
+        dequantize_layer.axis = 1
         set_layer_name(dequantize_layer, target, name + "_dequantize", source_ir)
         dequantize_layer.precision = trt.DataType.FP4
         dq_output = dequantize_layer.get_output(0)

tests/py/dynamo/models/test_models_export.py

@@ -199,10 +199,10 @@ def test_resnet18_half(ir):
     torch._dynamo.reset()
 
 
-@unittest.skipIf(
-    torch.cuda.get_device_capability() < (10, 0),
-    "FP4 quantization requires compute capability 10.0 or later",
-)
+# @unittest.skipIf(
+#     torch.cuda.get_device_capability() < (10, 0),
+#     "FP4 quantization requires compute capability 10.0 or later",
+# )
 @unittest.skipIf(
     not importlib.util.find_spec("modelopt"),
     "ModelOpt is required to run this test",
@@ -215,20 +215,17 @@ def test_base_fp4(ir):
     class SimpleNetwork(torch.nn.Module):
         def __init__(self):
             super(SimpleNetwork, self).__init__()
-            self.linear1 = torch.nn.Linear(in_features=32, out_features=16)
-            self.linear2 = torch.nn.Linear(in_features=16, out_features=1)
+            self.linear1 = torch.nn.Linear(in_features=16, out_features=5)
 
         def forward(self, x):
             x = self.linear1(x)
-            x = torch.nn.ReLU()(x)
-            x = self.linear2(x)
             return x
 
     def calibrate_loop(model):
         """Simple calibration function for testing."""
         model(input_tensor)
 
-    input_tensor = torch.randn(1, 32).cuda()
+    input_tensor = torch.randn(1, 16).cuda()
     model = SimpleNetwork().eval().cuda()
 
     quant_cfg = mtq.NVFP4_DEFAULT_CFG
@@ -283,7 +280,6 @@ def calibrate_loop(model):
 
     input_tensor = torch.randn(1, 10).cuda()
     model = SimpleNetwork().eval().cuda()
-
     quant_cfg = mtq.FP8_DEFAULT_CFG
     mtq.quantize(model, quant_cfg, forward_loop=calibrate_loop)
     # model has FP8 qdq nodes at this point