Merge pull request NVIDIA#11 from lcskrishna/cl/fused-optimizers-bfp16

[FusedOptimizers] Bug fixes in fused optimizers for fp16/bfp16.
lcskrishna · May 22, 2020 · 5cfdc01 · 5cfdc01
2 parents bdd481d + 9297be6
commit 5cfdc01
Show file tree

Hide file tree

Showing 3 changed files with 1 addition and 9 deletions.
diff --git a/csrc/multi_tensor_sgd_kernel.cu b/csrc/multi_tensor_sgd_kernel.cu
@@ -271,7 +271,7 @@ void multi_tensor_sgd_cuda(
         scale);
   }
   // Case 5. bfp16, bfp16, bfp16, No
-  if(grad_type == at::ScalarType::BFloat16 &&
+  else if(grad_type == at::ScalarType::BFloat16 &&
      weight_type == at::ScalarType::BFloat16 &&
      num_tensors == 3)
   {

diff --git a/tests/L0/run_amp/test_fused_sgd.py b/tests/L0/run_amp/test_fused_sgd.py
@@ -54,7 +54,6 @@ def tearDown(self):
         pass
 
     @unittest.skipIf(disabled, "amp_C is unavailable")
-    @skipIfRocm
     def test_2models2losses1optimizer(self):
         model0 = MyModel(1)
         model1 = MyModel(2)
@@ -187,7 +186,6 @@ def test_2models2losses1optimizer(self):
                             _amp_state.handle._deactivate()
 
     @unittest.skipIf(disabled, "amp_C is unavailable")
-    @skipIfRocm
     def test_3models2losses1optimizer(self):
 
         model0 = MyModel(1)
@@ -349,7 +347,6 @@ def test_3models2losses1optimizer(self):
                               _amp_state.handle._deactivate()
 
     @unittest.skipIf(disabled, "amp_C is unavailable")
-    @skipIfRocm
     def test_2models2losses2optimizers(self):
         model0 = MyModel(1)
         model1 = MyModel(2)
@@ -545,7 +542,6 @@ def what_got_skipped(which_iter, which_backward):
                             _amp_state.handle._deactivate()
 
     @unittest.skipIf(disabled, "amp_C is unavailable")
-    @skipIfRocm
     def test_3models2losses2optimizers(self):
         model0 = MyModel(1)
         model1 = MyModel(2)

diff --git a/tests/L0/run_amp/test_multiple_models_optimizers_losses.py b/tests/L0/run_amp/test_multiple_models_optimizers_losses.py
@@ -44,7 +44,6 @@ def setUp(self):
     def tearDown(self):
         pass
 
-    @skipIfRocm
     def test_2models2losses1optimizer(self):
         model0 = MyModel(1)
         model1 = MyModel(2)
@@ -170,7 +169,6 @@ def test_2models2losses1optimizer(self):
                       if opt_level == "O1":
                           _amp_state.handle._deactivate()
 
-    @skipIfRocm
     def test_3models2losses1optimizer(self):
 
         model0 = MyModel(1)
@@ -327,7 +325,6 @@ def test_3models2losses1optimizer(self):
                         if opt_level == "O1":
                             _amp_state.handle._deactivate()
 
-    @skipIfRocm
     def test_2models2losses2optimizers(self):
         model0 = MyModel(1)
         model1 = MyModel(2)
@@ -518,7 +515,6 @@ def what_got_skipped(which_iter, which_backward):
                       if opt_level == "O1":
                           _amp_state.handle._deactivate()
 
-    @skipIfRocm
     def test_3models2losses2optimizers(self):
         model0 = MyModel(1)
         model1 = MyModel(2)