SuperOffload - PCIe Optimization

csrc/adam/cpu_adam.cpp

@@ -8,8 +8,34 @@
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
 {
     using namespace pybind11::literals;
-    m.def("adam_update", &ds_adam_step, "DeepSpeed CPU Adam update (C++)");
-    m.def("adam_rollback", &ds_adam_rollback, "DeepSpeed CPU Adam rollback (C++)");
+    // The update, the rescale and the rollback release the GIL, so that a CPU Adam step in another
+    // thread lets the training thread run.
+    m.def("adam_update",
+          &ds_adam_step,
+          "DeepSpeed CPU Adam update (C++); grad_scale multiplies each gradient as it is read",
+          "optimizer_id"_a,
+          "step"_a,
+          "lr"_a,
+          "beta1"_a,
+          "beta2"_a,
+          "epsilon"_a,
+          "weight_decay"_a,
+          "bias_correction"_a,
+          "params"_a,
+          "grads"_a,
+          "exp_avg"_a,
+          "exp_avg_sq"_a,
+          "grad_scale"_a = 1.0f,
+          pybind11::call_guard<pybind11::gil_scoped_release>());
+    m.def("adam_rescale_step",
+          &ds_adam_rescale_step,
+          "DeepSpeed CPU Adam: turn the last step, taken on grads * old_scale, into the step on "
+          "grads * new_scale (C++)",
+          pybind11::call_guard<pybind11::gil_scoped_release>());
+    m.def("adam_rollback",
+          &ds_adam_rollback,
+          "DeepSpeed CPU Adam rollback (C++)",
+          pybind11::call_guard<pybind11::gil_scoped_release>());
     m.def("create_adam", &create_adam_optimizer, "DeepSpeed CPU Adam (C++)");
     m.def("destroy_adam", &destroy_adam_optimizer, "DeepSpeed CPU Adam destroy (C++)");
 

csrc/adam/cpu_adam_impl.cpp

@@ -32,9 +32,11 @@ static std::unordered_map<int, std::shared_ptr<void>> s_optimizers;
 
 // C++ interface
 
-template <typename ds_params_precision_t, typename ds_state_precision_t>
+template <typename ds_params_precision_t,
+          typename ds_state_precision_t,
+          typename ds_grad_precision_t>
 void Adam_Optimizer::Step_1(ds_params_precision_t* _params,
-                            ds_params_precision_t* grads,
+                            ds_grad_precision_t* grads,
                             ds_state_precision_t* _exp_avg,
                             ds_state_precision_t* _exp_avg_sq,
                             size_t _param_size,
@@ -62,6 +64,7 @@ void Adam_Optimizer::Step_1(ds_params_precision_t* _params,
 #pragma omp parallel for if (parallel)
             for (size_t k = t; k < offset; k++) {
                 float grad = (float)grads[k];
+                if (_grad_scale != 1.0f) { grad = grad * _grad_scale; }
                 float param = (float)_params[k];
                 float momentum = _exp_avg[k];
                 float variance = _exp_avg_sq[k];
@@ -86,9 +89,11 @@ void Adam_Optimizer::Step_1(ds_params_precision_t* _params,
     }
 }
 
-template <typename ds_params_precision_t, typename ds_state_precision_t>
+template <typename ds_params_precision_t,
+          typename ds_state_precision_t,
+          typename ds_grad_precision_t>
 void Adam_Optimizer::Step_4(ds_params_precision_t* _params,
-                            ds_params_precision_t* grads,
+                            ds_grad_precision_t* grads,
                             ds_state_precision_t* _exp_avg,
                             ds_state_precision_t* _exp_avg_sq,
                             size_t _param_size,
@@ -155,9 +160,11 @@ int create_adam_optimizer(int optimizer_id,
     return 0;
 }
 
-template <typename ds_params_precision_t, typename ds_state_precision_t>
+template <typename ds_params_precision_t,
+          typename ds_state_precision_t,
+          typename ds_grad_precision_t>
 void Adam_Optimizer::Step_8(ds_params_precision_t* _params,
-                            ds_params_precision_t* grads,
+                            ds_grad_precision_t* grads,
                             ds_state_precision_t* _exp_avg,
                             ds_state_precision_t* _exp_avg_sq,
                             size_t _param_size,
@@ -180,7 +187,9 @@ void Adam_Optimizer::Step_8(ds_params_precision_t* _params,
                parallel);
 }
 
-template <typename ds_params_precision_t, typename ds_state_precision_t>
+template <typename ds_params_precision_t,
+          typename ds_state_precision_t,
+          typename ds_grad_precision_t>
 void step_invoker(std::shared_ptr<Adam_Optimizer> opt,
                   void* _params,
                   void* grads,
@@ -190,25 +199,28 @@ void step_invoker(std::shared_ptr<Adam_Optimizer> opt,
                   bool parallel)
 {
     opt->Step_8((ds_params_precision_t*)(_params),
-                (ds_params_precision_t*)(grads),
+                (ds_grad_precision_t*)(grads),
                 (ds_state_precision_t*)(_exp_avg),
                 (ds_state_precision_t*)(_exp_avg_sq),
                 _param_size,
                 parallel);
 }
 
 std::map<
-    std::tuple<c10::ScalarType, c10::ScalarType>,
+    std::tuple<c10::ScalarType, c10::ScalarType, c10::ScalarType>,
     std::function<void(std::shared_ptr<Adam_Optimizer>, void*, void*, void*, void*, size_t, bool)>>
     invokers;
 
 // Fill map with template functions for each type
-template <class ds_params_precision_t, class ds_state_precision_t>
+template <class ds_params_precision_t,
+          class ds_state_precision_t,
+          class ds_grad_precision_t = ds_params_precision_t>
 void create_invoker()
 {
     invokers[std::tuple(c10::CppTypeToScalarType<ds_params_precision_t>(),
-                        c10::CppTypeToScalarType<ds_state_precision_t>())] =
-        step_invoker<ds_params_precision_t, ds_state_precision_t>;
+                        c10::CppTypeToScalarType<ds_state_precision_t>(),
+                        c10::CppTypeToScalarType<ds_grad_precision_t>())] =
+        step_invoker<ds_params_precision_t, ds_state_precision_t, ds_grad_precision_t>;
 }
 struct InvokerInitializer {
     InvokerInitializer()
@@ -218,6 +230,8 @@ struct InvokerInitializer {
         create_invoker<c10::BFloat16, float>();
         create_invoker<c10::BFloat16, c10::BFloat16>();
         create_invoker<float, float>();
+        // fp32 weights on bf16 gradients: SuperOffload copies them to the CPU in bf16.
+        create_invoker<float, float, c10::BFloat16>();
     }
 } _invoker_initializer;
 
@@ -231,11 +245,13 @@ void invoke(std::shared_ptr<Adam_Optimizer> opt,
 {
     c10::ScalarType params_type = at::typeMetaToScalarType(params.options().dtype());
     c10::ScalarType state_type = at::typeMetaToScalarType(exp_avg.options().dtype());
+    c10::ScalarType grad_type = at::typeMetaToScalarType(grads.options().dtype());
 
-    auto it = invokers.find(std::tuple(params_type, state_type));
+    auto it = invokers.find(std::tuple(params_type, state_type, grad_type));
     if (it == invokers.end()) {
         throw std::runtime_error("Adam optimizer with param type "s + c10::toString(params_type) +
-                                 " and state type "s + c10::toString(state_type) +
+                                 ", state type "s + c10::toString(state_type) + " and grad type "s +
+                                 c10::toString(grad_type) +
                                  " is not supported on current hardware"s);
     }
 
@@ -259,7 +275,8 @@ int ds_adam_step(int optimizer_id,
                  torch::Tensor& params,
                  torch::Tensor& grads,
                  torch::Tensor& exp_avg,
-                 torch::Tensor& exp_avg_sq)
+                 torch::Tensor& exp_avg_sq,
+                 float grad_scale)
 {
     auto params_c = params.contiguous();
     auto grads_c = grads.contiguous();
@@ -269,13 +286,110 @@ int ds_adam_step(int optimizer_id,
     std::shared_ptr<Adam_Optimizer> opt =
         std::static_pointer_cast<Adam_Optimizer>(s_optimizers[optimizer_id]);
     opt->IncrementStep(step, beta1, beta2);
-    opt->update_state(lr, epsilon, weight_decay, bias_correction);
+    opt->update_state(lr, epsilon, weight_decay, bias_correction, grad_scale);
 
     invoke(opt, params_c, grads_c, exp_avg_c, exp_avg_sq_c, params_c.numel());
 
     return 0;
 }
 
+template <typename ds_grad_precision_t>
+void Adam_Optimizer::Rescale_Step(float* _params,
+                                  ds_grad_precision_t* grads,
+                                  float* _exp_avg,
+                                  float* _exp_avg_sq,
+                                  size_t _param_size,
+                                  float old_scale,
+                                  float new_scale)
+{
+    if (_weight_decay > 0 && !_adamw_mode) {
+        // Adam's L2 decay adds the params of before the step to the gradient, which are gone.
+        throw std::runtime_error(
+            "Adam rescale step supports AdamW only, or Adam without weight decay");
+    }
+    size_t rounded_size = 0;
+#if defined(__AVX512__) or defined(__AVX256__) or defined(__NEON__)
+    Rescale_Step_AVX<8>(
+        &rounded_size, _params, grads, _exp_avg, _exp_avg_sq, _param_size, old_scale, new_scale);
+#endif
+    if (_param_size <= rounded_size) return;
+    float momentum_shift = (1 - _betta1) * (new_scale - old_scale);
+    float variance_shift = (1 - _betta2) * (new_scale * new_scale - old_scale * old_scale);
+    float step_size = -1 * _alpha / _bias_correction1;
+#pragma omp parallel for
+    for (size_t k = rounded_size; k < _param_size; k++) {
+        float grad = (float)grads[k];
+        float momentum = _exp_avg[k];
+        float variance = _exp_avg_sq[k];
+        float param = _params[k];
+        param = momentum / (sqrt(variance) * _bias_correction2 + _eps) * -step_size + param;
+        momentum = grad * momentum_shift + momentum;
+        // std::max(x, 0) keeps a NaN x.
+        variance = std::max(grad * grad * variance_shift + variance, 0.0f);
+        param = momentum / (sqrt(variance) * _bias_correction2 + _eps) * step_size + param;
+        _params[k] = param;
+        _exp_avg[k] = momentum;
+        _exp_avg_sq[k] = variance;
+    }
+}
+
+int ds_adam_rescale_step(int optimizer_id,
+                         size_t step,
+                         float lr,
+                         float beta1,
+                         float beta2,
+                         float epsilon,
+                         float weight_decay,
+                         bool bias_correction,
+                         torch::Tensor& params,
+                         torch::Tensor& grads,
+                         torch::Tensor& exp_avg,
+                         torch::Tensor& exp_avg_sq,
+                         float old_scale,
+                         float new_scale)
+{
+    TORCH_CHECK(params.scalar_type() == torch::kFloat32 &&
+                    exp_avg.scalar_type() == torch::kFloat32 &&
+                    exp_avg_sq.scalar_type() == torch::kFloat32,
+                "Adam rescale step supports float32 params and states only");
+    TORCH_CHECK(grads.scalar_type() == torch::kFloat32 || grads.scalar_type() == torch::kBFloat16,
+                "Adam rescale step supports float32 and bfloat16 gradients only");
+    TORCH_CHECK(params.is_contiguous() && grads.is_contiguous() && exp_avg.is_contiguous() &&
+                    exp_avg_sq.is_contiguous(),
+                "Adam rescale step updates its tensors in place, so they must be contiguous");
+    TORCH_CHECK(grads.numel() == params.numel() && exp_avg.numel() == params.numel() &&
+                    exp_avg_sq.numel() == params.numel(),
+                "Adam rescale step needs as many gradients and states as params");
+
+    std::shared_ptr<Adam_Optimizer> opt =
+        std::static_pointer_cast<Adam_Optimizer>(s_optimizers[optimizer_id]);
+    // As the step being rescaled had this step and these betas, beta1^step and beta2^step stay.
+    opt->IncrementStep(step, beta1, beta2);
+    opt->update_state(lr, epsilon, weight_decay, bias_correction);
+
+    float* params_ptr = params.data_ptr<float>();
+    float* exp_avg_ptr = exp_avg.data_ptr<float>();
+    float* exp_avg_sq_ptr = exp_avg_sq.data_ptr<float>();
+    if (grads.scalar_type() == torch::kBFloat16) {
+        opt->Rescale_Step(params_ptr,
+                          (c10::BFloat16*)grads.data_ptr(),
+                          exp_avg_ptr,
+                          exp_avg_sq_ptr,
+                          params.numel(),
+                          old_scale,
+                          new_scale);
+    } else {
+        opt->Rescale_Step(params_ptr,
+                          grads.data_ptr<float>(),
+                          exp_avg_ptr,
+                          exp_avg_sq_ptr,
+                          params.numel(),
+                          old_scale,
+                          new_scale);
+    }
+    return 0;
+}
+
 void adamw_rollback_inplace(float* params,
                             const float* grads,
                             float* momentum,
@@ -318,7 +432,10 @@ void adamw_rollback_inplace(float* params,
 
         const float grad = grads[i];
         momentum[i] = (momentum[i] - one_minus_beta1 * grad) / beta1;
-        variance[i] = (variance[i] - one_minus_beta2 * grad * grad) / beta2;
+        // exp_avg_sq is a sum of squares, so it is never negative. The inverse can leave a tiny
+        // negative rounding residual (in about half of the elements that were 0 before the step),
+        // and sqrt() of it in the next step makes the parameter NaN. std::max(x, 0) keeps a NaN x.
+        variance[i] = std::max((variance[i] - one_minus_beta2 * grad * grad) / beta2, 0.0f);
     }
 
     --step_count;
@@ -601,9 +718,9 @@ class ZenFlowAdam {
             torch::Tensor& M = grp.exp_avg[now_state];
             torch::Tensor& V = grp.exp_avg_sq[now_state];
 
-            auto it = invokers.find(std::tuple(P.scalar_type(), M.scalar_type()));
+            auto it = invokers.find(std::tuple(P.scalar_type(), M.scalar_type(), G.scalar_type()));
             TORCH_CHECK(it != invokers.end(),
-                        "ZenFlowAdam: unsupported param/state dtype combination");
+                        "ZenFlowAdam: unsupported param/state/grad dtype combination");
             auto fn = it->second;
 
             char* pp = static_cast<char*>(P.data_ptr());
@@ -612,13 +729,14 @@ class ZenFlowAdam {
             char* vp = static_cast<char*>(V.data_ptr());
             char* sp = grp.stale.defined() ? static_cast<char*>(grp.stale.data_ptr()) : nullptr;
             const size_t pe = P.element_size();
+            const size_t ge = G.element_size();
             const size_t se = M.element_size();
             const size_t numel = P.numel();
 
             pool_->parallel_for(numel, kZenAdamAlign, [=](size_t b, size_t e) {
                 const size_t len = e - b;
                 // parallel=false: each pinned thread runs its slice serially.
-                fn(opt, pp + b * pe, gp + b * pe, mp + b * se, vp + b * se, len, false);
+                fn(opt, pp + b * pe, gp + b * ge, mp + b * se, vp + b * se, len, false);
                 if (sp) std::memcpy(sp + b * pe, pp + b * pe, len * pe);
             });
         }