Repository: deepspeedai/DeepSpeed
Pull request: Fix cpu-adam AVX performance
This fixes https://github.com/microsoft/DeepSpeed/issues/1635

Problem statement:
[BUG] CPU adam too slow even with AVX512
**Describe the bug**
Optimizer step is taking unexpectedly long. I am running a GPT-like model with 1.35B parameters with ZeRO-3, optimizer CPU offload enabled.

I have added extra timers to measure exact timing for optimizer step:

```
 time (ms) | update_optims: 21766.25 | unscale_grads: 506.45 | optimizer_step: 22850.66
 time (ms) | forward_microstep: 376.60 | backward_microstep: 864.56 | backward_inner_microstep: 806.86 | backward_allreduce_microstep: 57.62 | step_microstep: 22880.87
time (ms) | forward: 376.60 | backward: 864.56 | backward_inner: 806.87 | backward_allreduce: 57.62 | step: 22880.87
```
I believe update computation using cpu-adam should be much faster - less than 5 seconds but here it's taking more than 22 seconds ...

**System info**
 - OS: Ubuntu 18.04
 - One RTX3090 and  Xeon Silver 4210R(which has 1 AVX512 unit)
 - Python 3.9, Pytorch 1.7.1+cu110

Below is my deepspeed configuration
```
{
    "fp16": {
        "enabled": true,
        "loss_scale": 10,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },

    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 5e-05,
            "betas": [0.9, 0.999],
            "eps": 1e-08,
            "weight_decay": 0.0
        }
    },

    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": 0,
            "warmup_max_lr": 5e-05,
            "warmup_num_steps": 0
        }
    },

    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_parameter": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": "auto",
        "sub_group_size": 1e9,
        "stage3_prefetch_bucket_size": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_fp16_weights_on_model_save": true
    },
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": 1,
    "wall_clock_breakdown": true,
    "memory_breakdown": false
}
```

I see that CPU peak utilization never increases beyond 50%. However, when two GPUs are used CPU utilization also doubles, but still its latency never reaches expected value.

Hints from issue discussion:
I meet the same problem in DeepSpeedExamples
https://github.com/microsoft/DeepSpeedExamples/issues/151
Hi @PeiqiWang, @nobellant215 

Thanks for reporting this issue. I will look into this.

@PeiqiWang, can I ask what model-size you were using in your experiments?

@nobellant215, the difference that you mentioned is really huge (5 vs 22). I am wondering how you are measuring the time here. We have a performance check for CPU-Adam (https://github.com/microsoft/DeepSpeed/blob/master/tests/perf/adam_test.py), is it possible that you ran it with the same number of parameters as you mentioned and compare the latency numbers? I see 0.34 sec on my side with 1.35B parameter size!

Thanks,
Reza 
Thanks for the reply.
Previously measured time was obtained by just placing synchronized clock before and after the optimizer step in stage3.py.
I have tried the performance check and I am now seeing 12 ~ 13 sec even with torch.num_thread set to 128.  

I think I was running with the scalar CPU-Adam as the AVX512 was not found. But, after switching the AVX512 on, I see serious slowdown. I think this might have caused with some of the new PR we had for the CPU-Adam implementation to add other optimizers. I will investigate this and give a solution for this soon.

Best,
Reza

Commit message: fixing the loop-unrolling for the avx operations

Patch to review:
diff --git a/csrc/includes/simd.h b/csrc/includes/simd.h
index 216d281a2672..44c4da83e01c 100644
--- a/csrc/includes/simd.h
+++ b/csrc/includes/simd.h
@@ -60,7 +60,7 @@ union AVX_Data {
 template <int span>
 inline void simd_store(float* dst, AVX_Data* src, bool half_precision)
 {
-#pragma omp parallel for
+#pragma unroll
     for (size_t i = 0; i < span; ++i) {
         SIMD_STORE2(dst + SIMD_WIDTH * i, src[i].data, half_precision);
     }
@@ -68,7 +68,7 @@ inline void simd_store(float* dst, AVX_Data* src, bool half_precision)
 template <int span>
 inline void simd_load(AVX_Data* dst, float* src, bool half_precision)
 {
-#pragma omp parallel for
+#pragma unroll
     for (size_t i = 0; i < span; ++i) {
         dst[i].data = SIMD_LOAD2(src + SIMD_WIDTH * i, half_precision);
     }
@@ -76,7 +76,7 @@ inline void simd_load(AVX_Data* dst, float* src, bool half_precision)
 template <int span>
 inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data src_m_r, AVX_Data* src_a)
 {
-#pragma omp parallel for
+#pragma unroll
     for (size_t i = 0; i < span; ++i) {
         dst[i].data = SIMD_FMA(src_m_l[i].data, src_m_r.data, src_a[i].data);
     }
@@ -84,7 +84,7 @@ inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data src_m_r, AVX_Dat
 template <int span>
 inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data src_m_r, AVX_Data src_a)
 {
-#pragma omp parallel for
+#pragma unroll
     for (size_t i = 0; i < span; ++i) {
         dst[i].data = SIMD_FMA(src_m_l[i].data, src_m_r.data, src_a.data);
     }
@@ -92,7 +92,7 @@ inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data src_m_r, AVX_Dat
 template <int span>
 inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data* src_m_r, AVX_Data* src_a)
 {
-#pragma omp parallel for
+#pragma unroll
     for (size_t i = 0; i < span; ++i) {
         dst[i].data = SIMD_FMA(src_m_l[i].data, src_m_r[i].data, src_a[i].data);
     }
@@ -100,37 +100,37 @@ inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data* src_m_r, AVX_Da
 template <int span>
 inline void simd_sqrt(AVX_Data* dst, AVX_Data* src)
 {
-#pragma omp parallel for
+#pragma unroll
     for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_SQRT(src[i].data); }
 }
 template <int span>
 inline void simd_add(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r)
 {
-#pragma omp parallel for
+#pragma unroll
     for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_ADD(src_a_l[i].data, src_a_r.data); }
 }
 template <int span>
 inline void simd_add(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r)
 {
-#pragma omp parallel for
+#pragma unroll
     for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_ADD(src_a_l[i].data, src_a_r[i].data); }
 }
 template <int span>
 inline void simd_mul(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r)
 {
-#pragma omp parallel for
+#pragma unroll
     for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_MUL(src_a_l[i].data, src_a_r.data); }
 }
 template <int span>
 inline void simd_mul(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r)
 {
-#pragma omp parallel for
+#pragma unroll
     for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_MUL(src_a_l[i].data, src_a_r[i].data); }
 }
 template <int span>
 inline void simd_div(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r)
 {
-#pragma omp parallel for
+#pragma unroll
     for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_DIV(src_a_l[i].data, src_a_r[i].data); }
 }
 
diff --git a/deepspeed/ops/transformer/inference/transformer_inference.py b/deepspeed/ops/transformer/inference/transformer_inference.py
index 36d2a7d432dc..3161a90edcc9 100755
--- a/deepspeed/ops/transformer/inference/transformer_inference.py
+++ b/deepspeed/ops/transformer/inference/transformer_inference.py
@@ -146,6 +146,31 @@ def _transpose_for_context(x):
                                       (hidden_size_per_partition,)
             return x.view(*new_x_layer_shape)
 
+        def backup_attention(mixed_query, key_layer, value_layer, input_mask):
+            if layer_past is not None:
+                past_key, past_value = layer_past
+                key_layer = torch.cat((past_key.type_as(key_layer), key_layer), dim=-2)
+                value_layer = torch.cat((past_value.type_as(value_layer),
+                                         value_layer),
+                                        dim=-2)
+            query = _transpose_for_scores(mixed_query, False)
+            key = _transpose_for_scores(key_layer, True)
+            value = _transpose_for_scores(value_layer, False)
+            p = torch.matmul(query, key)
+
+            ds_softmax = inference_cuda_module.softmax_fp16 if config.fp16 else \
+                            inference_cuda_module.softmax_fp32
+            p = ds_softmax(p / (float(key.size(-2))**0.5),
+                           input_mask,
+                           True,
+                           False,
+                           False,
+                           256)
+            p = p.to(value.dtype)
+            context_layer = torch.matmul(p, value)
+            context_layer = _transpose_for_context(context_layer)
+            return context_layer, key_layer, value_layer
+
         def compute_attention(qkv_out, input_mask):
             score_context_func = inference_cuda_module.softmax_context_fp32 if (not config.fp16) else \
                                     inference_cuda_module.softmax_context_fp16
@@ -176,7 +201,10 @@ def compute_attention(qkv_out, input_mask):
             if no_masking:
                 input_mask = torch.empty(1)
             head_size = (mixed_query.shape[-1] // num_attention_heads_per_partition)
-
+            #return backup_attention(mixed_query,
+            #     key_layer,
+            #     value_layer,
+            #     input_mask)
             unfused_mode = not config.specialized_mode or \
                                 mixed_query.shape[1] >= 32 or head_size > 128
 


Code review comments:
diff_hunk: @@ -146,6 +146,31 @@ def _transpose_for_context(x):
                                       (hidden_size_per_partition,)
             return x.view(*new_x_layer_shape)
 
+        def backup_attention(mixed_query, key_layer, value_layer, input_mask):
path: deepspeed/ops/transformer/inference/transformer_inference.py
text: @user1:
@author, was this an accidental addition? seems unrelated to the PR right?

@author:
Thanks Jeff, will remove it!