CUDA optional deepspeed ops

.github/workflows/nv-torch-latest-cpu.yml

@@ -0,0 +1,52 @@
+name: nv-torch-latest-cpu
+
+on:
+  push:
+    branches:
+      - 'master'
+      - 'staging**'
+    paths-ignore:
+      - 'docs/**'
+  pull_request:
+    paths-ignore:
+      - 'docs/**'
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  unit-tests:
+    runs-on: ubuntu-20.04
+
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: environment
+        run: |
+          echo "JobID: $AISC_NODE_INSTANCE_ID"
+          which python
+          python --version
+          pip install --upgrade pip
+          pip uninstall --yes torch torchvision triton
+          pip install torch==1.12.0+cpu torchvision==0.13.0+cpu torchaudio==0.12.0 --extra-index-url https://download.pytorch.org/whl/cpu
+          python -c "import torch; print('torch:', torch.__version__, torch)"
+          python -c "import torch; print('CUDA available:', torch.cuda.is_available())"
+
+      - name: Install deepspeed
+        run: |
+          pip uninstall --yes deepspeed
+          pip install .[dev,autotuning]
+          ds_report
+
+      - name: Python environment
+        run: |
+          pip list
+
+      - name: Unit tests
+        run: |
+          unset TORCH_CUDA_ARCH_LIST # only jit compile for current arch
+          if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi
+          cd tests
+          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -n 4 unit/ --torch_ver="1.12"
+          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -m 'sequential' unit/ --torch_ver="1.12"

csrc/adagrad/cpu_adagrad.cpp

@@ -1,16 +1,16 @@
 #include "cpu_adagrad.h"
-#include <cuda_runtime_api.h>
-#include <math.h>
-#include <omp.h>
 #include <torch/extension.h>
 #include <iostream>
 #include <memory>
 #include <type_traits>
 #include <unordered_map>
+#if defined(__ENABLE_CUDA__)
+#include <cuda_runtime_api.h>
 #include "cublas_v2.h"
 #include "cuda.h"
 #include "curand.h"
 #include "custom_cuda_layers.h"
+#endif
 
 static std::unordered_map<int, std::shared_ptr<void>> s_optimizers;
 
@@ -20,7 +20,7 @@ void Adagrad_Optimizer::Step_1(float* _params,
                                float* grads,
                                float* _exp_avg_sq,
                                size_t _param_size,
-                               __half* dev_params,
+                               ds_half_precision_t* dev_params,
                                bool half_precision)
 {
     size_t rounded_size = 0;
@@ -30,17 +30,19 @@ void Adagrad_Optimizer::Step_1(float* _params,
 #endif
     if (_param_size > rounded_size) {
         float step_size = -1 * _alpha;
-        __half* grads_cast_h;
-        __half* params_cast_h;
+        ds_half_precision_t* grads_cast_h;
+        ds_half_precision_t* params_cast_h;
         if (half_precision) {
-            grads_cast_h = reinterpret_cast<__half*>(grads);
-            params_cast_h = reinterpret_cast<__half*>(_params);
+            grads_cast_h = reinterpret_cast<ds_half_precision_t*>(grads);
+            params_cast_h = reinterpret_cast<ds_half_precision_t*>(_params);
         }
         for (size_t t = rounded_size; t < _param_size; t += TILE) {
             size_t copy_size = TILE;
             if ((t + TILE) > _param_size) copy_size = _param_size - t;
             size_t offset = copy_size + t;
+#if defined(__ENABLE_CUDA__)
             if ((t / TILE) >= 2) { cudaStreamSynchronize(_streams[_buf_index]); }
+#endif
 #pragma omp parallel for
             for (size_t k = t; k < offset; k++) {
                 float grad = half_precision ? (float)grads_cast_h[k] : grads[k];
@@ -55,21 +57,24 @@ void Adagrad_Optimizer::Step_1(float* _params,
                 grad += _eps;
                 grad = momentum / grad;
                 param = grad * step_size + param;
+#if defined(__ENABLE_CUDA__)
                 if (dev_params) _doubled_buffer[_buf_index][k - t] = param;
-
+#endif
                 if (half_precision)
-                    params_cast_h[k] = (__half)param;
+                    params_cast_h[k] = (ds_half_precision_t)param;
                 else
                     _params[k] = param;
                 // STORE UPDATE TERM TO GRAD'S MEMORY
                 grads[k] = grad * step_size;
                 _exp_avg_sq[k] = variance;
             }
+#if defined(__ENABLE_CUDA__)
             if (dev_params) {
                 launch_param_update(
                     _doubled_buffer[_buf_index], dev_params + t, (copy_size), _streams[_buf_index]);
                 _buf_index = !_buf_index;
             }
+#endif
         }
     }
 }
@@ -78,7 +83,7 @@ void Adagrad_Optimizer::Step_4(float* _params,
                                float* grads,
                                float* _exp_avg_sq,
                                size_t _param_size,
-                               __half* dev_params,
+                               ds_half_precision_t* dev_params,
                                bool half_precision)
 {
     size_t rounded_size = 0;
@@ -130,7 +135,7 @@ void Adagrad_Optimizer::Step_8(float* _params,
                                float* grads,
                                float* _exp_avg_sq,
                                size_t _param_size,
-                               __half* dev_params,
+                               ds_half_precision_t* dev_params,
                                bool half_precision)
 {
     size_t rounded_size = 0;
@@ -170,7 +175,9 @@ int ds_adagrad_step(int optimizer_id,
     opt->update_state(lr, epsilon, weight_decay);
     opt->Step_8(params_ptr, grads_ptr, exp_avg_sq_ptr, params_c.size(0));
 
+#if defined(__ENABLE_CUDA__)
     opt->SynchronizeStreams();
+#endif
     return 0;
 }
 
@@ -184,14 +191,15 @@ int ds_adagrad_step_plus_copy(int optimizer_id,
                               torch::Tensor& exp_avg_sq,
                               torch::Tensor& gpu_params)
 {
+#if defined(__ENABLE_CUDA__)
     auto params_c = params.contiguous();
     auto gpu_params_c = gpu_params.contiguous();
     auto exp_avg_sq_c = exp_avg_sq.contiguous();
     auto grads_c = grads.contiguous();
 
     float* params_ptr = (float*)params_c.data_ptr();
     float* grads_ptr = (float*)grads_c.data_ptr();
-    __half* gpu_params_ptr = (__half*)gpu_params_c.data_ptr();
+    ds_half_precision_t* gpu_params_ptr = (ds_half_precision_t*)gpu_params_c.data_ptr();
     float* exp_avg_sq_ptr = (float*)exp_avg_sq_c.data_ptr();
 
     std::shared_ptr<Adagrad_Optimizer> opt =
@@ -206,6 +214,9 @@ int ds_adagrad_step_plus_copy(int optimizer_id,
                 (params.options().dtype() == at::kHalf));
 
     opt->SynchronizeStreams();
+#else
+    assert(false);
+#endif
     return 0;
 }
 

csrc/adam/cpu_adam.cpp

@@ -1,16 +1,18 @@
 #include "cpu_adam.h"
-#include <cuda_runtime_api.h>
-#include <math.h>
-#include <omp.h>
 #include <torch/extension.h>
+#include <cassert>
 #include <iostream>
 #include <memory>
 #include <type_traits>
 #include <unordered_map>
+
+#if defined(__ENABLE_CUDA__)
+#include <cuda_runtime_api.h>
 #include "cublas_v2.h"
 #include "cuda.h"
 #include "curand.h"
 #include "custom_cuda_layers.h"
+#endif
 
 static std::unordered_map<int, std::shared_ptr<void>> s_optimizers;
 
@@ -21,7 +23,7 @@ void Adam_Optimizer::Step_1(float* _params,
                             float* _exp_avg,
                             float* _exp_avg_sq,
                             size_t _param_size,
-                            __half* dev_params,
+                            ds_half_precision_t* dev_params,
                             bool half_precision)
 {
     size_t rounded_size = 0;
@@ -41,19 +43,20 @@ void Adam_Optimizer::Step_1(float* _params,
 
         float step_size = -1 * _alpha / _bias_correction1;
         float w_decay = -1 * _alpha * _weight_decay;
-        __half* grads_cast_h;
-        __half* params_cast_h;
+        ds_half_precision_t* grads_cast_h;
+        ds_half_precision_t* params_cast_h;
         if (half_precision) {
-            grads_cast_h = reinterpret_cast<__half*>(grads);
-            params_cast_h = reinterpret_cast<__half*>(_params);
+            grads_cast_h = reinterpret_cast<ds_half_precision_t*>(grads);
+            params_cast_h = reinterpret_cast<ds_half_precision_t*>(_params);
         }
 
         for (size_t t = rounded_size; t < _param_size; t += TILE) {
             size_t copy_size = TILE;
             if ((t + TILE) > _param_size) copy_size = _param_size - t;
             size_t offset = copy_size + t;
+#if defined(__ENABLE_CUDA__)
             if ((t / TILE) >= 2) { cudaStreamSynchronize(_streams[_buf_index]); }
-
+#endif
 #pragma omp parallel for
             for (size_t k = t; k < offset; k++) {
                 float grad = half_precision ? (float)grads_cast_h[k] : grads[k];
@@ -73,21 +76,24 @@ void Adam_Optimizer::Step_1(float* _params,
                 grad = momentum / grad;
                 if (_weight_decay > 0 && _adamw_mode) { param += w_decay * param; }
                 param = grad * step_size + param;
+#if defined(__ENABLE_CUDA__)
                 if (dev_params) _doubled_buffer[_buf_index][k - t] = param;
-
+#endif
                 if (half_precision)
-                    params_cast_h[k] = (__half)param;
+                    params_cast_h[k] = (ds_half_precision_t)param;
                 else
                     _params[k] = param;
                 _exp_avg[k] = momentum;
                 _exp_avg_sq[k] = variance;
             }
+#if defined(__ENABLE_CUDA__)
             if (dev_params) {
                 launch_param_update(
                     _doubled_buffer[_buf_index], dev_params + t, (copy_size), _streams[_buf_index]);
 
                 _buf_index = !_buf_index;
             }
+#endif
         }
     }
 }
@@ -97,7 +103,7 @@ void Adam_Optimizer::Step_4(float* _params,
                             float* _exp_avg,
                             float* _exp_avg_sq,
                             size_t _param_size,
-                            __half* dev_params,
+                            ds_half_precision_t* dev_params,
                             bool half_precision)
 {
     size_t rounded_size = 0;
@@ -166,7 +172,7 @@ void Adam_Optimizer::Step_8(float* _params,
                             float* _exp_avg,
                             float* _exp_avg_sq,
                             size_t _param_size,
-                            __half* dev_params,
+                            ds_half_precision_t* dev_params,
                             bool half_precision)
 {
     size_t rounded_size = 0;
@@ -228,7 +234,9 @@ int ds_adam_step(int optimizer_id,
                 nullptr,
                 (params.options().dtype() == at::kHalf));
 
+#if defined(__ENABLE_CUDA__)
     opt->SynchronizeStreams();
+#endif
     return 0;
 }
 
@@ -246,6 +254,7 @@ int ds_adam_step_plus_copy(int optimizer_id,
                            torch::Tensor& exp_avg_sq,
                            torch::Tensor& gpu_params)
 {
+#if defined(__ENABLE_CUDA__)
     auto params_c = params.contiguous();
     auto gpu_params_c = gpu_params.contiguous();
     auto exp_avg_c = exp_avg.contiguous();
@@ -254,7 +263,7 @@ int ds_adam_step_plus_copy(int optimizer_id,
 
     float* params_ptr = (float*)params_c.data_ptr();
     float* grads_ptr = (float*)grads_c.data_ptr();
-    __half* gpu_params_ptr = (__half*)gpu_params_c.data_ptr();
+    ds_half_precision_t* gpu_params_ptr = (ds_half_precision_t*)gpu_params_c.data_ptr();
     float* exp_avg_ptr = (float*)exp_avg_c.data_ptr();
     float* exp_avg_sq_ptr = (float*)exp_avg_sq_c.data_ptr();
 
@@ -271,6 +280,9 @@ int ds_adam_step_plus_copy(int optimizer_id,
                 (params.options().dtype() == at::kHalf));
 
     opt->SynchronizeStreams();
+#else
+    assert(false);
+#endif
     return 0;
 }