Remove Apex Dependency For Deepspeed FP16_Optimizer (#12077)

* remove apex dependency

* fix amd build

orttraining/orttraining/python/training/optim/_ds_modifier.py

@@ -37,7 +37,7 @@ class DeepSpeedZeROModifier(FP16OptimizerModifier):
 
         return self.check_requirements(
             ["has_overflow_serial", "get_grad_norm_direct", "has_overflow_partitioned_grads_serial"],
-            require_apex=True,
+            require_apex=False,
             require_torch_non_finite_check=True,
         )
 
@@ -45,7 +45,7 @@ class DeepSpeedZeROModifier(FP16OptimizerModifier):
         warnings.warn("DeepSpeed fp16_optimizer functions are overrided with faster implementation.", UserWarning)
 
         def get_grad_norm_direct(target, gradients, params, norm_type=2):
-            import amp_C
+            from onnxruntime.training.ortmodule.torch_cpp_extensions import fused_ops
 
             def is_model_parallel_parameter(p):
                 return hasattr(p, "model_parallel") and p.model_parallel
@@ -93,7 +93,10 @@ class DeepSpeedZeROModifier(FP16OptimizerModifier):
                     # Multi-tensor applier takes a function and a list of list
                     # and performs the operation on that list all in one kernel.
                     grad_norm, _ = multi_tensor_applier(
-                        amp_C.multi_tensor_l2norm, dummy_overflow_buf, [grads_for_norm], False  # no per-parameter norm
+                        fused_ops.multi_tensor_l2norm,
+                        dummy_overflow_buf,
+                        [fused_ops.TorchTensorVector(grads_for_norm)],
+                        False,  # no per-parameter norm
                     )
                     # Since we will be summing across data parallel groups,
                     # we need the pow(norm-type).

orttraining/orttraining/python/training/optim/_megatron_modifier.py

@@ -11,7 +11,7 @@
 
 import types
 import warnings
-from numpy import inf
+
 from ._modifier import FP16OptimizerModifier, check_overflow, clip_grad_norm_fp32
 
 
@@ -23,7 +23,7 @@ class LegacyMegatronLMModifier(FP16OptimizerModifier):
 
     def can_be_modified(self):
         return self.check_requirements(
-            ["_check_overflow", "clip_master_grads"], require_apex=True, require_torch_non_finite_check=True
+            ["_check_overflow", "clip_master_grads"], require_apex=False, require_torch_non_finite_check=True
         )
 
     def override_function(self):

orttraining/orttraining/python/training/optim/_modifier.py

@@ -8,9 +8,11 @@
 # - check_overflow_for_grads : https://github.com/NVIDIA/Megatron-LM/blob/5ac5571ba0265af4c491ee0af1508ca7589450c6/megatron/optimizer/optimizer.py#L341
 # --------------------------------------------------------------------------
 
-import torch
 import warnings
+
+import torch
 from numpy import inf
+
 from ._multi_tensor_apply import MultiTensorApply
 
 multi_tensor_applier = MultiTensorApply(2048 * 32)
@@ -66,7 +68,7 @@ def check_overflow_for_grads(grad_data):
 def clip_grad_norm_fp32(
     parameters, max_norm, norm_type, get_horizontal_model_parallel_rank=None, get_horizontal_model_parallel_group=None
 ):
-    import amp_C
+    from onnxruntime.training.ortmodule.torch_cpp_extensions import fused_ops
 
     horizontal_model_parallel_grad_norm_aggregation = False
     if get_horizontal_model_parallel_rank and get_horizontal_model_parallel_group:
@@ -120,7 +122,10 @@ def clip_grad_norm_fp32(
             # Multi-tensor applier takes a function and a list of list
             # and performs the operation on that list all in one kernel.
             grad_norm, _ = multi_tensor_applier(
-                amp_C.multi_tensor_l2norm, dummy_overflow_buf, [grads_for_norm], False  # no per-parameter norm
+                fused_ops.multi_tensor_l2norm,
+                dummy_overflow_buf,
+                [fused_ops.TorchTensorVector(grads_for_norm)],
+                False,  # no per-parameter norm
             )
 
             if not horizontal_model_parallel_grad_norm_aggregation:
@@ -145,6 +150,7 @@ def clip_grad_norm_fp32(
         # Filter parameters with gradients.
         grads = [p.grad for p in parameters if p.grad is not None]
         if clip_coef < 1.0:
-            multi_tensor_applier(amp_C.multi_tensor_scale, dummy_overflow_buf, [grads, grads], clip_coef)
+            grads_vec = fused_ops.TorchTensorVector(grads)
+            multi_tensor_applier(fused_ops.multi_tensor_scale, dummy_overflow_buf, [grads_vec, grads_vec], clip_coef)
 
     return total_norm

orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/fused_ops_frontend.cpp

@@ -42,6 +42,12 @@ void multi_tensor_axpby_cuda(int chunk_size,
                              float b,
                              int arg_to_check);
 
+// This function is adapted from NVIDIA/apex
+// https://github.com/NVIDIA/apex/blob/0c7d8e3fa9a095a1641a2290877436d0314b69c6/csrc/amp_C_frontend.cpp#L30
+std::tuple<at::Tensor, at::Tensor> multi_tensor_l2norm_cuda(int chunk_size, at::Tensor noop_flag,
+                                                            std::vector<std::vector<at::Tensor>> tensor_lists,
+                                                            at::optional<bool> per_tensor_python);
+
 class MemoryBuffer {
  public:
   MemoryBuffer(size_t numel, at::Tensor val) {
@@ -218,4 +224,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("unscale_fp16_grads_into_fp32_grads",
         &unscale_fp16_grads_into_fp32_grads,
         "Unscale those fp16 gradients into fp32 gradient buffers.");
+  m.def("multi_tensor_scale",
+        &multi_tensor_scale_cuda,
+        "Fused overflow check + scale for a list of contiguous tensors");
+  m.def("multi_tensor_l2norm",
+        &multi_tensor_l2norm_cuda,
+        "Computes L2 norm for a list of contiguous tensors");
 }

orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/multi_tensor_l2norm_kernel.cu

@@ -0,0 +1,167 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Copyright NVIDIA/apex
+// This file is adapted from NVIDIA/apex, commit 3ff1a10f72ec07067c4e44759442329804ac5162
+
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+#include <c10/cuda/CUDAGuard.h>
+// Another possibility:
+// #include <torch/all.h>
+
+#include <assert.h>
+
+#include "type_shim.h"
+#include "multi_tensor_apply.cuh"
+
+#define BLOCK_SIZE 512
+#define ILP 4
+
+template <typename T>
+__device__ __forceinline__ bool is_aligned(T* p) {
+  return ((uint64_t)p) % (ILP * sizeof(T)) == 0;
+}
+
+template <typename T>
+__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset) {
+  typedef typename std::aligned_storage<ILP * sizeof(T), ILP * alignof(T)>::type LT;
+  ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];
+}
+
+template <typename x_t>
+struct L2NormFunctor {
+  __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<1>& tl,
+                                             float* output, float* output_per_tensor, bool per_tensor,
+                                             int max_chunks_per_tensor) {
+    // I'd like this kernel to propagate infs/nans.
+    // if(*noop_gmem == 1)
+    //   return;
+
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    x_t* x = (x_t*)tl.addresses[0][tensor_loc];
+    x += chunk_idx * chunk_size;
+
+    n -= chunk_idx * chunk_size;
+
+    __shared__ float s_vals[512];
+
+    float vals[ILP];  // = {0}; // this probably works too but I want to be sure...
+    x_t r_x[ILP];
+    for (int i = 0; i < ILP; i++) {
+      vals[i] = 0.f;
+      r_x[i] = 0;
+    }
+
+    // to make things simple, we put aligned case in a different code path
+    if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {
+      for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {
+        // load
+        load_store(r_x, x, 0, i_start);
+#pragma unroll
+        for (int ii = 0; ii < ILP; ii++) {
+          float next = static_cast<float>(r_x[ii]);
+          vals[ii] += next * next;
+        }
+      }
+    } else {
+      for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {
+#pragma unroll
+        for (int ii = 0; ii < ILP; ii++) {
+          int i = i_start + threadIdx.x + ii * blockDim.x;
+          if (i < n && i < chunk_size) {
+            float next = static_cast<float>(x[i]);
+            vals[ii] += next * next;
+          }
+        }
+      }
+    }
+
+    float val = 0.f;
+    for (int i = 0; i < ILP; i++) val += vals[i];
+
+    float final = reduce_block_into_lanes(s_vals, val);
+
+    if (threadIdx.x == 0) {
+      if (!isfinite(final)) *noop_gmem = 1;  // Blindly fire off a write.  These will race but that's ok.
+      output[blockIdx.x] += final;
+      if (per_tensor)
+        output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor + chunk_idx] = final;
+    }
+  }
+};
+
+__global__ void cleanup(float* output, float* output_per_tensor, float* ret, float* ret_per_tensor, bool per_tensor,
+                        int max_chunks_per_tensor) {
+  __shared__ float vals[512];
+
+  if (blockIdx.x == 0) {
+    float val = 0;
+    if (threadIdx.x < 320) val = output[threadIdx.x];
+
+    float final = reduce_block_into_lanes(vals, val);
+
+    if (threadIdx.x == 0) *ret = sqrt(final);
+  }
+
+  if (per_tensor) {
+    float* output_this_tensor = output_per_tensor + blockIdx.x * max_chunks_per_tensor;
+
+    float val = 0;
+    for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x) val += output_this_tensor[i];
+
+    float final = reduce_block_into_lanes(vals, val);
+
+    if (threadIdx.x == 0) ret_per_tensor[blockIdx.x] = sqrt(final);
+  }
+}
+
+std::tuple<at::Tensor, at::Tensor> multi_tensor_l2norm_cuda(int chunk_size, at::Tensor noop_flag,
+                                                            std::vector<std::vector<at::Tensor>> tensor_lists,
+                                                            at::optional<bool> per_tensor_python) {
+  bool per_tensor = per_tensor_python.has_value() ? per_tensor_python.value() : false;
+
+  auto float_options = tensor_lists[0][0].options().dtype(at::kFloat);
+  auto output = at::zeros({320}, float_options);
+
+  at::Tensor output_per_tensor;
+  at::Tensor ret_per_tensor;
+
+  int ntensors = tensor_lists[0].size();
+  int max_chunks_per_tensor = -1;
+
+  if (per_tensor) {
+    for (int t = 0; t < ntensors; t++) {
+      int max_chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size;
+      if (max_chunks_this_tensor > max_chunks_per_tensor) max_chunks_per_tensor = max_chunks_this_tensor;
+    }
+    output_per_tensor = at::zeros({ntensors * max_chunks_per_tensor}, float_options);
+    ret_per_tensor = at::empty({ntensors}, float_options);
+  } else {
+    ret_per_tensor = at::empty({0}, float_options);
+  }
+
+  DISPATCH_DOUBLE_FLOAT_AND_HALF(
+      tensor_lists[0][0].scalar_type(), 0, "multi_tensor_l2norm_cuda",
+      multi_tensor_apply<1>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, L2NormFunctor<scalar_t_0>(),
+                            output.data_ptr<float>(), per_tensor ? output_per_tensor.data_ptr<float>() : nullptr,
+                            per_tensor, max_chunks_per_tensor);)
+
+  AT_CUDA_CHECK(cudaGetLastError());
+  // AT_CUDA_CHECK(cudaDeviceSynchronize());
+
+  // This involves one more small kernel launches, but will be negligible end to end.
+  // I could get rid of these by hacking the functor + multi tensor harness with persistence
+  // logic, but keeping it simple for now
+  auto ret = at::empty({1}, output.options());
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(output));
+  auto stream = at::cuda::getCurrentCUDAStream();
+  cleanup<<<per_tensor ? ntensors : 1, 512, 0, stream>>>(
+      output.data_ptr<float>(), per_tensor ? output_per_tensor.data_ptr<float>() : nullptr, ret.data_ptr<float>(),
+      per_tensor ? ret_per_tensor.data_ptr<float>() : nullptr, per_tensor, max_chunks_per_tensor);
+
+  return std::tuple<at::Tensor, at::Tensor>(ret, ret_per_tensor);
+}

orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/setup.py

@@ -15,6 +15,7 @@ filenames = [
     os.path.join(os.path.dirname(__file__), "multi_tensor_adam.cu"),
     os.path.join(os.path.dirname(__file__), "multi_tensor_scale_kernel.cu"),
     os.path.join(os.path.dirname(__file__), "multi_tensor_axpby_kernel.cu"),
+    os.path.join(os.path.dirname(__file__), "multi_tensor_l2norm_kernel.cu"),
 ]
 
 use_rocm = True if os.environ["ONNXRUNTIME_ROCM_VERSION"] else False

orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/type_shim.h

@@ -24,3 +24,47 @@
         }                                                                        \
         default: AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \
     }
+
+template <typename T>
+__device__ __forceinline__ T reduce_block_into_lanes(T* x, T val, int lanes = 1,
+                                                     bool share_result = false)  // lanes is intended to be <= 32.
+{
+  int tid = threadIdx.x + threadIdx.y * blockDim.x;
+  int blockSize = blockDim.x * blockDim.y;  // blockSize is intended to be a multiple of 32.
+
+  if (blockSize >= 64) {
+    x[tid] = val;
+    __syncthreads();
+  }
+
+#pragma unroll
+  for (int i = (blockSize >> 1); i >= 64; i >>= 1) {
+    if (tid < i) x[tid] = x[tid] + x[tid + i];
+    __syncthreads();
+  }
+
+  T final;
+
+  if (tid < 32) {
+    if (blockSize >= 64)
+      final = x[tid] + x[tid + 32];
+    else
+      final = val;
+      // __SYNCWARP();
+
+#pragma unroll
+#if defined(CUDA_VERSION) && CUDA_VERSION >= 9000
+    for (int i = 16; i >= lanes; i >>= 1) final = final + __shfl_down_sync(0xffffffff, final, i);
+#else
+    for (int i = 16; i >= lanes; i >>= 1) final = final + __shfl_down(final, i);
+#endif
+  }
+
+  if (share_result) {
+    if (tid < lanes) x[tid] = final;  // EpilogueOp
+    // Make sure the smem result is visible to all warps.
+    __syncthreads();
+  }
+
+  return final;
+}