🐛 [Bug] Encountered bug when using Torch-TensorRT

[stupiding]

Bug Description

I'm running a smallest test code. When I use torch_tensorrt to compile the model, it throw s shape mismatch error in TorchScript, but the torch.jit.trace itself didn't show the same error.
Can you please help me with this problem?

To Reproduce

Steps to reproduce the behavior:

1. python trt_sample.py

code sample:

`

    import time
    import random

    import torch
    import torch_tensorrt
    from torch import nn

    torch.manual_seed(0)


    class Model(torch.nn.Module):
        def __init__(self, n_in=128, n_out=64, norm='LN', act=True):
            super().__init__()

            self.linear = nn.Linear(n_in, n_out, bias=False)
            self.norm = nn.LayerNorm(n_out)
            self.relu = nn.ReLU(inplace=True)
            self.act = act

        def forward(self, cout, rand, emb):
            idcs = torch.where(rand > 0.2)[0]

            out = self.linear(emb.view(1, -1, 128))

            out = self.norm(out)

            if self.act:
                out = self.relu(out)

            cout[:, idcs.long()] = out
            # cout.index_add_(1, idcs, out)
            return cout


    # Init data
    rand = torch.rand(5740).cuda()
    s = torch.where(rand>0.2)[0]

    emb = torch.rand((1, len(s) * 128)).cuda()
    cout = torch.rand(1, 5740, 64).cuda()

    # Init model
    model = Model().eval().cuda()

    # torch jit trace
    traced = torch.jit.trace(model, (cout.clone(), rand.clone(), emb.clone()))
    traced_output = traced(cout.clone(), rand.clone(), emb.clone())
    print('traced output: ', traced_output.shape)

    # trt compile
    enabled_precisions = {torch.float}
    trt_ts_module = torch_tensorrt.compile(
        model,  # traced,
        inputs=(cout.clone(), rand.clone(), emb.clone()),
        enabled_precisions=enabled_precisions,
        torch_executed_ops=['aten::index', 'aten::index_put_', 'aten::where', 'aten::to'],
        truncate_long_and_double=True,
        debug=True
    )
    print('trt compiled')

    # infer with trt
    result = trt_ts_module(cout.clone(), rand.clone(), emb.clone())
    print('trt: ', time.time() - start_time)

`

error message:
`

Traceback (most recent call last):
  File "test_trace.py", line 52, in <module>
    trt_ts_module = torch_tensorrt.compile(
  File "/opt/conda/lib/python3.8/site-packages/torch_tensorrt/_compile.py", line 113, in compile
    return torch_tensorrt.ts.compile(ts_mod, inputs=inputs, enabled_precisions=enabled_precisions, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch_tensorrt/ts/_compiler.py", line 134, in compile
    compiled_cpp_mod = _C.compile_graph(module._c, _parse_compile_spec(spec))
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "test_trace.py", line 30, in forward
            out = self.relu(out)

        cout[:, idcs.long()] = out
        ~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        # cout.index_add_(1, idcs, out)
        return cout
RuntimeError: shape mismatch: value tensor of shape [1, 4582, 64] cannot be broadcast to indexing result of shape [1, 4631, 64]

`

Expected behavior

run without shape missmatch

Environment

docker image: nvcr.io/nvidia/pytorch:22.09-py3

[gs-olive]

The origin of the issue seems to be the rand argument in the forward function. During partitioning, Torch-TRT currently generates random inputs to the the forward function, and so when the line idcs = torch.where(rand > 0.2)[0] is executed on internally-generated inputs, its size differs from that of the externally-generated rand. We are working on adding an argument to specify sample input tensors and value ranges currently (RFC #1425). A temporary solution would be to change rand to s, so the indices are computed outside of the forward function, potentially like so:

def forward(self, cout, s, emb):
    idcs = s
    ...

...

s = torch.where(rand > 0.2)[0].int().cuda()

### Then replace all instances of rand.clone() with s.clone()

The above changes make the code and inference functional on my machine, using the current latest version of Torch-TRT (commit e3b9929). Let me know if this also works on your end.

[stupiding]

The origin of the issue seems to be the rand argument in the forward function. During partitioning, Torch-TRT currently generates random inputs to the the forward function, and so when the line idcs = torch.where(rand > 0.2)[0] is executed on internally-generated inputs, its size differs from that of the externally-generated rand. We are working on adding an argument to specify sample input tensors and value ranges currently (RFC #1425). A temporary solution would be to change rand to s, so the indices are computed outside of the forward function, potentially like so:

def forward(self, cout, s, emb):
    idcs = s
    ...

...

s = torch.where(rand > 0.2)[0].int().cuda()

### Then replace all instances of rand.clone() with s.clone()

The above changes make the code and inference functional on my machine, using the current latest version of Torch-TRT (commit e3b9929). Let me know if this also works on your end.

Great! This works for me, and looking forward to the specify-input-tensor update~

[stupiding]

Sorry but another wierd problem occurs when I move the cout initialization into the module's forward

    def forward(self, idcs, emb):
        bs = emb.shape[0]
        out = self.linear(emb.view(1, -1, 128))

        out = self.norm(out)

        if self.act:
            out = self.relu(out)

        cout = torch.zeros([bs, 5740, 64], device=emb.device)
        cout[:, idcs.long()] = out[:, idcs.long()]
        return cout  # cout[:, idcs.long()], out

if I compile the model directly, the results are all good

but when I trace the model first and use it in the compile argments, the forward result tensor is all zeros

trt_ts_module = torch_tensorrt.compile(
   　traced,  # model
    　inputs=(s.clone(), emb.clone()),
    　enabled_precisions=enabled_precisions,
    　torch_executed_ops=['aten::index', 'aten::index_put_', 'aten::where', 'aten::to'],
    　truncate_long_and_double=True,
    　debug=True
     )

the compilation says nothing wrong, but something like the code
cout[:, idcs.long()] = out
was not executed

[gs-olive]

When I substitute in the above code, I am seeing a CUDA out-of-bounds error, because it seems that the cout[:, idcs.long()] = out[:, idcs.long()] line in the forward function tries to access out-of-bounds indices on the out tensor (because idcs includes indices larger than the size of out). Could you send the entire script you are using which is displaying the behavior, including tensor initializations? This will help to isolate the problem

[stupiding]

import time
import random

import torch
import torch_tensorrt
from torch import nn

torch.manual_seed(0)


class Model(torch.nn.Module):
    def __init__(self, n_in=128, n_out=64, norm='LN', act=True):
        super().__init__()

        self.linear = nn.Linear(n_in, n_out, bias=False)
        self.norm = nn.LayerNorm(n_out)
        self.relu = nn.ReLU(inplace=True)
        self.act = act
        # self.register_buffer('zero_buffer', torch.zeros([4, 5740, 64]))

    def forward(self, idcs, emb):
        bs = emb.shape[0]
        out = self.linear(emb.view(1, -1, 128))

        out = self.norm(out)

        if self.act:
            out = self.relu(out)

        cout = torch.zeros([bs, 5740, 64], device=emb.device)
        cout[:, idcs.long()] = out[:, idcs.long()]
        return cout  # cout[:, idcs.long()], out


# Init data
rand = torch.rand(5740).cuda().float()
s = torch.where(rand>0.2)[0]

emb = torch.rand((1, 5740 * 128)).cuda()

cout = torch.zeros([1, 5740, 64], device=emb.device)

# Init model
model = Model().eval().cuda()

# torch jit trace
traced = torch.jit.trace(model, (s, emb))
print(traced.graph)

# trt compile
enabled_precisions = {torch.float}
trt_ts_module = torch_tensorrt.compile(
    traced,
    inputs=(s.int(), emb),
    enabled_precisions=enabled_precisions,
    torch_executed_ops=['aten::index', 'aten::index_put_', 'aten::where'],  # , 'aten::to'],
    truncate_long_and_double=True,
    debug=True
)
print('trt compiled')

emb = torch.rand((1, 5740 * 128)).cuda()

model_output = model(s, emb)

traced_output = traced(s, emb)
print('traced output: ', traced_output.shape)

# infer with trt
trt_output = trt_ts_module(s.int(), emb)
print('trt: ', trt_output.shape, trt_output.max())

diff = traced_output - trt_output
print(diff.max(), model_output.max(), traced_output.max(), trt_output.max())

[narendasan]

@bowang007 can you look at this case for the int64 <-> int32 partitioning?
@gs-olive Is looking at a lowering pass to handing in place index_put_ -> out-place index_put