fix: cumsum add_constant bug fix (add dtype for np zeros)

[chohk88]

Description

When compiling the roberta-base model from Hugging Face (https://huggingface.co/FacebookAI/roberta-base), a TypeError occurs in the cumsum operation. For static shape input, the default datatype of np.zeros(new_dims) function is np.float64 which is not handled properly by the create_constant utility function.

Fixes # (issue)

  File "/usr/local/lib/python3.10/dist-packages/torch_tensorrt/dynamo/conversion/aten_ops_converters.py", line 934, in aten_ops_cumsum
    return impl.slice.cumsum(
  File "/usr/local/lib/python3.10/dist-packages/torch_tensorrt/dynamo/conversion/impl/slice/ops.py", line 387, in cumsum
    zero_trttensor = get_trt_tensor(ctx, zeros, f"{name}_initial_value")
  File "/usr/local/lib/python3.10/dist-packages/torch_tensorrt/dynamo/conversion/converter_utils.py", line 388, in get_trt_tensor
    return create_constant(ctx, input_val, name, dtype, min_rank)
  File "/usr/local/lib/python3.10/dist-packages/torch_tensorrt/dynamo/conversion/converter_utils.py", line 349, in create_constant
    constant = ctx.net.add_constant(
torch._dynamo.exc.BackendCompilerFailed: backend='torch_tensorrt' raised:
TypeError: add_constant(): incompatible function arguments. The following argument types are supported:
    1. (self: tensorrt.tensorrt.INetworkDefinition, shape: tensorrt.tensorrt.Dims, weights: tensorrt.tensorrt.Weights) -> tensorrt.tensorrt.IConstantLayer
Invoked with: <tensorrt.tensorrt.INetworkDefinition object at 0x7fee84ebd770>, (1,), array([0.])

Reproduction Code:

# https://huggingface.co/FacebookAI/roberta-base
import torch
from transformers import RobertaTokenizer, RobertaModel
import torch_tensorrt

backend = "torch_tensorrt"
device = "cuda:0"

# Load tokenizer and model
tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
model = RobertaModel.from_pretrained('roberta-base')
model = model.to(device)

# Tokenize input text
text = "Replace me by any text you'd like."
encoded_input = tokenizer(text, return_tensors='pt')
encoded_input = {k: v.to(device) for k, v in encoded_input.items()} 

# Compile model with Torch-TensorRT
model = torch.compile(
    model,
    backend=backend,
    options={
        "truncate_long_and_double": True,
        "enabled_precisions": {torch.float16},
    },
    dynamic=False,
)

# Run inference
output = model(**encoded_input)
print(output)

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[peri044]

should this dtype be dependent on input dtype or always float32 ?

[chohk88]

Using np.float32 for the input works fine regardless of the input type.

However, if we trace the root cause of the error, we can see that in this line, the name truncate_double is used, but we are passing truncate_long_and_double as an argument to torch.compile, as shown here.

Because of this, at this point, the truncate_long_and_double argument is not handled and then removed, which leads to an error when trying to process the default type float64 of np.zeros.

According to this section, torch_tensorrt.dynamo.compile prefers truncate_double as the input but can also handle truncate_long_and_double.

What would be the best way to fix this issue?

[peri044]

Change the truncate_long_and_double to truncate_double in this example

TensorRT/examples/dynamo/torch_compile_stable_diffusion.py

Line 40 in 6d40ff1

"truncate_long_and_double": True,

and that should work right ?

[peri044]

truncate_long_and_double is deprecated. But it looks like it is not correctly handled if user provides this argument in torch.compile workflow. Can you add this check in

TensorRT/py/torch_tensorrt/dynamo/utils.py

Line 483 in 6d40ff1

valid_attrs = {attr.name for attr in fields(settings)}

? similar to https://github.com/pytorch/TensorRT/blob/main/py/torch_tensorrt/dynamo/_compiler.py#L180-L185