reduce optim memory allocation, finish adafactor

Author: ITcarrot

roundpipe/optim/csrc/adadelta.cpp

@@ -37,17 +37,7 @@ void adadelta_kernel(bool current_bool, Args... args) {
 void adadelta(vector<Tensor> params, vector<Tensor> grads, vector<Tensor> square_avg,
               vector<Tensor> acc_delta, vector<Tensor> state_steps, double lr,
               double rho, double eps, double weight_decay, bool maximize) {
-    vector<int64_t> numel(params.size());
-    vector<float *> params_ptr(params.size());
-    vector<const float *> grads_ptr(params.size());
-    vector<float *> square_avg_ptr(params.size());
-    vector<float *> acc_delta_ptr(params.size());
     for (size_t i = 0; i < params.size(); ++i) {
-        numel[i] = params[i].numel();
-        params_ptr[i] = params[i].mutable_data_ptr<float>();
-        grads_ptr[i] = grads[i].const_data_ptr<float>();
-        square_avg_ptr[i] = square_avg[i].mutable_data_ptr<float>();
-        acc_delta_ptr[i] = acc_delta[i].mutable_data_ptr<float>();
         // The step is tracked for state-dict compatibility but does not enter the math.
         state_steps[i].add_(1);
     }
@@ -56,12 +46,16 @@ void adadelta(vector<Tensor> params, vector<Tensor> grads, vector<Tensor> square
         int rank = omp_get_thread_num();
         int nthreads = omp_get_num_threads();
         for (size_t i = 0; i < params.size(); ++i) {
-            int64_t block_size = numel[i] / nthreads + (rank < (numel[i] % nthreads));
+            int64_t numel = params[i].numel();
+            int64_t block_size = numel / nthreads + (rank < (numel % nthreads));
             int64_t offset =
-                (numel[i] / nthreads) * rank + min<int64_t>(rank, numel[i] % nthreads);
-            adadelta_kernel(maximize, weight_decay == 0.0, params_ptr[i] + offset,
-                            grads_ptr[i] + offset, square_avg_ptr[i] + offset,
-                            acc_delta_ptr[i] + offset, lr, rho, eps, weight_decay,
+                (numel / nthreads) * rank + min<int64_t>(rank, numel % nthreads);
+            float *params_ptr = params[i].mutable_data_ptr<float>() + offset;
+            const float *grads_ptr = grads[i].const_data_ptr<float>() + offset;
+            float *square_avg_ptr = square_avg[i].mutable_data_ptr<float>() + offset;
+            float *acc_delta_ptr = acc_delta[i].mutable_data_ptr<float>() + offset;
+            adadelta_kernel(maximize, weight_decay == 0.0, params_ptr, grads_ptr,
+                            square_avg_ptr, acc_delta_ptr, lr, rho, eps, weight_decay,
                             block_size);
         }
     }

roundpipe/optim/csrc/adafactor.cpp

@@ -17,41 +17,38 @@ using namespace torch;
 template <bool maximize, bool zero_weight_decay>
 static void adafactor_matrix(float *__restrict param, const float *__restrict grad,
                              float *__restrict row_var, float *__restrict col_var,
-                             int64_t B, int64_t R, int64_t C, float f_omb2,
-                             double rho_t, float f_eps1, float f_eps1sq, float f_eps2,
-                             float f_d, float f_wd_factor) {
+                             float *__restrict cs, float *__restrict inv_rvm,
+                             float *__restrict upd, int64_t B, int64_t R, int64_t C,
+                             float f_omb2, double rho_t, double eps1, double eps2,
+                             double d, float f_wd_factor) {
     const int64_t RC = R * C;
     const int64_t numel = B * RC;
-    if (numel == 0) {
-        return; // A zero-sized dim leaves the empty parameter untouched.
-    }
     const float gsign = maximize ? -1.0f : 1.0f;
-    const float f_inv_C = 1.0f / (float)C;
-    const float f_inv_R = 1.0f / (float)R;
+    float f_eps1 = eps1;
+    float f_eps1sq = eps1 * eps1;
+    const float f_inv_C = 1.0 / C;
+    const float f_inv_R = 1.0 / R;
     const int64_t BR = B * R;
     const int64_t BC = B * C;
 
     // --- RMS(param) for the step size, computed before the weight decay. ---
-    double param_sq = 0.0;
+    float param_sq = 0.0;
 #pragma omp parallel for simd reduction(+ : param_sq) schedule(static)
     for (int64_t o = 0; o < numel; ++o) {
-        param_sq += (double)param[o] * param[o];
+        param_sq += param[o] * param[o];
     }
-    double rms_param = sqrt(param_sq / (double)numel);
-    float f_alpha = (float)(max((double)f_eps2, rms_param) * rho_t);
+    double rms_param = sqrt((double)param_sq / numel);
+    double alpha = max(eps2, rms_param) * rho_t;
 
     // --- Row sums of grad^2 (and the row_var lerp) in parallel over rows; the
     // column sums are a reduction across rows, done with an OpenMP array-section
     // reduction so each thread accumulates its own colsum copy. ---
-    vector<float> colsum((size_t)BC, 0.0f);
-    float *__restrict cs = colsum.data();
 #pragma omp parallel for reduction(+ : cs[ : BC]) schedule(static)
     for (int64_t br = 0; br < BR; ++br) {
         const int64_t b = br / R;
         const float *__restrict grow = grad + br * C;
         float *__restrict csb = cs + b * C;
         float racc = 0.0f;
-#pragma omp simd reduction(+ : racc)
         for (int64_t c = 0; c < C; ++c) {
             float g2 = grow[c] * grow[c];
             racc += g2;
@@ -61,50 +58,51 @@ static void adafactor_matrix(float *__restrict param, const float *__restrict gr
     }
 
     // --- col_var lerp (parallel over columns) and the per-matrix mean of row_var
-    // that normalizes the outer product. B is tiny, so inv_rvm is serial. ---
-#pragma omp parallel for simd schedule(static)
-    for (int64_t bc = 0; bc < BC; ++bc) {
-        col_var[bc] += f_omb2 * (cs[bc] * f_inv_R - col_var[bc]);
-    }
-    vector<float> inv_rvm(B);
-    for (int64_t b = 0; b < B; ++b) {
-        const float *__restrict rv = row_var + b * R;
-        double rv_sum = 0.0;
-        for (int64_t r = 0; r < R; ++r) {
-            rv_sum += rv[r];
+    // that normalizes the outer product. ---
+#pragma omp parallel
+    {
+#pragma omp for simd schedule(static) nowait
+        for (int64_t bc = 0; bc < BC; ++bc) {
+            col_var[bc] += f_omb2 * (cs[bc] * f_inv_R - col_var[bc]);
+        }
+#pragma omp for schedule(static)
+        for (int64_t b = 0; b < B; ++b) {
+            const float *__restrict rv = row_var + b * R;
+            float rv_sum = 0.0;
+#pragma omp simd reduction(+ : rv_sum)
+            for (int64_t r = 0; r < R; ++r) {
+                rv_sum += rv[r];
+            }
+            inv_rvm[b] = 1.0f / max(rv_sum * f_inv_R, f_eps1);
         }
-        inv_rvm[b] = 1.0f / max((float)(rv_sum * (double)f_inv_R), f_eps1);
     }
 
     // --- Build the update U = grad / sqrt(max(var_estimate, eps1^2)) and its RMS,
     // parallel over rows. var_estimate[b,r,c] = row_var[b,r] * col_var[b,c] *
     // inv_rvm[b].
-    vector<float> update(numel);
-    float *__restrict upd = update.data();
-    double update_sq = 0.0;
+    float update_sq = 0.0;
 #pragma omp parallel for reduction(+ : update_sq) schedule(static)
     for (int64_t br = 0; br < BR; ++br) {
         const int64_t b = br / R;
         float pre = row_var[br] * inv_rvm[b];
         const float *__restrict grow = grad + br * C;
         const float *__restrict cvb = col_var + b * C;
         float *__restrict urow = upd + br * C;
-        double usq = 0.0;
-#pragma omp simd reduction(+ : usq)
+        float usq = 0.0;
         for (int64_t c = 0; c < C; ++c) {
             float ve = pre * cvb[c];
             ve = ve > f_eps1sq ? ve : f_eps1sq;
             float u = gsign * grow[c] / sqrt(ve);
             urow[c] = u;
-            usq += (double)u * u;
+            usq += u * u;
         }
         update_sq += usq;
     }
 
     // --- Clip by RMS(update)/d and apply the (decoupled) weight decay in one pass. ---
-    double rms_update = sqrt(update_sq / (double)numel);
-    float f_denom = (float)max(1.0, rms_update / (double)f_d);
-    float f_coeff = -f_alpha / f_denom;
+    double rms_update = sqrt((double)update_sq / numel);
+    double denom = max(1.0, rms_update / d);
+    float f_coeff = -alpha / denom;
 #pragma omp parallel for simd schedule(static)
     for (int64_t o = 0; o < numel; ++o) {
         param[o] =
@@ -117,76 +115,54 @@ static void adafactor_matrix(float *__restrict param, const float *__restrict gr
 // same elementwise-plus-two-reductions shape as the other RoundPipe optimizers.
 template <bool maximize, bool zero_weight_decay>
 static void adafactor_vector(float *__restrict param, const float *__restrict grad,
-                             float *__restrict variance, int64_t numel, float f_omb2,
-                             double rho_t, float f_eps1sq, float f_eps2, float f_d,
-                             float f_wd_factor) {
+                             float *__restrict variance, float *__restrict upd,
+                             int64_t numel, float f_omb2, double rho_t, double eps1,
+                             double eps2, double d, float f_wd_factor) {
     if (numel == 0) {
         return;
     }
+    float f_eps1sq = eps1 * eps1;
     const float gsign = maximize ? -1.0f : 1.0f;
 
-    double param_sq = 0.0;
-#pragma omp parallel for simd reduction(+ : param_sq) schedule(static)
-    for (int64_t o = 0; o < numel; ++o) {
-        param_sq += (double)param[o] * param[o];
-    }
-    double rms_param = sqrt(param_sq / (double)numel);
-    float f_alpha = (float)(max((double)f_eps2, rms_param) * rho_t);
-
-    vector<float> update(numel);
-    float *__restrict upd = update.data();
-    double update_sq = 0.0;
-#pragma omp parallel for simd reduction(+ : update_sq) schedule(static)
+    float param_sq = 0.0, update_sq = 0.0;
+#pragma omp parallel for simd reduction(+ : param_sq, update_sq) schedule(static)
     for (int64_t o = 0; o < numel; ++o) {
+        param_sq += param[o] * param[o];
         float g = gsign * grad[o];
         variance[o] += f_omb2 * (g * g - variance[o]);
-        float ve = variance[o] > f_eps1sq ? variance[o] : f_eps1sq;
-        float u = g / sqrt(ve);
+        float u = g / sqrt(variance[o] > f_eps1sq ? variance[o] : f_eps1sq);
         upd[o] = u;
-        update_sq += (double)u * u;
+        update_sq += u * u;
     }
 
-    double rms_update = sqrt(update_sq / (double)numel);
-    float f_denom = (float)max(1.0, rms_update / (double)f_d);
-    float f_coeff = -f_alpha / f_denom;
+    double rms_param = sqrt((double)param_sq / numel);
+    double alpha = max(eps2, rms_param) * rho_t;
+    double rms_update = sqrt((double)update_sq / numel);
+    double denom = max(1.0, rms_update / d);
+    float f_coeff = -alpha / denom;
 #pragma omp parallel for simd schedule(static)
     for (int64_t o = 0; o < numel; ++o) {
         param[o] =
             (zero_weight_decay ? param[o] : param[o] * f_wd_factor) + f_coeff * upd[o];
     }
 }
 
-// Step every parameter with the maximize / weight-decay branches resolved at
-// compile time, matching the templated style of the other RoundPipe kernels.
-template <bool maximize, bool zero_weight_decay>
-static void adafactor_impl(const vector<float *> &mP, const vector<const float *> &mG,
-                           const vector<float *> &mRV, const vector<float *> &mCV,
-                           const vector<int64_t> &mB, const vector<int64_t> &mR,
-                           const vector<int64_t> &mC, const vector<double> &m_omb2,
-                           const vector<double> &m_rho, const vector<float *> &vP,
-                           const vector<const float *> &vG, const vector<float *> &vVar,
-                           const vector<int64_t> &vN, const vector<double> &v_omb2,
-                           const vector<double> &v_rho, float f_eps1, float f_eps1sq,
-                           float f_eps2, float f_d, float f_wd) {
-    for (size_t i = 0; i < mP.size(); ++i) {
-        adafactor_matrix<maximize, zero_weight_decay>(
-            mP[i], mG[i], mRV[i], mCV[i], mB[i], mR[i], mC[i], (float)m_omb2[i],
-            m_rho[i], f_eps1, f_eps1sq, f_eps2, f_d, f_wd);
-    }
-    for (size_t i = 0; i < vP.size(); ++i) {
-        adafactor_vector<maximize, zero_weight_decay>(vP[i], vG[i], vVar[i], vN[i],
-                                                      (float)v_omb2[i], v_rho[i],
-                                                      f_eps1sq, f_eps2, f_d, f_wd);
+// Unpack the runtime booleans into template parameters.
+template <bool... FixedBools, typename... Args>
+static void adafactor_matrix(bool current_bool, Args &&...args) {
+    if (current_bool) {
+        adafactor_matrix<FixedBools..., true>(std::forward<Args>(args)...);
+    } else {
+        adafactor_matrix<FixedBools..., false>(std::forward<Args>(args)...);
     }
 }
 
-// Unpack the runtime booleans into template parameters.
 template <bool... FixedBools, typename... Args>
-static void adafactor_impl(bool current_bool, Args &&...args) {
+static void adafactor_vector(bool current_bool, Args &&...args) {
     if (current_bool) {
-        adafactor_impl<FixedBools..., true>(std::forward<Args>(args)...);
+        adafactor_vector<FixedBools..., true>(std::forward<Args>(args)...);
     } else {
-        adafactor_impl<FixedBools..., false>(std::forward<Args>(args)...);
+        adafactor_vector<FixedBools..., false>(std::forward<Args>(args)...);
     }
 }
 
@@ -200,52 +176,43 @@ void adafactor(vector<Tensor> params, vector<Tensor> grads, vector<Tensor> row_v
                vector<Tensor> col_vars, vector<Tensor> variances,
                vector<Tensor> state_steps, double lr, double beta2_decay, double eps1,
                double eps2, double d, double weight_decay, bool maximize) {
-    vector<float *> mP, mRV, mCV, vP, vVar;
-    vector<const float *> mG, vG;
-    vector<int64_t> mB, mR, mC, vN;
-    vector<double> m_omb2, m_rho, v_omb2, v_rho;
+    float f_wd_factor = 1.0 - (lr * weight_decay);
+    bool zero_weight_decay = weight_decay == 0.0;
+    vector<float> colsum;
+    vector<float> inv_rvm;
+    vector<float> update;
 
     for (size_t i = 0; i < params.size(); ++i) {
+        Tensor &p = params[i];
         state_steps[i].add_(1);
+        if (p.numel() == 0) {
+            continue; // A zero-sized parameter is a no-op.
+        }
         double s = state_steps[i].item<double>();
         double omb2 = pow(s, beta2_decay);
         double rho = min(lr, 1.0 / sqrt(s));
-        Tensor p = params[i];
         if (p.dim() > 1) {
             int64_t nd = p.dim();
             int64_t Cc = p.size(nd - 1);
             int64_t Rr = p.size(nd - 2);
             int64_t rc = Rr * Cc;
-            // A zero-sized last dim makes rc == 0; guard the integer division.
-            mB.push_back(rc == 0 ? 0 : p.numel() / rc);
-            mR.push_back(Rr);
-            mC.push_back(Cc);
-            mP.push_back(params[i].mutable_data_ptr<float>());
-            mG.push_back(grads[i].const_data_ptr<float>());
-            mRV.push_back(row_vars[i].mutable_data_ptr<float>());
-            mCV.push_back(col_vars[i].mutable_data_ptr<float>());
-            m_omb2.push_back(omb2);
-            m_rho.push_back(rho);
+            int64_t B = p.numel() / rc;
+            colsum.assign(B * Cc, 0.0f);
+            inv_rvm.reserve(B);
+            update.reserve(p.numel());
+            adafactor_matrix(
+                maximize, zero_weight_decay, p.mutable_data_ptr<float>(),
+                grads[i].const_data_ptr<float>(), row_vars[i].mutable_data_ptr<float>(),
+                col_vars[i].mutable_data_ptr<float>(), colsum.data(), inv_rvm.data(),
+                update.data(), B, Rr, Cc, omb2, rho, eps1, eps2, d, f_wd_factor);
         } else {
-            vN.push_back(p.numel());
-            vP.push_back(params[i].mutable_data_ptr<float>());
-            vG.push_back(grads[i].const_data_ptr<float>());
-            vVar.push_back(variances[i].mutable_data_ptr<float>());
-            v_omb2.push_back(omb2);
-            v_rho.push_back(rho);
+            update.reserve(p.numel());
+            adafactor_vector(maximize, zero_weight_decay, p.mutable_data_ptr<float>(),
+                             grads[i].const_data_ptr<float>(),
+                             variances[i].mutable_data_ptr<float>(), update.data(),
+                             p.numel(), omb2, rho, eps1, eps2, d, f_wd_factor);
         }
     }
-
-    float f_eps1 = eps1;
-    float f_eps1sq = (float)eps1 * (float)eps1;
-    float f_eps2 = eps2;
-    float f_d = d;
-    float f_wd = 1.0f - (float)(lr * weight_decay);
-    bool zero_weight_decay = weight_decay == 0.0;
-
-    adafactor_impl(maximize, zero_weight_decay, mP, mG, mRV, mCV, mB, mR, mC, m_omb2,
-                   m_rho, vP, vG, vVar, vN, v_omb2, v_rho, f_eps1, f_eps1sq, f_eps2,
-                   f_d, f_wd);
 }
 
 #if PYBIND11_VERSION_HEX >= 0x020D0000 // pybind11 >= 2.13

roundpipe/optim/csrc/adagrad.cpp

@@ -31,30 +31,25 @@ void adagrad_kernel(bool current_bool, Args... args) {
 void adagrad(vector<Tensor> params, vector<Tensor> grads, vector<Tensor> state_sum,
              vector<Tensor> state_steps, double lr, double lr_decay, double eps,
              double weight_decay, bool maximize) {
-    vector<int64_t> numel(params.size());
-    vector<float *> params_ptr(params.size());
-    vector<const float *> grads_ptr(params.size());
-    vector<float *> state_sum_ptr(params.size());
     for (size_t i = 0; i < params.size(); ++i) {
-        numel[i] = params[i].numel();
-        params_ptr[i] = params[i].mutable_data_ptr<float>();
-        grads_ptr[i] = grads[i].const_data_ptr<float>();
-        state_sum_ptr[i] = state_sum[i].mutable_data_ptr<float>();
         state_steps[i].add_(1);
     }
 #pragma omp parallel
     {
         int rank = omp_get_thread_num();
         int nthreads = omp_get_num_threads();
         for (size_t i = 0; i < params.size(); ++i) {
-            int64_t block_size = numel[i] / nthreads + (rank < (numel[i] % nthreads));
+            int64_t numel = params[i].numel();
+            int64_t block_size = numel / nthreads + (rank < (numel % nthreads));
             int64_t offset =
-                (numel[i] / nthreads) * rank + min<int64_t>(rank, numel[i] % nthreads);
+                (numel / nthreads) * rank + min<int64_t>(rank, numel % nthreads);
             double step = state_steps[i].item<double>();
             double clr = lr / (1.0 + (step - 1.0) * lr_decay);
-            adagrad_kernel(maximize, weight_decay == 0.0, params_ptr[i] + offset,
-                           grads_ptr[i] + offset, state_sum_ptr[i] + offset, clr, eps,
-                           weight_decay, block_size);
+            float *params_ptr = params[i].mutable_data_ptr<float>() + offset;
+            const float *grads_ptr = grads[i].const_data_ptr<float>() + offset;
+            float *state_sum_ptr = state_sum[i].mutable_data_ptr<float>() + offset;
+            adagrad_kernel(maximize, weight_decay == 0.0, params_ptr, grads_ptr,
+                           state_sum_ptr, clr, eps, weight_decay, block_size);
         }
     }
 }