How to convert a AMP trained model to get best performance and speed?

[JohnHerry]

According to the doc: https://docs.pytorch.org/TensorRT/user_guide/mixed_precision.html We can convert model with this project where the param precision are explicitly said in the code. But when I train a model with torch AMP GradScaler where no value precision tagged in model code, Can we use this method to get a conerted chackpoint with best performance and inference speedup?

In fect, we had tried the torch pt->onnx-> tensorrt fp16 pipeline to convert pytorch AMP trained checkpoint into trt model format, but the inference results are noisey. while pt->onnx-> tensorrt fp32 pipeline will get a trt fp32 model the inference slower then what we need.

[narendasan]

I'm not sure I fully understand your question and I also don't know too deeply what ONNX-TensorRT does for supporting AMP graphs. For Torch-TRT, there are two modes you can compile with Torch-TensorRT, implicit (default) and explicit typing. Implicit typing ignores things like AMP or .to(dtype=...) casting and allows TensorRT to auto tune precision for best performance. Explicit typing preserves your own specified types such as those that come from AMP (after lowering). GradScalar iirc is for the backwards pass where as TRT only cares about the forward pass so the relevant information would be like what you would get out of torch.amp.autocast. If your code does not preserve these types, then I would recommend implicit typing. Note: in the future, only explicit typing will be supported.

Here is an example

import torch
import torch_tensorrt

class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.a_float32 = torch.rand((2, 2), device="cuda")
        self.b_float32 = torch.rand((2, 2), device="cuda")
        self.d_float32 = torch.rand((2, 2), device="cuda")

    def forward(self, x, y):
        # USE AUTOCAST TO SELECT OPERATIONS WHICH SHOULD RUN IN FP16
        with torch.autocast(x.device.type, enabled=True):
            e_float16 = torch.mm(self.a_float32, self.b_float32)
            f_float16 = torch.mm(self.d_float32, e_float16)
            g_float16 = torch.mm(f_float16, x)
            h_float16 = torch.add(g_float16, y)
        return h_float16

module = MyModule().to("cuda")
inp = (torch.tensor([[0.5, 0.5]], device="cuda").T, torch.tensor([0.5], device="cuda"))

ep = torch.export.export(module, inp)
out = ep.module()(*inp)
print(ep)

with torch_tensorrt.dynamo.Debugger(
    "graphs",
    logging_dir=".",
    capture_fx_graph_before=[torch_tensorrt.dynamo.lowering.ATEN_PRE_LOWERING_PASSES.passes[0].__name__],
    capture_fx_graph_after=[torch_tensorrt.dynamo.lowering.ATEN_POST_LOWERING_PASSES.passes[-1].__name__],
    save_engine_profile=True,
    profile_format="trex",
    engine_builder_monitor=False,
):
      # HERE WE COMPILE THE MODEL WITH TENSORRT BUT USE AMP'S TYPES 
      trt_mod = torch_tensorrt.compile(ep.module(), arg_inputs=inp, use_explicit_types=True, min_block_size=1)

print(trt_mod)
out_trt = trt_mod(*inp)
print(out - out_trt)

Here you can see before the PyTorch graph gets lowered, that there is an autocast subgraph in fp16 and an output in fp32.

And after Torch-TRT lowering this gets simplified to this graph where constants are folded there are some operators running in FP16 and others in FP32. With use_explicit_typing=True these types are reflected in the TRT directly with cast layers inserted according to PyTorch type promotion rules:

TRT and PyTorch don't use identical numerics which might be the cause of your inference instability but in practice we don't see that effecting model performance (if you do file a bug with https://github.com/NVIDIA/TensorRT). In Torch-TRT specifically we do things like use_fp32_acc=True to keep numerics closer to PyTorch. Unsure about what ONNX does here.