Dynamic shapes support for tensorrt backend in torch.compile

[peri044]

TL DR:

This doc lists the existing way of dynamic shapes (DS) support in Torch-TRT using torch.compile and suggested improvements.
For a dynamic input shaped graph to be optimized by TensorRT , we would need the (min, max) shape info for the nodes. With Pytorch 2.X, we are able to get these ranges using the symbolic shape data in the node.meta and ensure TensorRT can use this data. However, there are some usability concerns in torch.compile workflow which are listed below.

Current DS workflow in export

class MyModule(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.conv = torch.nn.Conv2d(3, 16, 3, stride=1, bias=True)
            self.relu = torch.nn.ReLU()

        def forward(self, x):
            out = self.conv(x)
            out = self.relu(out)
            return out

model = MyModule().eval().cuda()
input = torch.randn((1, 3, 224, 224)).to("cuda")
   
dyn_dim = torch.export.Dim("x", min=2, max=8)
ep = torch.export.export(model, (input,), dynamic_shapes = ({0: dyn_dim},))

The dynamic_shapes argument is very helpful to provide all the DS info we need for any number of inputs and their dimensions.

Current DS workflow in compile

Based on this discussion thread, the current workflow that works is

class MyModule(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.conv = torch.nn.Conv2d(3, 16, 3, stride=1, bias=True)
            self.relu = torch.nn.ReLU()

        def forward(self, x):
            torch._check(x.size()[0] >= 1)
            torch._check(x.size()[0] <= 8)
            out = self.conv(x)
            out = self.relu(out)
            return out

model = MyModule().eval().cuda()
input_bs4 = torch.randn((4, 3, 224, 224)).to("cuda")
torch._dynamo.mark_dynamic(input_bs4, 0)

# Compile the model
trt_model = torch.compile(model, backend="tensorrt", options=compile_spec)
trt_model(input_bs4)

# Works for other batch sizes as well
input_bs6 = torch.randn((6, 3, 224, 224)).to("cuda")
trt_model(input_bs6)

In torch.compile workflow, we would have to add the torch._checks and mark the input dynamic via torch._dynamo.mark_dynamic

Limitations in the torch.compile workflow

We don't always have access to the model directly. Models sometimes get wrapped (if we use external third party libraries like huggingface). In the following GPT2 code, we have to modify the forward function in the source code which isn't always straightforward.

model = AutoModelForCausalLM.from_pretrained("gpt2", pad_token_id=tokenizer.eos_token_id).to(torch_device)

In general, it would be great if users don't have to modify their source code for their models.

Proposal

A straightforward way to handle this would be in torch._dynamo.mark_dynamic. If we could have an API like torch._dynamo.mark_dynamic(input, dimension, min, max), it would be easier for end-users. An end-end example is as follows

class MyModule(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.conv = torch.nn.Conv2d(3, 16, 3, stride=1, bias=True)
            self.relu = torch.nn.ReLU()

        def forward(self, x):
            out = self.conv(x)
            out = self.relu(out)
            return out

model = MyModule().eval().cuda()
input_bs4 = torch.randn((4, 3, 224, 224)).to("cuda")
torch._dynamo.mark_dynamic(input_bs4, 0, min=2, max=8)

# Compile the model
trt_model = torch.compile(model, backend="tensorrt", options=compile_spec)
trt_model(input_bs4)

# Works for other batch sizes as well
input_bs6 = torch.randn((6, 3, 224, 224)).to("cuda")
trt_model(input_bs6)

[peri044]

• mark_dynamic adds special attribute . look at impl. constrain_as_range API. shape_env reads of dynamic_dim which should read min, max.