How to enable fp16 with torch-tensorrt backend?

[vedantroy]

When trying to use fp16, I get the following error:

[2024-06-17 17:25:08,351][torch_tensorrt [TensorRT Conversion Context]][ERROR] - 4: [network.cpp::validate::3399] Error Code 4: Internal Error (fp16 precision has been set for a layer or layer output, but fp16 is not configured in the builder)

I'm configuring this as follows:

    with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.float16, enabled=True):
        model = torch.compile(model, fullgraph=True, backend="tensorrt")
        outputs = model(inputs)

[peri044]

You should use enabled_precisions as follows

model = torch.compile(model, fullgraph=True, 
                                       backend="tensorrt", 
                                       options={
                                         "enabled_precisions": {torch.float16},
                                       },)

[IbrokhimN]

autocast isn't really doing anything here, it just controls precision in eager pytorch mode and that info never reaches the tensorrt builder. by default the builder is set up for fp32, so when the graph ends up with fp16 layers from autocast it complains.

you need to explicitly pass enabled_precisions in options:

import torch
import torch_tensorrt

with torch.no_grad():
    model = torch.compile(
        model,
        backend="tensorrt",
        options={"enabled_precisions": {torch.float16}},
        fullgraph=True,
    )
    outputs = model(inputs)

you can drop autocast entirely, not needed. if you want a mix of fp32/fp16 just put both in the set:

options={"enabled_precisions": {torch.float32, torch.float16}}

also worth noting, if your input tensors are already fp16, that's separate from enabled_precisions, the latter just lets tensorrt use fp16 kernels internally, doesn't set the input dtype.