patching nvfuser conv cudnn test numerics mismatch (#2048)

Tests failed on upstream, not yet in our devel branch. Disabling TF32 in the test, which creates numerical issue when validating outputs.
csarofeen · Oct 11, 2022 · 7117a7e · 7117a7e
1 parent 65af1a4
commit 7117a7e
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/torch/csrc/jit/codegen/cuda/test/test_gpu_shift.cpp b/torch/csrc/jit/codegen/cuda/test/test_gpu_shift.cpp
@@ -2976,6 +2976,7 @@ TEST_F(NVFuserTest, FusionConv2D_CUDA) {
 TEST_F(NVFuserTest, FusionConv2DNoPadding_CUDA) {
   Fusion fusion;
   FusionGuard fg(&fusion);
+  ContextCudnnTF32Disabled disabling_tf32_cudnn;
 
   // Input: [C, H, W]
   auto inp = makeSymbolicTensor(3);

diff --git a/torch/csrc/jit/codegen/cuda/test/test_utils.h b/torch/csrc/jit/codegen/cuda/test/test_utils.h
@@ -8,6 +8,7 @@
 #include <torch/csrc/jit/codegen/cuda/lower_magic_zero.h>
 #include <torch/csrc/jit/codegen/cuda/transform_replay.h>
 
+#include <ATen/Context.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDACachingAllocator.h>
 #include <torch/torch.h>
@@ -340,6 +341,21 @@ struct TransformPropagatorWithCheck : public TransformPropagator {
 
 } // namespace
 
+class ContextCudnnTF32Disabled {
+ public:
+  ContextCudnnTF32Disabled() {
+    flag_ = at::globalContext().allowTF32CuDNN();
+    at::globalContext().setAllowTF32CuDNN(false);
+  }
+
+  ~ContextCudnnTF32Disabled() {
+    at::globalContext().setAllowTF32CuDNN(flag_);
+  }
+
+ private:
+  bool flag_;
+};
+
 // Fixture class must be uniquely identified, i.e., can't be in an
 // anonymous namespace
 class NVFuserTest : public ::testing::Test {