Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
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
6 changes: 3 additions & 3 deletions MODULE.bazel
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
133 changes: 67 additions & 66 deletions tests/modules/hub.py
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
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