alibaba · tanyokwok · Mar 17, 2023 · Mar 16, 2023 · Mar 16, 2023
diff --git a/.github/workflows/pytorch200_cpu.yml b/.github/workflows/pytorch200_cpu.yml
@@ -0,0 +1,27 @@
+name: pytorch-2.0.0-cpu
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+  schedule:
+    - cron: '0 17 * * *' # 1:00 am UTC+8:00
+  workflow_dispatch:
+
+jobs:
+  build:
+    uses: ./.github/workflows/resuable_cpu_build.yml
+    with:
+      remote_runtime_docker: bladedisc:latest-runtime-torch-2.0.0-cpu
+      develop_base_image: nvidia/cuda:11.7.1-cudnn8-devel-ubuntu20.04
+      runtime_base_image: nvidia/cuda:11.7.1-cudnn8-runtime-ubuntu20.04
+      extra_build_args: --build-arg PYTHON_VERSION=PYTHON3.8 --build-arg ENABLE_FIND_FASTEST_APT_SOURCE=OFF
+      extra_envs: -e TORCH_BLADE_BUILD_WITH_CUDA_SUPPORT=OFF
+        -e TORCH_BLADE_CI_BUILD_TORCH_VERSION=2.0.0+cpu
+      exec_command: bash ./scripts/ci/test_pytorch_blade.sh
+      deploy_command: TORCH_VERSION=2.0.0+cpu bash ./scripts/ci/deploy_pytorch_blade.sh
+    secrets:
+      ALIYUN_DOCKER_USERNAME: ${{ secrets.ALIYUN_DOCKER_USERNAME }}
+      ALIYUN_DOCKER_PASSWORD: ${{ secrets.ALIYUN_DOCKER_PASSWORD }}
+      DOCKER_USERNAME: ${{ secrets.DOCKER_USERNAME }}
+      DOCKER_PASSWORD: ${{ secrets.DOCKER_PASSWORD }}
diff --git a/.github/workflows/pytorch200_gpu.yml b/.github/workflows/pytorch200_gpu.yml
@@ -0,0 +1,28 @@
+name: pytorch_2.0.0-cuda11_7
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+  schedule:
+    - cron: '0 17 * * *' # 1:00 am UTC+8:00
+  workflow_dispatch:
+
+jobs:
+  build:
+    uses: ./.github/workflows/resuable_gpu_build.yml
+    with:
+      cuda_version: cu117
+      remote_runtime_docker: bladedisc:latest-runtime-torch-2.0.0-cu117
+      develop_base_image: nvidia/cuda:11.7.1-cudnn8-devel-ubuntu20.04
+      runtime_base_image: nvidia/cuda:11.7.1-cudnn8-devel-ubuntu20.04
+      extra_envs: -e TORCH_BLADE_BUILD_TENSORRT_STATIC=OFF
+          -e TORCH_BLADE_CI_BUILD_TORCH_VERSION=2.0.0+cu117
+      extra_build_args: --build-arg PYTHON_VERSION=PYTHON3.8
+      exec_command: bash ./scripts/ci/test_pytorch_blade.sh
+      deploy_command: TORCH_VERSION=2.0.0+cu117 bash ./scripts/ci/deploy_pytorch_blade.sh
+    secrets:
+      ALIYUN_DOCKER_USERNAME: ${{ secrets.ALIYUN_DOCKER_USERNAME }}
+      ALIYUN_DOCKER_PASSWORD: ${{ secrets.ALIYUN_DOCKER_PASSWORD }}
+      DOCKER_USERNAME: ${{ secrets.DOCKER_USERNAME }}
+      DOCKER_PASSWORD: ${{ secrets.DOCKER_PASSWORD }}
diff --git a/pytorch_blade/benchmark/TorchBench/requirements_cpu.txt b/pytorch_blade/benchmark/TorchBench/requirements_cpu.txt
@@ -1,8 +1,8 @@
 --pre --extra-index-url https://download.pytorch.org/whl/nightly/cpu
 pandas
 librosa
-torchaudio==2.0.0.dev20230130+cpu
-torchtext==0.15.0.dev20230130
+torchaudio==2.0.0.dev20230315+cpu
+torchtext==0.15.0.dev20230315
 onnx
 onnxruntime
 pyyaml
diff --git a/pytorch_blade/benchmark/TorchBench/requirements_cuda.txt b/pytorch_blade/benchmark/TorchBench/requirements_cuda.txt
@@ -1,6 +1,6 @@
 --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu117 --extra-index-url https://download.pytorch.org/whl/nightly/cpu
 pandas
 librosa
-torchaudio==2.0.0.dev20230130+cu117
-torchtext==0.15.0.dev20230130
+torchaudio==2.0.0.dev20230315+cu117
+torchtext==0.15.0.dev20230315
 pyyaml
diff --git a/pytorch_blade/pytorch_blade/compiler/jit/torch/shape_analysis.cpp b/pytorch_blade/pytorch_blade/compiler/jit/torch/shape_analysis.cpp
@@ -1632,11 +1632,15 @@ class ShapePropagator : public PropertyPropBase {
 #if PYTORCH_VERSION_GE(1, 14)
             "aten::sum.dim_IntList(Tensor self, int[]? dim, bool keepdim, *, int? dtype) -> Tensor",
             "aten::mean.dim(Tensor self, int[]? dim, bool keepdim, *, int? dtype) -> Tensor",
-            "aten::var.correction(Tensor self, int[1]? dim, *, int? correction, bool keepdim=False) -> Tensor",
             "aten::amax(Tensor self, int[1] dim=[], bool keepdim=False) -> Tensor",
 #else
             "aten::sum(Tensor self, int[] dim, bool keepdim, *, int? dtype) -> Tensor",
             "aten::mean(Tensor self, int[] dim, bool keepdim, *, int? dtype) -> Tensor",
+#endif
+#if PYTORCH_VERSION_GE(2, 1)
+            "aten::var.correction(Tensor self, int[1]? dim=None, *, Scalar? correction=None, bool keepdim=False) -> Tensor",
+#elif PYTORCH_VERSION_GE(1, 14)
+            "aten::var.correction(Tensor self, int[1]? dim, *, int? correction, bool keepdim=False) -> Tensor",
 #endif
         },
         [](Node* node) -> type_vec_t {
@@ -1647,7 +1651,14 @@ class ShapePropagator : public PropertyPropBase {
           }
 
           at::optional<IValue> opt_dtype;
-#if PYTORCH_VERSION_GE(1, 14)
+#if PYTORCH_VERSION_GE(2, 1)
+          if (!(node->matches(
+                    "aten::var.correction(Tensor self, int[1]? dim=None, *, Scalar? correction=None, bool keepdim=False) -> Tensor") ||
+                node->matches(
+                    "aten::amax(Tensor self, int[1] dim=[], bool keepdim=False) -> Tensor"))) {
+            opt_dtype = node->get(attr::dtype);
+          }
+#elif PYTORCH_VERSION_GE(1, 14)
           if (!(node->matches(
                     "aten::var.correction(Tensor self, int[1]? dim, *, int? correction, bool keepdim=False) -> Tensor") ||
                 node->matches(
@@ -1656,7 +1667,7 @@ class ShapePropagator : public PropertyPropBase {
           }
 #else
           opt_dtype = node->get(attr::dtype);
-#endif // PYTORCH_VERSION_GE(1, 14)
+#endif
           auto dims = constant_as<c10::List<int64_t>>(list);
           return multidim_reduce_with_keepdim(
               node,

diff --git a/pytorch_blade/scripts/pip/requirements-dev-2.0.0+cpu.txt b/pytorch_blade/scripts/pip/requirements-dev-2.0.0+cpu.txt
@@ -0,0 +1,11 @@
+wget
+pytest==5.0.0
+networkx
+onnx==1.12.0
+hypothesis
+expecttest
+py-cpuinfo
+aliyun-log-python-sdk==0.6.48.6
+cryptography
+torch==2.0.0+cpu
+torchvision==0.15.0+cpu
diff --git a/pytorch_blade/scripts/pip/requirements-dev-2.0.0+cu117.txt b/pytorch_blade/scripts/pip/requirements-dev-2.0.0+cu117.txt
@@ -0,0 +1,11 @@
+wget
+pytest==5.0.0
+networkx
+onnx==1.12.0
+hypothesis
+expecttest
+py-cpuinfo
+aliyun-log-python-sdk==0.6.48.6
+cryptography
+torch==2.0.0+cu117
+torchvision==0.15.0+cu117
diff --git a/pytorch_blade/scripts/pip/requirements-dev-pre+cpu.txt b/pytorch_blade/scripts/pip/requirements-dev-pre+cpu.txt
@@ -8,6 +8,6 @@ expecttest
 aliyun-log-python-sdk==0.6.48.6
 cryptography
 # please update pytorch_blade/benchmark/TorchBench/requirements_cpu.txt both
-torch==2.0.0.dev20230130+cpu
-torchvision==0.15.0.dev20230130+cpu
+torch==2.1.0.dev20230315+cpu
+torchvision==0.15.0.dev20230315+cpu
 --pre --extra-index-url https://download.pytorch.org/whl/nightly/cpu
diff --git a/pytorch_blade/scripts/pip/requirements-dev-pre+cu117.txt b/pytorch_blade/scripts/pip/requirements-dev-pre+cu117.txt
@@ -8,6 +8,6 @@ py-cpuinfo
 aliyun-log-python-sdk==0.6.48.6
 cryptography
 # please update pytorch_blade/benchmark/TorchBench/requirements_cuda.txt both
-torch==2.0.0.dev20230130+cu117
-torchvision==0.15.0.dev20230130+cu117
+torch==2.1.0.dev20230315+cu117
+torchvision==0.15.0.dev20230315+cu117
 --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu117
diff --git a/pytorch_blade/torch_blade/dynamo/__init__.py b/pytorch_blade/torch_blade/dynamo/__init__.py
@@ -11,9 +11,10 @@
 
 import torch_blade.dynamo.patch_user_defined
 
-from torch._dynamo.optimizations.training import aot_autograd
-from torch._dynamo.optimizations.backends import BACKENDS, create_backend
-from torch._dynamo.optimizations.subgraph import SubGraph
+from torch._dynamo.backends.common import aot_autograd
+from torch._dynamo.backends.registry import register_backend
+from torch._dynamo.utils import torchscript
+
 from torch._functorch import compilers
 from functorch.compile import min_cut_rematerialization_partition
 
@@ -90,16 +91,14 @@ def disc_compile(fx_g: fx.GraphModule, inps, use_ts=False) -> Callable:
     return _disc_compile(fx_g, inps, use_ts=False)
 
 def disc(fx_g: fx.GraphModule, inps) -> Callable:
-    import tempfile
-    with tempfile.TemporaryDirectory() as tmp:
-        scripted = SubGraph(fx_g, inps, tmp).scripted
-        torch._C._jit_pass_remove_mutation(scripted.graph)
-        f = torch.jit.freeze(scripted.eval())
-        cfg = torch_blade.Config()
-        cfg.disable_optimization_for_inference = False
-        with cfg:
-            f = torch_blade.optimize(f, True, tuple(inps))
-        return f
+    scripted = torchscript(fx_g, inps)
+    torch._C._jit_pass_remove_mutation(scripted.graph)
+    f = torch.jit.freeze(scripted.eval())
+    cfg = torch_blade.Config()
+    cfg.disable_optimization_for_inference = False
+    with cfg:
+        f = torch_blade.optimize(f, True, tuple(inps))
+    return f
 
 @compilers.make_boxed_compiler
 def disc_compile_ts(fx_g: fx.GraphModule, inps, use_ts=False) -> Callable:
@@ -212,6 +211,7 @@ def _get_disc_decomp():
     decompositions=_get_disc_decomp(),
     partition_fn=min_cut_rematerialization_partition)
 
-BACKENDS["disc"] = disc
-BACKENDS["aot_disc"] = aot_disc
-BACKENDS["aot_disc_debug"] = aot_disc_debug
+
+register_backend(name="disc", compiler_fn=disc)
+register_backend(name="aot_disc", compiler_fn=aot_disc)
+register_backend(name="aot_disc_debug", compiler_fn=aot_disc_debug)
diff --git a/pytorch_blade/torch_blade/testing/common_utils.py b/pytorch_blade/torch_blade/testing/common_utils.py
@@ -17,12 +17,17 @@
 from numbers import Number
 
 import torch
-from torch._six import string_classes
 from torch.testing import assert_allclose
 from torch.testing._internal.common_utils import TestCase as TorchTestCase
 from torch.testing._internal.common_utils import is_iterable
 from torch_blade import version
 
+try:
+    # used when PT 1.x
+    from torch._six import string_classes
+except:
+    string_classes = str
+
 __all__ = ['benchmark', 'assert_almost_equal', 'TestCase']