Merge pull request #851 from kexinyu/master

make FusedLAMB async

apex/optimizers/fused_lamb.py

@@ -132,7 +132,7 @@ class FusedLAMB(torch.optim.Optimizer):
         global_grad_norm = multi_tensor_applier(self.multi_tensor_l2norm,
                                                 self._dummy_overflow_buf,
                                                 [[g_norm_32, g_norm_16]],
-                                                False)[0].item()
+                                                False)[0]
         max_grad_norm = self.defaults['max_grad_norm']
 
         for group in self.param_groups:

csrc/amp_C_frontend.cpp

@@ -42,7 +42,7 @@ void multi_tensor_lamb_stage1_cuda(
     const float beta1,
     const float beta2,
     const float epsilon,
-    const float global_grad_norm,
+    at::Tensor global_grad_norm,
     const float max_global_grad_norm);
 
 void multi_tensor_lamb_stage2_cuda(
@@ -108,7 +108,7 @@ void multi_tensor_lamb_cuda(
   const float weight_decay,
   const int grad_averaging,
   const int mode,
-  const float global_grad_norm,
+  at::Tensor global_grad_norm,
   const float max_grad_norm,
   at::optional<bool> use_nvlamb_python);
 

csrc/multi_tensor_lamb.cu

@@ -52,7 +52,7 @@ struct LAMBStage1Functor
     const float epsilon,
     adamMode_t mode,
     const float decay,
-    const float global_grad_norm,
+    const float* global_grad_norm,
     const float max_global_grad_norm)
   {
     // I'd like this kernel to propagate infs/nans.
@@ -63,7 +63,7 @@ struct LAMBStage1Functor
     int chunk_idx = tl.block_to_chunk[blockIdx.x];
     int n = tl.sizes[tensor_loc];
 
-    float clipped_global_grad_norm = global_grad_norm > max_global_grad_norm ? global_grad_norm / max_global_grad_norm : 1.0f;
+    float clipped_global_grad_norm = (*global_grad_norm) > max_global_grad_norm ? (*global_grad_norm) / max_global_grad_norm : 1.0f;
 
     T* g = (T*)tl.addresses[0][tensor_loc];
     g += chunk_idx*chunk_size;
@@ -342,7 +342,7 @@ void multi_tensor_lamb_cuda(
   const float weight_decay,
   const int grad_averaging,
   const int mode,
-  const float global_grad_norm,
+  at::Tensor global_grad_norm,
   const float max_grad_norm,
   at::optional<bool> use_nvlamb_python)
 {
@@ -387,7 +387,7 @@ void multi_tensor_lamb_cuda(
         epsilon,
         (adamMode_t) mode,
         weight_decay,
-        global_grad_norm,
+        global_grad_norm.DATA_PTR<float>(),
         max_grad_norm); )
 
   // Compute update norms

csrc/multi_tensor_lamb_stage_1.cu

@@ -118,12 +118,13 @@ void multi_tensor_lamb_stage1_cuda(
   const float beta1,
   const float beta2,
   const float epsilon,
-  const float global_grad_norm,
+  at::Tensor global_grad_norm,
   const float max_global_grad_norm)
 {
   using namespace at;
 
-  float clipped_global_grad_norm = global_grad_norm > max_global_grad_norm ? global_grad_norm / max_global_grad_norm : 1.0f;
+  const float* g_grad_norm = global_grad_norm.DATA_PTR<float>();
+  float clipped_global_grad_norm = *(g_grad_norm) > max_global_grad_norm ? *(g_grad_norm) / max_global_grad_norm : 1.0f;
   float next_step = float(step+1);
   float beta1_correction = 1.0f - std::pow(beta1, next_step);
   float beta2_correction = 1.0f - std::pow(beta2, next_step);