Skip to content

fix ts fe test error. #3619

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

Merged
merged 5 commits into from
Jun 26, 2025
Merged

fix ts fe test error. #3619

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
ae1c7ea
fix bug
lanluo-nvidia Jun 25, 2025
d0caab9
fix flashinfer issue
lanluo-nvidia Jun 25, 2025
ca84874
resolve comments
lanluo-nvidia Jun 26, 2025
6d5fcd1
test
lanluo-nvidia Jun 26, 2025
c902569
test
lanluo-nvidia Jun 26, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/scripts/install-torch-tensorrt.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ pip install --pre -r ${PWD}/tests/py/requirements.txt
# eg. timm will install the latest torchvision, however we want to use the torchvision from nightly
# reinstall torch torchvisionto make sure we have the correct version
pip uninstall -y torch torchvision
pip install --force-reinstall --pre ${TORCHVISION} --index-url ${INDEX_URL} --no-deps
pip install --force-reinstall --pre ${TORCHVISION} --index-url ${INDEX_URL}
pip install --force-reinstall --pre ${TORCH} --index-url ${INDEX_URL}


Expand Down
6 changes: 3 additions & 3 deletions MODULE.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ new_local_repository(
new_local_repository(
name = "cuda_win",
build_file = "@//third_party/cuda:BUILD",
path = "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v12.8/",
path = "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v12.9/",
)

http_archive = use_repo_rule("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
Expand All @@ -66,7 +66,7 @@ http_archive(
name = "libtorch",
build_file = "@//third_party/libtorch:BUILD",
strip_prefix = "libtorch",
urls = ["https://download.pytorch.org/libtorch/nightly/cu128/libtorch-shared-with-deps-latest.zip"],
urls = ["https://download.pytorch.org/libtorch/nightly/cu129/libtorch-shared-with-deps-latest.zip"],
)

# in aarch64 platform you can get libtorch via either local or wheel file
Expand All @@ -84,7 +84,7 @@ http_archive(
name = "libtorch_win",
build_file = "@//third_party/libtorch:BUILD",
strip_prefix = "libtorch",
urls = ["https://download.pytorch.org/libtorch/nightly/cu128/libtorch-win-shared-with-deps-latest.zip"],
urls = ["https://download.pytorch.org/libtorch/nightly/cu129/libtorch-win-shared-with-deps-latest.zip"],
)

http_archive(
Expand Down
1 change: 0 additions & 1 deletion setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -215,7 +215,6 @@ def build_libtorchtrt_cxx11_abi(
print("Jetpack build")

if IS_SBSA:
cmd.append("--platforms=//toolchains:sbsa")
print("SBSA build")

if CI_BUILD:
Expand Down
133 changes: 67 additions & 66 deletions tests/modules/hub.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,16 +1,14 @@
import importlib
import importlib.util
import json
import os

import custom_models as cm
import torch

if importlib.util.find_spec("torchvision"):
import timm
import torchvision.models as models

torch.hub._validate_not_a_forked_repo = lambda a, b, c: True


torch_version = torch.__version__

# Detect case of no GPU before deserialization of models on GPU
Expand All @@ -22,7 +20,71 @@
# Downloads all model files again if manifest file is not present
MANIFEST_FILE = "model_manifest.json"

to_test_models = {}
to_test_models = {
"pooling": {"model": cm.Pool(), "path": "trace"},
"module_fallback": {"model": cm.ModuleFallbackMain(), "path": "script"},
"loop_fallback_eval": {"model": cm.LoopFallbackEval(), "path": "script"},
"loop_fallback_no_eval": {"model": cm.LoopFallbackNoEval(), "path": "script"},
"conditional": {"model": cm.FallbackIf(), "path": "script"},
"inplace_op_if": {"model": cm.FallbackInplaceOPIf(), "path": "script"},
"standard_tensor_input": {"model": cm.StandardTensorInput(), "path": "script"},
"tuple_input": {"model": cm.TupleInput(), "path": "script"},
"list_input": {"model": cm.ListInput(), "path": "script"},
"tuple_input_output": {"model": cm.TupleInputOutput(), "path": "script"},
"list_input_output": {"model": cm.ListInputOutput(), "path": "script"},
"list_input_tuple_output": {
"model": cm.ListInputTupleOutput(),
"path": "script",
},
# "bert_base_uncased": {"model": cm.BertModule(), "path": "trace"},
}

if importlib.util.find_spec("torchvision"):
import timm
import torchvision.models as models

torchvision_models = {
"alexnet": {"model": models.alexnet(pretrained=True), "path": "both"},
"vgg16": {"model": models.vgg16(pretrained=True), "path": "both"},
"squeezenet": {"model": models.squeezenet1_0(pretrained=True), "path": "both"},
"densenet": {"model": models.densenet161(pretrained=True), "path": "both"},
"inception_v3": {"model": models.inception_v3(pretrained=True), "path": "both"},
"shufflenet": {
"model": models.shufflenet_v2_x1_0(pretrained=True),
"path": "both",
},
"mobilenet_v2": {"model": models.mobilenet_v2(pretrained=True), "path": "both"},
"resnext50_32x4d": {
"model": models.resnext50_32x4d(pretrained=True),
"path": "both",
},
"wideresnet50_2": {
"model": models.wide_resnet50_2(pretrained=True),
"path": "both",
},
"mnasnet": {"model": models.mnasnet1_0(pretrained=True), "path": "both"},
"resnet18": {
"model": torch.hub.load(
"pytorch/vision:v0.9.0", "resnet18", pretrained=True
),
"path": "both",
},
"resnet50": {
"model": torch.hub.load(
"pytorch/vision:v0.9.0", "resnet50", pretrained=True
),
"path": "both",
},
"efficientnet_b0": {
"model": timm.create_model("efficientnet_b0", pretrained=True),
"path": "script",
},
"vit": {
"model": timm.create_model("vit_base_patch16_224", pretrained=True),
"path": "script",
},
}
to_test_models.update(torchvision_models)


def get(n, m, manifest):
Expand Down Expand Up @@ -77,67 +139,6 @@ def download_models(version_matches, manifest):


def main():
if not importlib.util.find_spec("torchvision"):
print(f"torchvision is not installed, skip models download")
return

to_test_models = {
"alexnet": {"model": models.alexnet(pretrained=True), "path": "both"},
"vgg16": {"model": models.vgg16(pretrained=True), "path": "both"},
"squeezenet": {"model": models.squeezenet1_0(pretrained=True), "path": "both"},
"densenet": {"model": models.densenet161(pretrained=True), "path": "both"},
"inception_v3": {"model": models.inception_v3(pretrained=True), "path": "both"},
"shufflenet": {
"model": models.shufflenet_v2_x1_0(pretrained=True),
"path": "both",
},
"mobilenet_v2": {"model": models.mobilenet_v2(pretrained=True), "path": "both"},
"resnext50_32x4d": {
"model": models.resnext50_32x4d(pretrained=True),
"path": "both",
},
"wideresnet50_2": {
"model": models.wide_resnet50_2(pretrained=True),
"path": "both",
},
"mnasnet": {"model": models.mnasnet1_0(pretrained=True), "path": "both"},
"resnet18": {
"model": torch.hub.load(
"pytorch/vision:v0.9.0", "resnet18", pretrained=True
),
"path": "both",
},
"resnet50": {
"model": torch.hub.load(
"pytorch/vision:v0.9.0", "resnet50", pretrained=True
),
"path": "both",
},
"efficientnet_b0": {
"model": timm.create_model("efficientnet_b0", pretrained=True),
"path": "script",
},
"vit": {
"model": timm.create_model("vit_base_patch16_224", pretrained=True),
"path": "script",
},
"pooling": {"model": cm.Pool(), "path": "trace"},
"module_fallback": {"model": cm.ModuleFallbackMain(), "path": "script"},
"loop_fallback_eval": {"model": cm.LoopFallbackEval(), "path": "script"},
"loop_fallback_no_eval": {"model": cm.LoopFallbackNoEval(), "path": "script"},
"conditional": {"model": cm.FallbackIf(), "path": "script"},
"inplace_op_if": {"model": cm.FallbackInplaceOPIf(), "path": "script"},
"standard_tensor_input": {"model": cm.StandardTensorInput(), "path": "script"},
"tuple_input": {"model": cm.TupleInput(), "path": "script"},
"list_input": {"model": cm.ListInput(), "path": "script"},
"tuple_input_output": {"model": cm.TupleInputOutput(), "path": "script"},
"list_input_output": {"model": cm.ListInputOutput(), "path": "script"},
"list_input_tuple_output": {
"model": cm.ListInputTupleOutput(),
"path": "script",
},
# "bert_base_uncased": {"model": cm.BertModule(), "path": "trace"},
}

manifest = None
version_matches = False
Expand Down
9 changes: 6 additions & 3 deletions tests/py/dynamo/automatic_plugin/test_flashinfer_rmsnorm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,7 @@
import pytest

flashinfer = pytest.importorskip("flashinfer")
import importlib
import unittest

import pytest
import torch
import torch.nn as nn
import torch_tensorrt
Expand All @@ -12,6 +11,9 @@

from ..conversion.harness import DispatchTestCase

if importlib.util.find_spec("flashinfer"):
import flashinfer


@torch.library.custom_op("flashinfer::rmsnorm", mutates_args=()) # type: ignore[misc]
def flashinfer_rmsnorm(
Expand All @@ -31,6 +33,7 @@ def _(input: torch.Tensor, weight: torch.Tensor, b: float = 1e-6) -> torch.Tenso


@unittest.skip("Not Available")
@unittest.skipIf(not importlib.util.find_spec("flashinfer"), "flashinfer not installed")
class TestAutomaticPlugin(DispatchTestCase):
@parameterized.expand(
[
Expand Down
Loading