Merge branch 'FDecaYed-deyuf/fused_optimizer_v2'

csrc/welford.cu

@@ -591,6 +591,58 @@ template <
     int PARALLEL_LOADS>
 __global__ void batchnorm_forward_c_last_kernel(
       const scalar_t* __restrict__ input,
+      const scalar_t* __restrict__ z,
+      const accscalar_t* __restrict__ mean,
+      const accscalar_t* __restrict__ inv_std,
+      const layerscalar_t* __restrict__ weight,
+      const layerscalar_t* __restrict__ shift,
+      scalar_t* __restrict__ out,
+      const int reduction_size,
+      const int stride,
+      const bool fuse_relu) {
+  // tensor dimension (m,c)
+  // loop along m dimension
+  int inner_loop_stride = blockDim.y * gridDim.y;
+
+  // offset along m dimension
+  int m_offset = blockIdx.y * blockDim.y + threadIdx.y;
+  int c_offset = blockIdx.x * blockDim.x + threadIdx.x;
+
+  auto m_c = mean[c_offset];
+  auto inv_std_c = static_cast<accscalar_t>(inv_std[c_offset]);
+  auto w_c = weight == NULL ? accscalar_t(1.0) : static_cast<accscalar_t>(weight[c_offset]);
+  auto s_c = shift == NULL ? accscalar_t(0.0) : static_cast<accscalar_t>(shift[c_offset]);
+
+  int loop_count = 1 + (reduction_size - 1) / (inner_loop_stride * PARALLEL_LOADS);
+  int address_base = m_offset * stride + c_offset;
+  int address_increment = inner_loop_stride * stride;
+
+  for (int i = 0; i < loop_count; i++) {
+#pragma unroll
+    for (int j = 0; j < PARALLEL_LOADS; j++) {
+      if (c_offset < stride && m_offset < reduction_size) {
+        auto tmp = w_c * (static_cast<accscalar_t>(input[address_base]) - m_c ) * inv_std_c + s_c;
+        if (z != NULL) {
+          tmp += z[address_base];
+        }
+        out[address_base] = (fuse_relu && tmp <= accscalar_t(0.0) ? scalar_t(0.0) : static_cast<scalar_t>(tmp));
+      }
+      m_offset += inner_loop_stride;
+      address_base += address_increment;
+    }
+  }
+}
+
+// elementwise BN kernel
+template <
+    typename scalar_t,
+    typename accscalar_t,
+    typename layerscalar_t,
+    int PARALLEL_LOADS>
+__global__ void relu_backward_c_last_kernel(
+      const scalar_t* __restrict__ grad_output,
+      const scalar_t* __restrict__ input,
+      const scalar_t* __restrict__ z,
       const accscalar_t* __restrict__ mean,
       const accscalar_t* __restrict__ inv_std,
       const layerscalar_t* __restrict__ weight,
@@ -619,9 +671,11 @@ __global__ void batchnorm_forward_c_last_kernel(
 #pragma unroll
     for (int j = 0; j < PARALLEL_LOADS; j++) {
       if (c_offset < stride && m_offset < reduction_size) {
-        out[address_base] = static_cast<scalar_t>(
-            w_c * (static_cast<accscalar_t>(input[address_base]) - m_c ) * inv_std_c + s_c
-          );
+        auto tmp = w_c * (static_cast<accscalar_t>(input[address_base]) - m_c ) * inv_std_c + s_c;
+        if (z != NULL) {
+          tmp += z[address_base];
+        }
+        out[address_base] = (tmp <= accscalar_t(0.0) ? scalar_t(0.0) : grad_output[address_base]);
       }
       m_offset += inner_loop_stride;
       address_base += address_increment;
@@ -1147,10 +1201,12 @@ std::vector<at::Tensor> welford_mean_var_c_last_CUDA(const at::Tensor input) {
 
 at::Tensor batchnorm_forward_c_last_CUDA(
     const at::Tensor input,
+    const at::optional<at::Tensor> z,
     const at::Tensor mean,
     const at::Tensor inv_std,
     const at::optional<at::Tensor> weight,
-    const at::optional<at::Tensor> shift) {
+    const at::optional<at::Tensor> shift,
+    const bool fuse_relu) {
   const auto stride = input.size(input.ndimension()-1);
   const auto reduction_size = input.numel() / stride;
 
@@ -1170,13 +1226,15 @@ at::Tensor batchnorm_forward_c_last_CUDA(
       batchnorm_forward_c_last_kernel<scalar_t_0, accscalar_t, accscalar_t, ELEMENTS_PER_ITER>
           <<<grid, block, 0, stream>>>(
           input.data<scalar_t_0>(),
+          z.has_value() ? z.value().data<scalar_t_0>() : NULL,
           mean.data<accscalar_t>(),
           inv_std.data<accscalar_t>(),
           weight.has_value() ? weight.value().data<accscalar_t>() : NULL,
           shift.has_value() ? shift.value().data<accscalar_t>(): NULL,
           out.data<scalar_t_0>(),
           reduction_size,
-          stride);
+          stride,
+          fuse_relu);
     );
   } else {
     if (weight.has_value()) {
@@ -1189,13 +1247,15 @@ at::Tensor batchnorm_forward_c_last_CUDA(
       batchnorm_forward_c_last_kernel<scalar_t_0, accscalar_t, scalar_t_0, ELEMENTS_PER_ITER>
           <<<grid, block, 0, stream>>>(
           input.data<scalar_t_0>(),
+          z.has_value() ? z.value().data<scalar_t_0>() : NULL,
           mean.data<accscalar_t>(),
           inv_std.data<accscalar_t>(),
           weight.has_value() ? weight.value().data<scalar_t_0>() : NULL,
           shift.has_value() ? shift.value().data<scalar_t_0>(): NULL,
           out.data<scalar_t_0>(),
           reduction_size,
-          stride);
+          stride,
+          fuse_relu);
     );
   }
   return out;
@@ -1351,3 +1411,66 @@ at::Tensor batchnorm_backward_c_last_CUDA(
  
   return grad_input;
 }
+
+at::Tensor relu_backward_c_last_CUDA(
+    const at::Tensor grad_output,
+    const at::Tensor input,
+    const at::optional<at::Tensor> z,
+    const at::Tensor mean,
+    const at::Tensor inv_std,
+    const at::optional<at::Tensor> weight,
+    const at::optional<at::Tensor> shift) {
+
+  const auto stride = input.size(input.ndimension()-1);
+  const auto reduction_size = input.numel() / stride;
+
+  at::Tensor out = at::empty_like(input);
+
+  dim3 block;
+  dim3 grid;
+  flexible_launch_configs(reduction_size, stride, block, grid);
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  if (input.scalar_type() == at::ScalarType::Half
+      && weight.has_value() && weight.value().scalar_type() == at::ScalarType::Float) {
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "batchnorm_forward",
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      relu_backward_c_last_kernel<scalar_t_0, accscalar_t, accscalar_t, ELEMENTS_PER_ITER>
+          <<<grid, block, 0, stream>>>(
+          grad_output.data<scalar_t_0>(),
+          input.data<scalar_t_0>(),
+          z.has_value() ? z.value().data<scalar_t_0>() : NULL,
+          mean.data<accscalar_t>(),
+          inv_std.data<accscalar_t>(),
+          weight.has_value() ? weight.value().data<accscalar_t>() : NULL,
+          shift.has_value() ? shift.value().data<accscalar_t>(): NULL,
+          out.data<scalar_t_0>(),
+          reduction_size,
+          stride);
+    );
+  } else {
+    if (weight.has_value()) {
+      TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),
+          "input.scalar_type() is not supported with weight.scalar_type()");
+    }
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "batchnorm_forward",
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      relu_backward_c_last_kernel<scalar_t_0, accscalar_t, scalar_t_0, ELEMENTS_PER_ITER>
+          <<<grid, block, 0, stream>>>(
+          grad_output.data<scalar_t_0>(),
+          input.data<scalar_t_0>(),
+          z.has_value() ? z.value().data<scalar_t_0>() : NULL,
+          mean.data<accscalar_t>(),
+          inv_std.data<accscalar_t>(),
+          weight.has_value() ? weight.value().data<scalar_t_0>() : NULL,
+          shift.has_value() ? shift.value().data<scalar_t_0>(): NULL,
+          out.data<scalar_t_0>(),
+          reduction_size,
+          stride);
+    );
+  }
+  return out;
+}

examples/imagenet/main_amp.py

@@ -25,54 +25,6 @@ try:
 except ImportError:
     raise ImportError("Please install apex from https://www.github.com/nvidia/apex to run this example.")
 
-model_names = sorted(name for name in models.__dict__
-                     if name.islower() and not name.startswith("__")
-                     and callable(models.__dict__[name]))
-
-parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
-parser.add_argument('data', metavar='DIR',
-                    help='path to dataset')
-parser.add_argument('--arch', '-a', metavar='ARCH', default='resnet18',
-                    choices=model_names,
-                    help='model architecture: ' +
-                    ' | '.join(model_names) +
-                    ' (default: resnet18)')
-parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
-                    help='number of data loading workers (default: 4)')
-parser.add_argument('--epochs', default=90, type=int, metavar='N',
-                    help='number of total epochs to run')
-parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
-                    help='manual epoch number (useful on restarts)')
-parser.add_argument('-b', '--batch-size', default=256, type=int,
-                    metavar='N', help='mini-batch size per process (default: 256)')
-parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
-                    metavar='LR', help='Initial learning rate.  Will be scaled by <global batch size>/256: args.lr = args.lr*float(args.batch_size*args.world_size)/256.  A warmup schedule will also be applied over the first 5 epochs.')
-parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
-                    help='momentum')
-parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
-                    metavar='W', help='weight decay (default: 1e-4)')
-parser.add_argument('--print-freq', '-p', default=10, type=int,
-                    metavar='N', help='print frequency (default: 10)')
-parser.add_argument('--resume', default='', type=str, metavar='PATH',
-                    help='path to latest checkpoint (default: none)')
-parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
-                    help='evaluate model on validation set')
-parser.add_argument('--pretrained', dest='pretrained', action='store_true',
-                    help='use pre-trained model')
-
-parser.add_argument('--prof', default=-1, type=int,
-                    help='Only run 10 iterations for profiling.')
-parser.add_argument('--deterministic', action='store_true')
-
-parser.add_argument("--local_rank", default=0, type=int)
-parser.add_argument('--sync_bn', action='store_true',
-                    help='enabling apex sync BN.')
-
-parser.add_argument('--opt-level', type=str)
-parser.add_argument('--keep-batchnorm-fp32', type=str, default=None)
-parser.add_argument('--loss-scale', type=str, default=None)
-
-cudnn.benchmark = True
 
 def fast_collate(batch):
     imgs = [img[0] for img in batch]
@@ -90,24 +42,75 @@ def fast_collate(batch):
         
     return tensor, targets
 
-best_prec1 = 0
-args = parser.parse_args()
 
-print("opt_level = {}".format(args.opt_level))
-print("keep_batchnorm_fp32 = {}".format(args.keep_batchnorm_fp32), type(args.keep_batchnorm_fp32))
-print("loss_scale = {}".format(args.loss_scale), type(args.loss_scale))
-
-print("\nCUDNN VERSION: {}\n".format(torch.backends.cudnn.version()))
+def parse():
+    model_names = sorted(name for name in models.__dict__
+                     if name.islower() and not name.startswith("__")
+                     and callable(models.__dict__[name]))
 
-if args.deterministic:
-    cudnn.benchmark = False
-    cudnn.deterministic = True
-    torch.manual_seed(args.local_rank)
-    torch.set_printoptions(precision=10)
+    parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
+    parser.add_argument('data', metavar='DIR',
+                        help='path to dataset')
+    parser.add_argument('--arch', '-a', metavar='ARCH', default='resnet18',
+                        choices=model_names,
+                        help='model architecture: ' +
+                        ' | '.join(model_names) +
+                        ' (default: resnet18)')
+    parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
+                        help='number of data loading workers (default: 4)')
+    parser.add_argument('--epochs', default=90, type=int, metavar='N',
+                        help='number of total epochs to run')
+    parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
+                        help='manual epoch number (useful on restarts)')
+    parser.add_argument('-b', '--batch-size', default=256, type=int,
+                        metavar='N', help='mini-batch size per process (default: 256)')
+    parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
+                        metavar='LR', help='Initial learning rate.  Will be scaled by <global batch size>/256: args.lr = args.lr*float(args.batch_size*args.world_size)/256.  A warmup schedule will also be applied over the first 5 epochs.')
+    parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
+                        help='momentum')
+    parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
+                        metavar='W', help='weight decay (default: 1e-4)')
+    parser.add_argument('--print-freq', '-p', default=10, type=int,
+                        metavar='N', help='print frequency (default: 10)')
+    parser.add_argument('--resume', default='', type=str, metavar='PATH',
+                        help='path to latest checkpoint (default: none)')
+    parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
+                        help='evaluate model on validation set')
+    parser.add_argument('--pretrained', dest='pretrained', action='store_true',
+                        help='use pre-trained model')
+
+    parser.add_argument('--prof', default=-1, type=int,
+                        help='Only run 10 iterations for profiling.')
+    parser.add_argument('--deterministic', action='store_true')
+
+    parser.add_argument("--local_rank", default=0, type=int)
+    parser.add_argument('--sync_bn', action='store_true',
+                        help='enabling apex sync BN.')
+
+    parser.add_argument('--opt-level', type=str)
+    parser.add_argument('--keep-batchnorm-fp32', type=str, default=None)
+    parser.add_argument('--loss-scale', type=str, default=None)
+    args = parser.parse_args()
+    return args
 
 def main():
     global best_prec1, args
 
+    args = parse()
+    print("opt_level = {}".format(args.opt_level))
+    print("keep_batchnorm_fp32 = {}".format(args.keep_batchnorm_fp32), type(args.keep_batchnorm_fp32))
+    print("loss_scale = {}".format(args.loss_scale), type(args.loss_scale))
+
+    print("\nCUDNN VERSION: {}\n".format(torch.backends.cudnn.version()))
+
+    cudnn.benchmark = True
+    best_prec1 = 0
+    if args.deterministic:
+        cudnn.benchmark = False
+        cudnn.deterministic = True
+        torch.manual_seed(args.local_rank)
+        torch.set_printoptions(precision=10)
+
     args.distributed = False
     if 'WORLD_SIZE' in os.environ:
         args.distributed = int(os.environ['WORLD_SIZE']) > 1

setup.py

@@ -85,11 +85,15 @@ if "--cuda_ext" in sys.argv:
         ext_modules.append(
             CUDAExtension(name='amp_C',
                           sources=['csrc/amp_C_frontend.cpp',
+                                   'csrc/multi_tensor_sgd_kernel.cu',
                                    'csrc/multi_tensor_scale_kernel.cu',
                                    'csrc/multi_tensor_axpby_kernel.cu',
                                    'csrc/multi_tensor_l2norm_kernel.cu',
                                    'csrc/multi_tensor_lamb_stage_1.cu',
-                                   'csrc/multi_tensor_lamb_stage_2.cu'],
+                                   'csrc/multi_tensor_lamb_stage_2.cu',
+                                   'csrc/multi_tensor_adam.cu',
+                                   'csrc/multi_tensor_novograd.cu',
+                                   'csrc/multi_tensor_lamb.cu'],
                           extra_compile_args={'cxx': ['-O3'],
                                               'nvcc':['-lineinfo',
                                                       '-O3',
@@ -115,6 +119,58 @@ if "--cuda_ext" in sys.argv:
                                                       '-O3',
                                                       '--use_fast_math'] + version_ge_1_1}))
 
+if "--bnp" in sys.argv:
+    from torch.utils.cpp_extension import CUDAExtension
+    sys.argv.remove("--bnp")
+
+    from torch.utils.cpp_extension import BuildExtension
+    cmdclass['build_ext'] = BuildExtension
+
+    if torch.utils.cpp_extension.CUDA_HOME is None:
+        raise RuntimeError("--bnp was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
+    else:
+        # Set up macros for forward/backward compatibility hack around
+        # https://github.com/pytorch/pytorch/commit/4404762d7dd955383acee92e6f06b48144a0742e
+        version_ge_1_1 = []
+        if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 0):
+            version_ge_1_1 = ['-DVERSION_GE_1_1']
+        ext_modules.append(
+            CUDAExtension(name='bnp',
+                          sources=['apex/contrib/csrc/groupbn/batch_norm.cu',
+                                   'apex/contrib/csrc/groupbn/ipc.cu',
+                                   'apex/contrib/csrc/groupbn/interface.cpp',
+                                   'apex/contrib/csrc/groupbn/batch_norm_add_relu.cu'],
+                          extra_compile_args={'cxx': [] + version_ge_1_1,
+                                              'nvcc':['-DCUDA_HAS_FP16=1',
+                                                      '-D__CUDA_NO_HALF_OPERATORS__',
+                                                      '-D__CUDA_NO_HALF_CONVERSIONS__',
+                                                      '-D__CUDA_NO_HALF2_OPERATORS__',
+                                                      '-gencode',
+                                                      'arch=compute_70,code=sm_70'] + version_ge_1_1}))
+
+if "--xentropy" in sys.argv:
+    from torch.utils.cpp_extension import CUDAExtension
+    sys.argv.remove("--xentropy")
+
+    from torch.utils.cpp_extension import BuildExtension
+    cmdclass['build_ext'] = BuildExtension
+
+    if torch.utils.cpp_extension.CUDA_HOME is None:
+        raise RuntimeError("--xentropy was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
+    else:
+        # Set up macros for forward/backward compatibility hack around
+        # https://github.com/pytorch/pytorch/commit/4404762d7dd955383acee92e6f06b48144a0742e
+        version_ge_1_1 = []
+        if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 0):
+            version_ge_1_1 = ['-DVERSION_GE_1_1']
+        ext_modules.append(
+            CUDAExtension(name='xentropy_cuda',
+                          sources=['apex/contrib/csrc/xentropy/interface.cpp',
+                                   'apex/contrib/csrc/xentropy/xentropy_kernel.cu'],
+                          include_dirs=['csrc'],
+                          extra_compile_args={'cxx': ['-O3'] + version_ge_1_1,
+                                              'nvcc':['-O3'] + version_ge_1_1}))
+
 setup(
     name='apex',
     version='0.1',

tests/L0/run_amp/test_basic_casts.py

@@ -137,26 +137,6 @@ class TestTensorCasts(unittest.TestCase):
         fn = lambda x: x.sum()
         run_layer_test(self, [fn], ALWAYS_FLOAT, (self.b, self.h))
 
-class TestDisabledCasts(unittest.TestCase):
-    def setUp(self):
-        self.handle = amp.init(enabled=False)
-        common_init(self)
-
-    def test_disabled_linear(self):
-        m = nn.Linear(self.h, self.h)
-        f = ft.partial(F.linear, weight=m.weight, bias=m.bias)
-        input_shape = (self.b, self.h)
-
-        for fn in [m, f]:
-            x = torch.randn(input_shape, dtype=torch.float).requires_grad_()
-            y = fn(x)
-            self.assertEqual(y.type(), FLOAT)
-            y.sum().backward()
-            self.assertEqual(x.grad.type(), FLOAT)
-
-            x = torch.randn(input_shape, dtype=torch.half).requires_grad_()
-            self.assertRaises(RuntimeError, fn, x)
-
     # TODO: maybe more tests on disabled casting?
 
 if __name__ == '__main__':

tests/L0/run_amp/test_fused_sgd.py

+import unittest
+
+import functools as ft
+import itertools as it
+
+from apex import amp
+from apex.amp import _amp_state
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn import Parameter
+
+from utils import common_init, HALF, FLOAT,\
+    ALWAYS_HALF, ALWAYS_FLOAT, MATCH_INPUT
+
+
+try:
+  import amp_C
+  disabled = False
+  from apex.optimizers import FusedSGD as FusedSGD
+except ImportError as err:
+  print("amp_C fused kernels unavailable, disabling TestMultiTensorApply.  ImportError was ", err)
+  disabled = True
+
+
+class MyModel(torch.nn.Module):
+    def __init__(self, unique):
+        super(MyModel, self).__init__()
+        self.weight0 = Parameter(unique +
+            torch.arange(2, device='cuda', dtype=torch.float32))
+        self.weight1 = Parameter(1. + unique + torch.arange(2, device='cuda', dtype=torch.float16))
+
+    @staticmethod
+    def ops(input, weight0, weight1):
+        return ((input*(weight0.float()))*(weight1.float())).sum()
+
+    def forward(self, input):
+        return self.ops(input, self.weight0, self.weight1)
+
+# Abandon all hope, ye who enter here.
+
+# This is hands down the ugliest code I have ever written, but it succeeds in testing
+# multiple models/optimizers/losses fairly thoroughly.  Many of the different test cases
+# require slightly divergent code in a way that seems near-impossible to genericize into a simple
+# cross product or nested loops.
+
+class TestMultipleModelsOptimizersLosses(unittest.TestCase):
+    def setUp(self):
+        self.x = torch.ones((2), device='cuda', dtype=torch.float32)
+        common_init(self)
+
+    def tearDown(self):
+        pass
+
+    @unittest.skipIf(disabled, "amp_C is unavailable")
+    def test_2models2losses1optimizer(self):
+        model0 = MyModel(1)
+        model1 = MyModel(2)
+
+        optimizer = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                     {'params' : model1.parameters(), 'lr' : 0.5}],
+                                    momentum=0.125)
+
+        reference_grads = []
+        for i in range(2):
+            optimizer.zero_grad()
+            loss0 = model0(self.x)
+            loss1 = model1(self.x)
+            loss0.backward()
+            loss1.backward()
+
+            reference_grads.append([param.grad.data.clone() for param in model0.parameters()] +
+                                   [param.grad.data.clone() for param in model1.parameters()])
+
+            optimizer.step()
+
+        final_params = [param.data.clone() for param in model0.parameters()] + \
+                       [param.data.clone() for param in model1.parameters()]
+
+        for materialize_master_grads in (False, True):
+          for opt_level in ("O0", "O1", "O2", "O3"):
+            for how_to_zero in ("none", "model", "optimizer"):
+              for use_multiple_loss_scalers in (False, True):
+                if opt_level == "O1" or opt_level == "O2":
+                    inject_inf_iters = (-1, 0, 1)
+                else:
+                    inject_inf_iters = (-1,)
+
+                for inject_inf in inject_inf_iters:
+                  if inject_inf >= 0:
+                     inject_inf_locs = ("fp16", "fp32")
+                     which_backwards = (0, 1)
+                  else:
+                     inject_inf_locs = ("fdsa",)
+                     which_backwards = (None,)
+
+                  for inject_inf_loc in inject_inf_locs:
+                    for which_backward in which_backwards:
+                        if use_multiple_loss_scalers:
+                            num_losses = 2
+                            loss_ids = [0, 1]
+                        else:
+                            num_losses = 1
+                            loss_ids = [0, 0]
+
+                        if inject_inf >= 0:
+                            iters = 3
+                        else:
+                            iters = 2
+
+                        model0 = MyModel(1)
+                        model1 = MyModel(2)
+
+                        models = [model0, model1]
+
+                        optimizer = FusedSGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                              {'params' : model1.parameters(), 'lr' : 0.5}],
+                                             momentum=0.125,
+                                             materialize_master_grads=materialize_master_grads)
+
+                        _amp_state.allow_incoming_model_not_fp32 = True
+                        [model0, model1], optimizer = amp.initialize(
+                            [model0, model1],
+                            optimizer,
+                            opt_level=opt_level,
+                            verbosity=0,
+                            cast_model_type=False,
+                            num_losses=num_losses)
+                        _amp_state.allow_incoming_model_not_fp32 = False
+
+                        _amp_state.loss_scalers[0]._loss_scale = 4.0
+                        if use_multiple_loss_scalers:
+                            _amp_state.loss_scalers[1]._loss_scale = 16.0
+
+                        unskipped = 0
+                        for i in range(iters):
+                            if how_to_zero == "none":
+                                for model in models:
+                                    for param in model.parameters():
+                                        param.grad = None
+                            elif how_to_zero == "model":
+                                for model in models:
+                                    model.zero_grad()
+                            else:
+                                optimizer.zero_grad()
+
+                            loss0 = model0(self.x)
+                            loss1 = model1(self.x)
+
+                            with amp.scale_loss(loss0, optimizer, loss_id=loss_ids[0]) as scaled_loss:
+                                scaled_loss.backward()
+                                if i == inject_inf and which_backward == 0:
+                                    if inject_inf_loc == "fp32":
+                                        model0.weight0.grad[0] = float('inf')
+                                    elif inject_inf_loc == "fp16":
+                                        model0.weight1.grad[0] = float('inf')
+                            with amp.scale_loss(loss1, optimizer, loss_id=loss_ids[1]) as scaled_loss:
+                                scaled_loss.backward()
+                                if i == inject_inf and which_backward == 1:
+                                    if inject_inf_loc == "fp32":
+                                        model1.weight0.grad[0] = float('inf')
+                                    elif inject_inf_loc == "fp16":
+                                        model1.weight1.grad[0] = float('inf')
+
+                            if i != inject_inf:
+                                master_params = amp.master_params(optimizer)
+                                for param, reference_grad in zip(master_params, reference_grads[unskipped]):
+                                    if opt_level == "O2" and not materialize_master_grads:
+                                        continue
+                                    else:
+                                        self.assertTrue(torch.allclose(param.grad.float(), reference_grad.float()),
+                                                        "opt_level {} i {} inject_inf {} which_backward {} inject_inf_loc {} use_multiple_loss_scalers {}".format(opt_level, i, inject_inf, which_backward, inject_inf_loc, use_multiple_loss_scalers))
+                                unskipped += 1
+                            optimizer.step()
+
+                        model_params = [p for p in model0.parameters()] + [p for p in model1.parameters()]
+                        for model, master, reference in zip(
+                                model_params,
+                                amp.master_params(optimizer),
+                                final_params):
+                            self.assertTrue(torch.allclose(model, reference))
+                            self.assertTrue(torch.allclose(model, master.to(model.dtype)))
+
+                        if opt_level == "O1":
+                            _amp_state.handle._deactivate()
+
+    @unittest.skipIf(disabled, "amp_C is unavailable")
+    def test_3models2losses1optimizer(self):
+
+        model0 = MyModel(1)
+        model1 = MyModel(2)
+        model2 = MyModel(3)
+
+        optimizer = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                     {'params' : model1.parameters(), 'lr' : 0.5},
+                                     {'params' : model2.parameters(), 'lr' : 0.125}],
+                                     momentum=0.125)
+
+        reference_grads = []
+        for i in range(2):
+            optimizer.zero_grad()
+            loss0 = model0(self.x) + model2(self.x)
+            loss1 = model1(self.x) + model2(self.x)
+            loss0.backward()
+            loss1.backward()
+
+            reference_grads.append([param.grad.data.clone() for param in model0.parameters()] +
+                                   [param.grad.data.clone() for param in model1.parameters()] +
+                                   [param.grad.data.clone() for param in model2.parameters()])
+
+            optimizer.step()
+
+
+        final_params = [param.data.clone() for param in model0.parameters()] + \
+                       [param.data.clone() for param in model1.parameters()] + \
+                       [param.data.clone() for param in model2.parameters()]
+
+        for materialize_master_grads in (False, True):
+          for opt_level in ("O0", "O1", "O2", "O3"):
+            for how_to_zero in ("none", "model", "optimizer"):
+              for use_multiple_loss_scalers in (False, True):
+                if opt_level == "O1" or opt_level == "O2":
+                    inject_inf_iters = (-1, 0, 1)
+                else:
+                    inject_inf_iters = (-1,)
+
+                for inject_inf in inject_inf_iters:
+                  if inject_inf >= 0:
+                     inject_inf_locs = ("fp16", "fp32")
+                     which_backwards = (0, 1)
+                  else:
+                     inject_inf_locs = ("fdsa",)
+                     which_backwards = (None,)
+
+                  for inject_inf_loc in inject_inf_locs:
+                    for which_backward in which_backwards:
+                      if use_multiple_loss_scalers:
+                          num_losses = 2
+                          loss_ids = [0, 1]
+                      else:
+                          num_losses = 1
+                          loss_ids = [0, 0]
+
+                      if inject_inf >= 0:
+                          iters = 3
+                          if which_backward == 0:
+                              which_models = (0, 2)
+                          elif which_backward == 1:
+                              which_models = (1, 2)
+                      else:
+                          iters = 2
+                          which_models = (None,)
+
+                      for which_model in which_models:
+                          model0 = MyModel(1)
+                          model1 = MyModel(2)
+                          model2 = MyModel(3)
+
+                          models = [model0, model1, model2]
+
+                          optimizer = FusedSGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                               {'params' : model1.parameters(), 'lr' : 0.5},
+                                               {'params' : model2.parameters(), 'lr' : 0.125}],
+                                               momentum=0.125,
+                                               materialize_master_grads=materialize_master_grads)
+
+                          _amp_state.allow_incoming_model_not_fp32 = True
+                          [model0, model1, model2], optimizer = amp.initialize(
+                              [model0, model1, model2],
+                              optimizer,
+                              opt_level=opt_level,
+                              verbosity=0,
+                              cast_model_type=False,
+                              num_losses=num_losses)
+                          _amp_state.allow_incoming_model_not_fp32 = False
+
+                          _amp_state.loss_scalers[0]._loss_scale = 4.0
+                          if use_multiple_loss_scalers:
+                              _amp_state.loss_scalers[1]._loss_scale = 16.0
+
+                          unskipped = 0
+                          for i in range(iters):
+                              if how_to_zero == "none":
+                                  for model in models:
+                                      for param in model.parameters():
+                                          param.grad = None
+                              elif how_to_zero == "model":
+                                  for model in models:
+                                      model.zero_grad()
+                              else:
+                                  optimizer.zero_grad()
+
+                              loss0 = model0(self.x) + model2(self.x)
+                              loss1 = model1(self.x) + model2(self.x)
+
+                              with amp.scale_loss(loss0, optimizer, loss_id=loss_ids[0]) as scaled_loss:
+                                  scaled_loss.backward()
+                                  if i == inject_inf and which_backward == 0:
+                                      if which_model == 0:
+                                          inj_model = model0
+                                      elif which_model == 2:
+                                          inj_model = model2
+                                      else:
+                                          raise RuntimeError(which_model + " invalid for loss 0")
+                                      if inject_inf_loc == "fp32":
+                                          inj_model.weight0.grad[0] = float('inf')
+                                      elif inject_inf_loc == "fp16":
+                                          inj_model.weight1.grad[0] = float('inf')
+                              with amp.scale_loss(loss1, optimizer, loss_id=loss_ids[1]) as scaled_loss:
+                                  scaled_loss.backward()
+                                  if i == inject_inf and which_backward == 1:
+                                      if which_model == 1:
+                                          inj_model = model1
+                                      elif which_model == 2:
+                                          inj_model = model2
+                                      else:
+                                          raise RuntimeError(which_model + " invalid for loss 1 ")
+                                      if inject_inf_loc == "fp32":
+                                          inj_model.weight0.grad[0] = float('inf')
+                                      elif inject_inf_loc == "fp16":
+                                          inj_model.weight1.grad[0] = float('inf')
+
+                              if i != inject_inf:
+                                  master_params = amp.master_params(optimizer)
+                                  for param, reference_grad in zip(master_params, reference_grads[unskipped]):
+                                      if opt_level == "O2" and not materialize_master_grads:
+                                          continue
+                                      else:
+                                          self.assertTrue(torch.allclose(param.grad.float(), reference_grad.float()),
+                                          "opt_level {} i {} inject_inf {} which_backward {} inject_inf_loc {} which_model {} use_multiple_loss_scalers {}".format(opt_level, i, inject_inf, which_backward, inject_inf_loc, which_model, use_multiple_loss_scalers))
+                                  unskipped += 1
+
+                              optimizer.step()
+
+                          model_params = [p for p in model0.parameters()] + \
+                                         [p for p in model1.parameters()] + \
+                                         [p for p in model2.parameters()]
+                          for model, master, reference in zip(
+                                  model_params,
+                                  amp.master_params(optimizer),
+                                  final_params):
+                              self.assertTrue(torch.allclose(model, reference))
+                              self.assertTrue(torch.allclose(model, master.to(model.dtype)))
+
+                          if opt_level == "O1":
+                              _amp_state.handle._deactivate()
+
+    @unittest.skipIf(disabled, "amp_C is unavailable")
+    def test_2models2losses2optimizers(self):
+        model0 = MyModel(1)
+        model1 = MyModel(2)
+
+        optimizer0 = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25}],
+                                      momentum=0.125)
+        optimizer1 = torch.optim.SGD([{'params' : model1.parameters(), 'lr' : 0.5}],
+                                      momentum=0.25)
+
+        # Don't do it like this:  reference_grads = [[]]*5
+        # because then it creates a list of 5 references to the same "[]" and appending
+        # to any of them effectively makes you append to all of them, which multiplies
+        # the resulting size of reference_grads by 5x and needless to say makes the test fail.
+        reference_grads = [[], [], [], [], []]
+        final_params = [None, None, None, None, None]
+        for i in range(2):
+            optimizer0.zero_grad()
+            optimizer1.zero_grad()
+            loss0 = model0(self.x)
+            loss1 = model1(self.x)
+            loss0.backward()
+            loss1.backward()
+
+            reference_grads[0].append([param.grad.data.clone() for param in model0.parameters()] +
+                                   [param.grad.data.clone() for param in model1.parameters()])
+
+            optimizer0.step()
+            optimizer1.step()
+
+        final_params[0] = [param.data.clone() for param in model0.parameters()] + \
+                          [param.data.clone() for param in model1.parameters()]
+
+        def what_got_skipped(which_iter, which_backward):
+            if which_iter == 0 and which_backward == 0:
+                return 1
+            if which_iter == 0 and which_backward == 1:
+                return 2
+            if which_iter == 1 and which_backward == 0:
+                return 3
+            if which_iter == 1 and which_backward == 1:
+                return 4
+            return 0
+
+        for which_iter in (0,1):
+            for which_backward in (0,1):
+                model0 = MyModel(1)
+                model1 = MyModel(2)
+
+                optimizer0 = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25}],
+                                              momentum=0.125)
+                optimizer1 = torch.optim.SGD([{'params' : model1.parameters(), 'lr' : 0.5}],
+                                              momentum=0.25)
+
+                for i in range(3):
+                    optimizer0.zero_grad()
+                    optimizer1.zero_grad()
+                    loss0 = model0(self.x)
+                    loss1 = model1(self.x)
+                    loss0.backward()
+                    loss1.backward()
+
+                    if i != which_iter:
+                        reference_grads[what_got_skipped(which_iter, which_backward)].append(
+                            [param.grad.data.clone() for param in model0.parameters()] +
+                            [param.grad.data.clone() for param in model1.parameters()])
+
+                    if i == which_iter:
+                        if which_backward == 0:
+                            optimizer1.step()
+                        else:
+                            optimizer0.step()
+                    else:
+                        optimizer0.step()
+                        optimizer1.step()
+
+                final_params[what_got_skipped(which_iter, which_backward)] = \
+                    [param.data.clone() for param in model0.parameters()] + \
+                    [param.data.clone() for param in model1.parameters()]
+
+        for materialize_master_grads in (False, True):
+          for opt_level in ("O0", "O1", "O2", "O3"):
+            for how_to_zero in ("none", "model", "optimizer"):
+              for use_multiple_loss_scalers in (False, True):
+                if opt_level == "O1" or opt_level == "O2":
+                    inject_inf_iters = (-1, 0, 1)
+                else:
+                    inject_inf_iters = (-1,)
+
+                for inject_inf in inject_inf_iters:
+                  if inject_inf >= 0:
+                     inject_inf_locs = ("fp16", "fp32")
+                     which_backwards = (0, 1)
+                  else:
+                     inject_inf_locs = ("fdsa",)
+                     which_backwards = (None,)
+
+                  for inject_inf_loc in inject_inf_locs:
+                    for which_backward in which_backwards:
+                        if use_multiple_loss_scalers:
+                            num_losses = 2
+                            loss_ids = [0, 1]
+                        else:
+                            num_losses = 1
+                            loss_ids = [0, 0]
+
+                        if inject_inf >= 0:
+                            iters = 3
+                        else:
+                            iters = 2
+
+                        model0 = MyModel(1)
+                        model1 = MyModel(2)
+
+                        models = [model0, model1]
+
+                        optimizer0 = FusedSGD([{'params' : model0.parameters(), 'lr' : 0.25}],
+                                              momentum=0.125, materialize_master_grads=materialize_master_grads)
+                        optimizer1 = FusedSGD([{'params' : model1.parameters(), 'lr' : 0.5}],
+                                              momentum=0.25, materialize_master_grads=materialize_master_grads)
+
+                        _amp_state.allow_incoming_model_not_fp32 = True
+                        [model0, model1], [optimizer0, optimizer1] = amp.initialize(
+                            [model0, model1],
+                            [optimizer0, optimizer1],
+                            opt_level=opt_level,
+                            verbosity=0,
+                            cast_model_type=False,
+                            num_losses=num_losses)
+                        _amp_state.allow_incoming_model_not_fp32 = False
+
+                        _amp_state.loss_scalers[0]._loss_scale = 4.0
+                        if use_multiple_loss_scalers:
+                            _amp_state.loss_scalers[1]._loss_scale = 16.0
+
+                        unskipped = 0
+                        for i in range(iters):
+                            if how_to_zero == "none":
+                                for model in models:
+                                    for param in model.parameters():
+                                        param.grad = None
+                            elif how_to_zero == "model":
+                                for model in models:
+                                    model.zero_grad()
+                            else:
+                                optimizer0.zero_grad()
+                                optimizer1.zero_grad()
+
+                            loss0 = model0(self.x)
+                            loss1 = model1(self.x)
+
+                            with amp.scale_loss(loss0, optimizer0, loss_id=loss_ids[0]) as scaled_loss:
+                                scaled_loss.backward()
+                                if i == inject_inf and which_backward == 0:
+                                    if inject_inf_loc == "fp32":
+                                        model0.weight0.grad[0] = float('inf')
+                                    elif inject_inf_loc == "fp16":
+                                        model0.weight1.grad[0] = float('inf')
+                            with amp.scale_loss(loss1, optimizer1, loss_id=loss_ids[1]) as scaled_loss:
+                                scaled_loss.backward()
+                                if i == inject_inf and which_backward == 1:
+                                    if inject_inf_loc == "fp32":
+                                        model1.weight0.grad[0] = float('inf')
+                                    elif inject_inf_loc == "fp16":
+                                        model1.weight1.grad[0] = float('inf')
+
+                            # print("opt_level {} i {} inject_inf {} which_backward {} inject_inf_loc {} use_multiple_loss_scalers {}".format(opt_level, i, inject_inf, which_backward, inject_inf_loc, use_multiple_loss_scalers))
+
+                            if i != inject_inf:
+                                master_params = list(amp.master_params(optimizer0)) + \
+                                                list(amp.master_params(optimizer1))
+                                for param, reference_grad in zip(master_params,
+                                        reference_grads[what_got_skipped(inject_inf, which_backward)][unskipped]):
+                                    if opt_level == "O2" and not materialize_master_grads:
+                                        continue
+                                    else:
+                                        self.assertTrue(torch.allclose(param.grad.float(), reference_grad.float()))
+                                unskipped += 1
+
+                            optimizer0.step()
+                            optimizer1.step()
+
+                        model_params = [p for p in model0.parameters()] + [p for p in model1.parameters()]
+                        master_params = [p for p in amp.master_params(optimizer0)] + \
+                                        [p for p in amp.master_params(optimizer1)]
+                        for model, master, reference in zip(
+                                model_params,
+                                master_params,
+                                final_params[what_got_skipped(inject_inf, which_backward)]):
+                            self.assertTrue(torch.allclose(model, reference))
+                            self.assertTrue(torch.allclose(model, master.to(model.dtype)))
+
+                        if opt_level == "O1":
+                            _amp_state.handle._deactivate()
+
+    @unittest.skipIf(disabled, "amp_C is unavailable")
+    def test_3models2losses2optimizers(self):
+        model0 = MyModel(1)
+        model1 = MyModel(2)
+        model2 = MyModel(3)
+
+        optimizer0 = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                      {'params' : model1.parameters(), 'lr' : 1.0}],
+                                     momentum=0.5)
+        optimizer1 = torch.optim.SGD([{'params' : model2.parameters(), 'lr' : 0.5}],
+                                     momentum=0.25)
+
+        # Again, can't do this:  reference_grads = [[]]*9
+        reference_grads = [[], [], [], [], [], [], [], [], []]
+        final_params = [None, None, None, None, None, None, None, None, None]
+        for i in range(2):
+            optimizer0.zero_grad()
+            optimizer1.zero_grad()
+            loss0 = model0(self.x) + model1(self.x)
+            loss1 = model2(self.x) + model1(self.x)
+            loss0.backward()
+            loss1.backward()
+
+            reference_grads[0].append([param.grad.data.clone() for param in model0.parameters()] +
+                                   [param.grad.data.clone() for param in model1.parameters()])
+
+            optimizer0.step()
+            optimizer1.step()
+
+        final_params[0] = \
+            [param.data.clone() for param in model0.parameters()] + \
+            [param.data.clone() for param in model1.parameters()] + \
+            [param.data.clone() for param in model2.parameters()]
+
+        def what_got_skipped(which_iter, which_backward, which_model):
+            if which_iter == 0:
+                if which_backward == 0:
+                    if which_model == 0:
+                        return 1
+                    if which_model == 1:
+                        return 2
+                if which_backward == 1:
+                    if which_model == 2:
+                        return 3
+                    if which_model == 1:
+                        return 4
+            if which_iter == 1:
+                if which_backward == 0:
+                    if which_model == 0:
+                        return 5
+                    if which_model == 1:
+                        return 6
+                if which_backward == 1:
+                    if which_model == 2:
+                        return 7
+                    if which_model == 1:
+                        return 8
+            return 0
+
+        for which_iter in (0,1):
+            for which_backward in (0,1):
+                if which_backward == 0:
+                    which_models = (0,1)
+                if which_backward == 1:
+                    which_models = (2,1)
+                for which_model in which_models:
+
+                    model0 = MyModel(1)
+                    model1 = MyModel(2)
+                    model2 = MyModel(3)
+
+                    optimizer0 = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                                  {'params' : model1.parameters(), 'lr' : 1.0}],
+                                                 momentum=0.5)
+                    optimizer1 = torch.optim.SGD([{'params' : model2.parameters(), 'lr' : 0.5}],
+                                                 momentum=0.25)
+
+                    for i in range(3):
+                        optimizer0.zero_grad()
+                        optimizer1.zero_grad()
+                        loss0 = model0(self.x) + model1(self.x)
+                        loss1 = model2(self.x) + model1(self.x)
+                        loss0.backward()
+                        loss1.backward()
+
+                        if i != which_iter:
+                            reference_grads[what_got_skipped(which_iter,
+                                    which_backward, which_model)].append(
+                                [param.grad.data.clone() for param in model0.parameters()] +
+                                [param.grad.data.clone() for param in model1.parameters()])
+
+                        if i == which_iter:
+                            if which_backward == 0:
+                                # if which_model == 0:
+                                    optimizer1.step()
+                                # if which_model == 1:
+                                #     optimizer1.step()
+                            if which_backward == 1:
+                                # if which_model == 2:
+                                #     optimizer0.step()
+                                # if which_model == 1:
+                                    continue
+                        else:
+                            optimizer0.step()
+                            optimizer1.step()
+
+                    final_params[what_got_skipped(which_iter, which_backward, which_model)] = \
+                        [param.data.clone() for param in model0.parameters()] + \
+                        [param.data.clone() for param in model1.parameters()] + \
+                        [param.data.clone() for param in model2.parameters()]
+
+        for materialize_master_grads in (False, True):
+          for opt_level in ("O0", "O1", "O2", "O3"):
+            for how_to_zero in ("none", "model", "optimizer"):
+              for use_multiple_loss_scalers in (False, True):
+                if opt_level == "O1" or opt_level == "O2":
+                    inject_inf_iters = (-1, 0, 1)
+                else:
+                    inject_inf_iters = (-1,)
+
+                for inject_inf in inject_inf_iters:
+                  if inject_inf >= 0:
+                     inject_inf_locs = ("fp16", "fp32")
+                     which_backwards = (0, 1)
+                  else:
+                     inject_inf_locs = ("fdsa",)
+                     which_backwards = (None,)
+
+                  for inject_inf_loc in inject_inf_locs:
+                    for which_backward in which_backwards:
+                      if use_multiple_loss_scalers:
+                          num_losses = 2
+                          loss_ids = [0, 1]
+                      else:
+                          num_losses = 1
+                          loss_ids = [0, 0]
+
+                      if inject_inf >= 0:
+                          iters = 3
+                          if which_backward == 0:
+                              which_models = (0, 1)
+                          elif which_backward == 1:
+                              which_models = (2, 1)
+                      else:
+                          iters = 2
+                          which_models = (None,)
+
+                      for which_model in which_models:
+                          model0 = MyModel(1)
+                          model1 = MyModel(2)
+                          model2 = MyModel(3)
+
+                          models = [model0, model1, model2]
+
+                          optimizer0 = FusedSGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                            {'params' : model1.parameters(), 'lr' : 1.0}],
+                                            momentum=0.5, materialize_master_grads=materialize_master_grads)
+                          optimizer1 = FusedSGD([{'params' : model2.parameters(), 'lr' : 0.5}],
+                                                momentum=0.25, materialize_master_grads=materialize_master_grads)
+
+                          _amp_state.allow_incoming_model_not_fp32 = True
+                          [model0, model1, model2], [optimizer0, optimizer1] = amp.initialize(
+                              [model0, model1, model2],
+                              [optimizer0, optimizer1],
+                              opt_level=opt_level,
+                              verbosity=0,
+                              cast_model_type=False,
+                              num_losses=num_losses)
+                          _amp_state.allow_incoming_model_not_fp32 = False
+
+                          _amp_state.loss_scalers[0]._loss_scale = 4.0
+                          if use_multiple_loss_scalers:
+                              _amp_state.loss_scalers[1]._loss_scale = 16.0
+
+                          unskipped = 0
+                          for i in range(iters):
+                              if how_to_zero == "none":
+                                  for model in models:
+                                      for param in model.parameters():
+                                          param.grad = None
+                              elif how_to_zero == "model":
+                                  for model in models:
+                                      model.zero_grad()
+                              else:
+                                  optimizer0.zero_grad()
+                                  optimizer1.zero_grad()
+
+                              loss0 = model0(self.x) + model1(self.x)
+                              loss1 = model2(self.x) + model1(self.x)
+
+                              with amp.scale_loss(loss0, optimizer0, loss_id=loss_ids[0]) as scaled_loss:
+                                  scaled_loss.backward()
+                                  if i == inject_inf and which_backward == 0:
+                                      if which_model == 0:
+                                          inj_model = model0
+                                      elif which_model == 1:
+                                          inj_model = model1
+                                      else:
+                                          raise RuntimeError(which_model + " invalid for loss 0")
+                                      if inject_inf_loc == "fp32":
+                                          inj_model.weight0.grad[0] = float('inf')
+                                      elif inject_inf_loc == "fp16":
+                                          inj_model.weight1.grad[0] = float('inf')
+                              with amp.scale_loss(loss1, [optimizer0, optimizer1], loss_id=loss_ids[1]) as scaled_loss:
+                                  scaled_loss.backward()
+                                  if i == inject_inf and which_backward == 1:
+                                      if which_model == 2:
+                                          inj_model = model2
+                                      elif which_model == 1:
+                                          inj_model = model1
+                                      else:
+                                          raise RuntimeError(which_model + " invalid for loss 1 ")
+                                      if inject_inf_loc == "fp32":
+                                          inj_model.weight0.grad[0] = float('inf')
+                                      elif inject_inf_loc == "fp16":
+                                          inj_model.weight1.grad[0] = float('inf')
+
+                              if i != inject_inf:
+                                  master_params = list(amp.master_params(optimizer0)) + \
+                                                  list(amp.master_params(optimizer1))
+                                  for param, reference_grad in zip(master_params,
+                                        reference_grads[what_got_skipped(inject_inf,
+                                            which_backward, which_model)][unskipped]):
+                                      if opt_level == "O2" and not materialize_master_grads:
+                                          continue
+                                      else:
+                                          self.assertTrue(torch.allclose(param.grad.float(), reference_grad.float()))
+                                  unskipped += 1
+
+                              optimizer0.step()
+                              optimizer1.step()
+
+                          model_params = [p for p in model0.parameters()] + \
+                                         [p for p in model1.parameters()] + \
+                                         [p for p in model2.parameters()]
+                          master_params = [p for p in amp.master_params(optimizer0)] + \
+                                          [p for p in amp.master_params(optimizer1)]
+
+                          # print("opt_level {} i {} inject_inf {} which_backward {} inject_inf_loc {} use_multiple_loss_scalers {} which_model {}".format(opt_level, i, inject_inf, which_backward, inject_inf_loc, use_multiple_loss_scalers, which_model))
+
+                          for model, master, reference in zip(
+                                  model_params,
+                                  master_params,
+                                  final_params[what_got_skipped(inject_inf, which_backward, which_model)]):
+                              self.assertTrue(torch.allclose(model, reference))
+                              self.assertTrue(torch.allclose(model, master.to(model.dtype)))
+
+                          if opt_level == "O1":
+                              _amp_state.handle._deactivate()
+
+if __name__ == '__main__':
+    unittest.main()

tests/L0/run_mixed_adam/test_fp16_optimizer.py

-import unittest
-import torch
-import apex
-
-class TestFP16Optimizer(unittest.TestCase):
-    def setUp(self, max_abs_diff=1e-3, max_rel_diff=1, iters=7):
-        self.max_abs_diff = max_abs_diff
-        self.max_rel_diff = max_rel_diff
-        self.iters = iters
-        torch.cuda.manual_seed(13337)
-
-        N, D_in, D_out = 64, 1024, 16
-        self.N = N
-        self.D_in = D_in
-        self.D_out = D_out
-        self.x = torch.randn((N, D_in), dtype=torch.float16, device='cuda')
-        self.ref_model = torch.nn.Linear(D_in, D_out).cuda().half()
-        self.tst_model = torch.nn.Linear(D_in, D_out).cuda().half()
-        for p,q in zip(self.tst_model.parameters(), self.ref_model.parameters()):
-            p.data.copy_(q.data)
-
-    def get_max_diff(self, ref_param, tst_param):
-        max_abs_diff = max_rel_diff = 0
-        for p_ref, p_tst in zip(ref_param, tst_param):
-            max_abs_diff_p = (p_ref - p_tst).abs().max().item()
-            max_rel_diff_p = ((p_ref - p_tst) / p_ref).abs().max().item()
-
-            if max_abs_diff_p > max_abs_diff:  max_abs_diff = max_abs_diff_p
-            if max_rel_diff_p > max_rel_diff:  max_rel_diff = max_rel_diff_p
-
-        return max_abs_diff, max_rel_diff
-
-    def test_fp16_optimizer(self):
-
-        ref_optim = torch.optim.Adam(self.ref_model.parameters())
-        ref_optim = apex.fp16_utils.FP16_Optimizer(ref_optim, verbose=False)
-
-        tst_optim = apex.optimizers.FusedAdam(self.tst_model.parameters())
-        tst_optim = apex.optimizers.FP16_Optimizer(tst_optim)
-
-        for i in range(self.iters):
-            ref_loss = self.ref_model(self.x).sum()
-            ref_optim.backward(ref_loss)
-            ref_optim.step()
-
-            tst_loss = self.tst_model(self.x).sum()
-            tst_optim.backward(tst_loss)
-            tst_optim.step()
-
-            max_abs_diff, max_rel_diff = self.get_max_diff(self.ref_model.parameters(), self.tst_model.parameters())
-            self.assertLessEqual(max_abs_diff, self.max_abs_diff)
-            self.assertLessEqual(max_rel_diff, self.max_rel_diff)
-
-
-    def test_loss_scaling(self):
-
-        ref_optim = torch.optim.Adam(self.ref_model.parameters())
-        ref_optim = apex.fp16_utils.FP16_Optimizer(ref_optim, static_loss_scale=128.0, verbose=False)
-
-        tst_optim = apex.optimizers.FusedAdam(self.tst_model.parameters())
-        tst_optim = apex.optimizers.FP16_Optimizer(tst_optim, static_loss_scale=128.0)
-
-        for i in range(self.iters):
-            ref_loss = self.ref_model(self.x).sum()
-            ref_optim.backward(ref_loss)
-            ref_optim.step()
-
-            tst_loss = self.tst_model(self.x).sum()
-            tst_optim.backward(tst_loss)
-            tst_optim.step()
-
-            max_abs_diff, max_rel_diff = self.get_max_diff(self.ref_model.parameters(), self.tst_model.parameters())
-            self.assertLessEqual(max_abs_diff, self.max_abs_diff)
-            self.assertLessEqual(max_rel_diff, self.max_rel_diff)
-
-    def test_parameter_groups(self):
-
-        ref_groups = [{'params': [self.ref_model.weight]},{'params': [self.ref_model.bias]}]
-        ref_optim = torch.optim.Adam(ref_groups)
-        ref_optim = apex.fp16_utils.FP16_Optimizer(ref_optim, verbose=False)
-
-        tst_groups = [{'params': [self.tst_model.weight]},{'params': [self.tst_model.bias]}]
-        tst_optim = apex.optimizers.FusedAdam(tst_groups)
-        tst_optim = apex.optimizers.FP16_Optimizer(tst_optim)
-
-        for i in range(self.iters):
-            ref_loss = self.ref_model(self.x).sum()
-            ref_optim.backward(ref_loss)
-            ref_optim.step()
-
-            tst_loss = self.tst_model(self.x).sum()
-            tst_optim.backward(tst_loss)
-            tst_optim.step()
-
-            max_abs_diff, max_rel_diff = self.get_max_diff(self.ref_model.parameters(), self.tst_model.parameters())
-            self.assertLessEqual(max_abs_diff, self.max_abs_diff)
-            self.assertLessEqual(max_rel_diff, self.max_rel_diff)
-
-    def test_grad_clip(self):
-        ref_optim = torch.optim.Adam(self.ref_model.parameters())
-        ref_optim = apex.fp16_utils.FP16_Optimizer(ref_optim, verbose=False)
-
-        tst_optim = apex.optimizers.FusedAdam(self.tst_model.parameters(), max_grad_norm=0.01)
-        tst_optim = apex.optimizers.FP16_Optimizer(tst_optim)
-
-        for i in range(self.iters):
-            ref_loss = self.ref_model(self.x).sum()
-            ref_optim.backward(ref_loss)
-            ref_optim.clip_master_grads(0.01)
-            ref_optim.step()
-
-            tst_loss = self.tst_model(self.x).sum()
-            tst_optim.backward(tst_loss)
-            tst_optim.step()
-
-            max_abs_diff, max_rel_diff = self.get_max_diff(self.ref_model.parameters(), self.tst_model.parameters())
-            self.assertLessEqual(max_abs_diff, self.max_abs_diff)
-            self.assertLessEqual(max_rel_diff, self.max_rel_diff)
-
-    @unittest.skip('Not support grad being None')
-    def test_grad_None(self):
-        self.fail()
-
-    @unittest.skip('Not support same weight decay as pytorch')
-    def test_weight_decay(self):
-        self.fail()
-
-    @unittest.skip('Not support empty parameter groups')
-    def test_group_empty(self):
-        self.fail()
-
-if __name__ == '__main__':
-    script_path = os.path.dirname(os.path.realpath(__file__))
-    unittest.main()

tests/L0/run_mixed_adam/test_mixed_adam.py → tests/L0/run_optimizers/test_adam.py

@@ -24,7 +24,7 @@ class TestFusedAdam(unittest.TestCase):
 
         ref_optim = torch.optim.Adam(ref_param, **adam_option)
         tst_optim = apex.optimizers.FusedAdam(tst_param, **adam_option)
-       
+
         return (ref_param, tst_param, ref_optim, tst_optim)
 
     def gen_grad(self, ref_param, tst_param):
@@ -75,23 +75,9 @@ class TestFusedAdam(unittest.TestCase):
         self.gen_single_type_test(param_type=torch.float)
 
     def test_half(self):
-        nelem = 278011
-        adam_option = {'lr':5e-4, 'betas':(0.9, 0.999), 'eps':1e-08,
-            'weight_decay':0, 'amsgrad':False}
-
-        tensor = torch.rand(nelem, dtype=torch.float, device='cuda')
-        ref_param, tst_param, ref_optim, tst_optim = \
-            self.gen_param_optim([tensor], adam_option)
-
-        for i in range(self.iters):
-            half_grads = self.gen_mixed_grad(ref_param, tst_param)
-            ref_optim.step()
-            tst_optim.step(grads=half_grads)
-            max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param)
-
-            self.assertLessEqual(max_abs_diff, self.max_abs_diff)
-            self.assertLessEqual(max_rel_diff, self.max_rel_diff)
+        self.gen_single_type_test(param_type=torch.float16)
 
+    @unittest.skip('Disable until 8/1/2019 adam/adamw upstream picked')
     def test_multi_params(self):
         sizes = [[4096, 1024], [4096], [4096, 2048], [32320, 1024], [1]]
         adam_option = {'lr':5e-4, 'betas':(0.9, 0.999), 'eps':1e-08,
@@ -104,13 +90,14 @@ class TestFusedAdam(unittest.TestCase):
             self.gen_param_optim(tensors, adam_option)
 
         for i in range(self.iters):
-            half_grads = self.gen_mixed_grad(ref_param, tst_param)
+            self.gen_grad(ref_param, tst_param)
             ref_optim.step()
-            tst_optim.step(grads=half_grads)
+            tst_optim.step()
             max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param)
             self.assertLessEqual(max_abs_diff, self.max_abs_diff)
             self.assertLessEqual(max_rel_diff, self.max_rel_diff)
 
+    @unittest.skip('No longer support fuse scaling')
     def test_scale(self):
         nelem = 278011
         adam_option = {'lr':5e-4, 'betas':(0.9, 0.999), 'eps':1e-08,
@@ -130,6 +117,7 @@ class TestFusedAdam(unittest.TestCase):
             self.assertLessEqual(max_abs_diff, self.max_abs_diff)
             self.assertLessEqual(max_rel_diff, self.max_rel_diff)
 
+    @unittest.skip('No longer support output fp16 param')
     def test_fp16_output(self):
         nelem = 278011
         adam_option = {'lr':5e-4, 'betas':(0.9, 0.999), 'eps':1e-08,

tests/L0/run_test.py

 import unittest
 import sys
 
-test_dirs = ["run_amp", "run_fp16util", "run_mixed_adam", "run_fused_layer_norm", "run_pyprof_nvtx"]
+test_dirs = ["run_amp", "run_fp16util", "run_optimizers", "run_fused_layer_norm", "run_pyprof_nvtx"]
 
 runner = unittest.TextTestRunner(verbosity=2)
 

tests/L1/cross_product/run.sh

 #!/bin/bash
 
-DATADIR="/home/mcarilli/Desktop/pt18data/apex_stale/examples/imagenet/bare_metal_train_val/"
+# DATADIR="/home/mcarilli/Desktop/pt18data/apex_stale/examples/imagenet/bare_metal_train_val/"
 # DATADIR="/opt/home/apex/examples/imagenet/"
 cp ../common/* .
-bash run_test.sh single_gpu $1 $DATADIR yes
+bash run_test.sh single_gpu $1

tests/distributed/DDP/ddp_race_condition_test.py

@@ -35,8 +35,9 @@ class Model(Module):
 
 model = Model()
 # model = DDP(model, message_size=1, gradient_predivide_factor=8.0)
-model = DDP(model, delay_allreduce=True)
+# model = DDP(model, delay_allreduce=True)
 # model = DDP(model, message_size=1, allreduce_trigger_params=[model.b])
+model = DDP(model, message_size=1, allreduce_trigger_params=[model.b], num_allreduce_streams=3)
 
 x = torch.cuda.FloatTensor(4096*4096)