fix cpu adam compilation for AVX2

csrc/includes/cpu_adam.h

@@ -1,7 +1,9 @@
 #pragma once
 
+#include <cpuid.h>
 #include <cuda_fp16.h>
 #include <cuda_runtime_api.h>
+#include <stdio.h>
 #include <x86intrin.h>
 #include <cassert>
 #include "context.h"
@@ -18,10 +20,30 @@
         }                                                                                      \
     }
 
-#define SIMD_WIDTH 16
-
 #define TILE (1024 * 1024 * 1024)
 
+#if defined(__AVX512__)
+#define SIMD_STORE(a, d) _mm512_storeu_ps(a, d)
+#define SIMD_LOAD(x) _mm512_loadu_ps(x)
+#define SIMD_SET(x) _mm512_set1_ps(x)
+#define SIMD_MUL(x, y) _mm512_mul_ps(x, y)
+#define SIMD_FMA(x, y, c) _mm512_fmadd_ps(x, y, c)
+#define SIMD_SQRT(x) _mm512_sqrt_ps(x)
+#define SIMD_DIV(x, y) _mm512_div_ps(x, y)
+#define SIMD_WIDTH 16
+#else
+#if defined(__AVX256__)
+#define SIMD_STORE(a, d) _mm256_storeu_ps(a, d)
+#define SIMD_LOAD(x) _mm256_loadu_ps(x)
+#define SIMD_SET(x) _mm256_set1_ps(x)
+#define SIMD_MUL(x, y) _mm256_mul_ps(x, y)
+#define SIMD_FMA(x, y, c) _mm256_fmadd_ps(x, y, c)
+#define SIMD_SQRT(x) _mm256_sqrt_ps(x)
+#define SIMD_DIV(x, y) _mm256_div_ps(x, y)
+#define SIMD_WIDTH 8
+#endif
+#endif
+
 class Adam_Optimizer {
 public:
     Adam_Optimizer(float alpha = 1e-3,
@@ -64,12 +86,23 @@ class Adam_Optimizer {
                 float* _exp_avg_sq,
                 size_t _param_size,
                 __half* dev_params = nullptr);
+    inline void IncrementStep()
+    {
+        _betta1_t *= _betta1;
+        _betta2_t *= _betta2;
+    }
 
 private:
-    union AVX_512 {
+#if defined(__AVX512__) or defined(__AVX256__)
+    union AVX_Data {
+#if defined(__AVX512__)
         __m512 data;
+#else
+        __m256 data;
+#endif
         // float data_f[16];
     };
+#endif
 
     float _alpha;
     float _betta1;

docs/_pages/features.md

@@ -162,6 +162,12 @@ Please see the [core API doc](https://deepspeed.readthedocs.io/) for more detail
 With DeepSpeed, the user can choose to use a high performance implementation of ADAM from
 NVIDIA, or any training optimizer that extends torch's `torch.optim.Optimizer` class.
 
+### CPU-Adam: High-Performance vectorized implementation of Adam
+We introduce an efficient implementation of Adam optimizer on CPU that improves the parameter-update
+performance by nearly an order of magnitude. Comparing to torch-adam, we observe 5.1x to 6.5x 
+speedups considering the model-size between 1 to 10 billion parameters. For the CPU-Adam implementation, 
+we use the AVX SIMD instructions on Intel-x86 architecture. We support both AVX-512 and AVX-2 instruction sets. 
+
 ### Memory bandwidth optimized FP16 Optimizer
 Mixed precision training is handled by the DeepSpeed FP16 Optimizer. This optimizer not
 only handles FP16 training but is also highly efficient. The performance of weight update

docs/index.md

@@ -167,6 +167,7 @@ overview](/features/) for descriptions and usage.
     * Automatic loss scaling with mixed precision
 * [Training Optimizers](/features/#training-optimizers)
     * Fused Adam optimizer and arbitrary `torch.optim.Optimizer`
+    * CPU-Adam: High-Performance vectorized Adam
     * Memory bandwidth optimized FP16 Optimizer
     * Large Batch Training with LAMB Optimizer
     * Memory efficient Training with ZeRO Optimizer

install.sh

@@ -164,10 +164,10 @@ if [ ! -f $hostfile ]; then
     local_only=1
 fi
 
-#if [ "$skip_requirements" == "0" ]; then
-#    # Ensure dependencies are installed locally
-#    $PIP_SUDO $PIP_INSTALL -r requirements.txt
-#fi
+if [ "$skip_requirements" == "0" ]; then
+   # Ensure dependencies are installed locally
+   $PIP_SUDO $PIP_INSTALL -r requirements/requirements.txt
+fi
 
 # Build wheels
 if [ "$third_party_install" == "1" ]; then
@@ -220,10 +220,10 @@ else
     tmp_wheel_path="/tmp/deepspeed_wheels"
 
     pdsh -w $hosts "if [ -d $tmp_wheel_path ]; then rm $tmp_wheel_path/*.whl; else mkdir -pv $tmp_wheel_path; fi"
-    #pdcp -w $hosts requirements/*.txt ${tmp_wheel_path}/
-    #if [ "$skip_requirements" == "0" ]; then
-    #    pdsh -w $hosts "$PIP_SUDO $PIP_INSTALL -r ${tmp_wheel_path}/requirements.txt"
-    #fi
+    pdcp -w $hosts requirements/requirements.txt ${tmp_wheel_path}/
+    if [ "$skip_requirements" == "0" ]; then
+       pdsh -w $hosts "$PIP_SUDO $PIP_INSTALL -r ${tmp_wheel_path}/requirements.txt"
+    fi
     if [ "$third_party_install" == "1" ]; then
         pdsh -w $hosts "$PIP_SUDO pip uninstall -y apex"
         pdcp -w $hosts third_party/apex/dist/apex*.whl $tmp_wheel_path/

requirements/requirements.txt

@@ -2,4 +2,5 @@ torch>=1.2
 torchvision>=0.4.0
 tqdm
 psutil
+cpufeature
 tensorboardX==1.8

setup.py

@@ -96,6 +96,16 @@ def fetch_requirements(path):
     version_ge_1_5 = ['-DVERSION_GE_1_5']
 version_dependent_macros = version_ge_1_1 + version_ge_1_3 + version_ge_1_5
 
+import cpufeature
+d = cpufeature.CPUFeature
+
+SIMD_WIDTH = ''
+if d['AVX512f']:
+    SIMD_WIDTH = '-D__AVX512__'
+elif d['AVX2']:
+    SIMD_WIDTH = '-D__AVX256__'
+print("SIMD_WIDTH = ", SIMD_WIDTH)
+
 ext_modules = []
 
 ## Lamb ##
@@ -135,7 +145,8 @@ def fetch_requirements(path):
                               '-g',
                               '-Wno-reorder',
                               '-march=native',
-                              '-fopenmp'
+                              '-fopenmp',
+                              SIMD_WIDTH
                           ],
                           'nvcc': [
                               '-O3',