Skip to content

Add support for TensorRT-RTX #3753

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open

Add support for TensorRT-RTX #3753

wants to merge 5 commits into from
+2,784 −748

Conversation

lanluo-nvidia
Copy link
Collaborator

Description

Add initial support for TensorRT-RTX.
The following are the currently identified issues:
RTX team side:

  1. mobilenet/efficientnet test failed when input dtype is bfloat16(float16 and float32 is working)
    5439176
  2. Flux model FP4 and FP8 are not working (fp16 and fp32 is working)
    5400490
  3. NoneZero has not been supported yet in RTX(will add support in future)
    5407733
  4. Int8 quantization only support weights quantization, does not support activations quantization yet.
    5402295

Our side:

  1. a few test cases failed on strong typing
    PR is in Progress:
    fix: atan2 strong type support & bug fix for integer dynamic shape #3751
    add strong typing fix #3749
  2. batchnorm constant fold got nan (skip constant fold is able to get the result)
    🐛 [Bug] TensorRT-RTX BatchNorm constant fold got nan #3699
  3. refit test failed on missing eps weights due to we only capture (bias/weight/running_mean/running_var)
    🐛 [Bug] Refitter test failed when constant fold is disabled #3752

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • This change requires a documentation update

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

@github-actions github-actions bot added documentation Improvements or additions to documentation component: tests Issues re: Tests component: lowering Issues re: The lowering / preprocessing passes component: conversion Issues re: Conversion stage component: core Issues re: The core compiler component: converters Issues re: Specific op converters component: build system Issues re: Build system component: api [Python] Issues re: Python API component: api [C++] Issues re: C++ API labels Aug 6, 2025
@meta-cla meta-cla bot added the cla signed label Aug 6, 2025
@github-actions github-actions bot added component: evaluators Issues re: Specific op evaluators component: runtime component: partitioning component: fx component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Aug 6, 2025
@github-actions github-actions bot requested a review from narendasan August 6, 2025 18:49
@lanluo-nvidia lanluo-nvidia marked this pull request as ready for review August 6, 2025 18:50
@lanluo-nvidia lanluo-nvidia added the ciflow/binaries/all Build for all Python Versions label Aug 6, 2025
@pytorch-bot PyTorch Bot
Copy link

pytorch-bot bot commented Aug 6, 2025

No ciflow labels are configured for this repo.
For information on how to enable CIFlow bot see this wiki

narendasan
narendasan reviewed Aug 6, 2025
View reviewed changes
narendasan
narendasan reviewed Aug 6, 2025
View reviewed changes
Copy link
Collaborator

@narendasan narendasan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@lanluo-nvidia just remove the PTQ Calibrator feature from python and C++ and put in deprecation errors.

core/conversion/conversionctx/ConversionCtx.cpp
@@ -59,11 +59,16 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
for (auto p = settings.enabled_precisions.begin(); p != settings.enabled_precisions.end(); ++p) {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cant we condition this full block? since enabled precisions doesnt matter to RTX? Also the else case of the macro should be a warning that the setting enabled precision does not do anything for TRT RTX

core/conversion/conversion.cpp
@@ -204,7 +204,7 @@ void AddInputs(ConversionCtx* ctx, c10::ArrayRef<const torch::jit::Value*> input
"Optimization profile is invalid, please check the input range provided (conversion.AddInputs)");

ctx->cfg->addOptimizationProfile(profile);
#if NV_TENSORRT_MAJOR > 7 || (NV_TENSORRT_MAJOR == 7 && NV_TENSORRT_MINOR >= 1)
#ifndef TRT_MAJOR_RTX && (NV_TENSORRT_MAJOR > 7 || (NV_TENSORRT_MAJOR == 7 && NV_TENSORRT_MINOR >= 1))
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this work? like a ifndef + some other conditional?

cpp/bin/torchtrtc/main.cpp
@@ -334,8 +336,12 @@ int main(int argc, char** argv) {
if (calibration_cache_file) {
calibration_cache_file_path = torchtrtc::fileio::resolve_path(args::get(calibration_cache_file));
}

#ifndef TRT_MAJOR_RTX
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same thing here, the else cases should log a warning if the user tries to use the API. Or in this case maybe a error

py/torch_tensorrt/fx/converters/acc_ops_converters.py
@@ -12,6 +12,7 @@
import torch
from torch.fx.immutable_collections import immutable_list
from torch.fx.node import Argument, Target
from torch_tensorrt._utils import is_tensorrt_version_supported
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need these FX frontend updates?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@narendasan narendasan narendasan left review comments

Assignees
No one assigned
Labels
ciflow/binaries/all Build for all Python Versions cla signed component: api [C++] Issues re: C++ API component: api [Python] Issues re: Python API component: build system Issues re: Build system component: conversion Issues re: Conversion stage component: converters Issues re: Specific op converters component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: evaluators Issues re: Specific op evaluators component: fx component: lowering Issues re: The lowering / preprocessing passes component: partitioning component: runtime component: tests Issues re: Tests documentation Improvements or additions to documentation fx
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@lanluo-nvidia @narendasan @facebook-github-bot