Support L2 regularization and decoupled weight decay in rowwise adagrad (#718)

Summary:
Pull Request resolved: #718

Add two kinds of weight decay in rowwise adagrad:

L2 regularization:
```
g' = g + weight_decay * w
multiplier = lr / (sqrt(v) + eps)
w = w - lr * g' / (sqrt(v) + eps)
    = w - lr * g / (sqrt(v) + eps) - lr * weight_decay * w / (sqrt(v) + eps)
    = (1 - multiplier * weight_decay) * w - multiplier * g
```

Decoupled weight decay:
```
multiplier = lr / (sqrt(v) + eps)
w = w - lr * (g / (sqrt(v) + eps) + weight_decay * w)
    = w - lr * g / (sqrt(v) + eps) - lr * weight_decay * w
    = (1 - lr * weight_decay) * w - multiplier * g
```

Reviewed By: choudharydhruv

Differential Revision: D31285351

fbshipit-source-id: e361627f8426856021badef0410455e23620f21b

Author: Yuxi Hu

fbgemm_gpu/codegen/embedding_backward_code_generator.py

@@ -367,49 +367,90 @@ def table_info_precomputation(momentum_prefix: str = "momentum1") -> str:
 
 def rowwise_adagrad() -> None:
     split_weight_update = """
-      weight_new.fma_(grad, -multiplier);
+        weight_new.acc.x = correction * weight_new.acc.x - multiplier * grad.acc.x;
+        weight_new.acc.y = correction * weight_new.acc.y - multiplier * grad.acc.y;
+        weight_new.acc.z = correction * weight_new.acc.z - multiplier * grad.acc.z;
+        weight_new.acc.w = correction * weight_new.acc.w - multiplier * grad.acc.w;
     """
     split_precomputation = """
     acc_type<cache_t, true> g_local_sum_square = 0.0;
     #pragma unroll kMaxVecsPerThread
     for (int32_t i = 0;
         i < kMaxVecsPerThread && 4 * kWarpSize * i + threadIdx.x * 4 < D;
         ++i) {
-    g_local_sum_square += grad_sum[i].acc.x * grad_sum[i].acc.x +
-        grad_sum[i].acc.y * grad_sum[i].acc.y +
-        grad_sum[i].acc.z * grad_sum[i].acc.z +
-        grad_sum[i].acc.w * grad_sum[i].acc.w;
+        auto gx = grad_sum[i].acc.x;
+        auto gy = grad_sum[i].acc.y;
+        auto gz = grad_sum[i].acc.z;
+        auto gw = grad_sum[i].acc.w;
+        if (weight_decay_mode == 0) {
+            // L2 regularization
+            int32_t d = 4 * kWarpSize * i + threadIdx.x * 4;
+            Vec4T<acc_type<cache_t, true>> weight = weight_row_template.load(d, qparams_template);
+            gx += weight_decay * weight.acc.x;
+            gy += weight_decay * weight.acc.y;
+            gz += weight_decay * weight.acc.z;
+            gw += weight_decay * weight.acc.w;
+        }
+        g_local_sum_square += gx * gx + gy * gy + gz * gz + gw * gw;
     }
     const acc_type<cache_t, true> g_avg_square =
         warpReduceAllSum<acc_type<cache_t, true>>(g_local_sum_square) / D;
 
     acc_type<cache_t, true> multiplier;
+    acc_type<cache_t, true> correction = 1.0;
     if (threadIdx.x == 0) {
         acc_type<cache_t, true> new_sum_square_grads = momentum1[idx] + g_avg_square;
         momentum1[idx] = new_sum_square_grads;
         multiplier = learning_rate / (sqrtf(new_sum_square_grads) + eps);
+        if (weight_decay_mode == 0) {
+            // L2 regularization
+            correction = 1 - multiplier * weight_decay;
+        } else if (weight_decay_mode == 1){
+            // Decoupled weight decay
+            correction = 1 - learning_rate * weight_decay;
+        }
     }
     multiplier = __shfl_sync(0xFFFFFFFF, multiplier, 0);
+    correction = __shfl_sync(0xFFFFFFFF, correction, 0);
     """
     split_weight_update_cpu = """
         acc_type<scalar_t, true> g_local_sum_square = 0.0;
         for (int64_t d = 0; d < D; ++d) {
-            g_local_sum_square += grad_buffer[d] * grad_buffer[d];
+            auto grad = grad_buffer[d];
+            if (weight_decay_mode == 0) {
+                // L2 regularization
+                grad += weight_decay * host_weights_data[embedding_begin + d];
+            }
+            g_local_sum_square += grad * grad;
         }
         auto g_avg_square = g_local_sum_square / D;
         acc_type<scalar_t, true> new_sum_square_grads = momentum1_host[momentum1_offsets_data[feature_begin] + idx] + g_avg_square;
         momentum1_host[momentum1_offsets_data[feature_begin] + idx] = new_sum_square_grads;
         acc_type<scalar_t, true> multiplier;
         multiplier = learning_rate / (sqrtf(new_sum_square_grads) + eps);
+        acc_type<scalar_t, true> correction = 1.0;
+        if (weight_decay_mode == 0) {
+            // L2 regularization
+            correction = 1 - multiplier * weight_decay;
+        } else if (weight_decay_mode == 1) {
+            // Decoupled weight decay
+            correction = 1 - learning_rate * weight_decay;
+        }
         for (int64_t d = 0; d < D; ++d) {
-            host_weights_data[embedding_begin + d] -= grad_buffer[d] * multiplier;
+            host_weights_data[embedding_begin + d] = correction * host_weights_data[embedding_begin + d] - grad_buffer[d] * multiplier;
         }
     """
 
     generate(
         optimizer="rowwise_adagrad",
         args=make_args(
-            [(TENSOR, "momentum1"), (FLOAT, "eps"), (FLOAT, "learning_rate")]
+            [
+                (TENSOR, "momentum1"),
+                (FLOAT, "eps"),
+                (FLOAT, "learning_rate"),
+                (FLOAT, "weight_decay"),
+                (INT, "weight_decay_mode"),
+            ]
         ),
         split_precomputation=split_precomputation,
         split_weight_update=split_weight_update,
@@ -425,7 +466,13 @@ def rowwise_adagrad() -> None:
     generate(
         optimizer="approx_rowwise_adagrad",
         args=make_args(
-            [(TENSOR, "momentum1"), (FLOAT, "eps"), (FLOAT, "learning_rate")]
+            [
+                (TENSOR, "momentum1"),
+                (FLOAT, "eps"),
+                (FLOAT, "learning_rate"),
+                (FLOAT, "weight_decay"),
+                (INT, "weight_decay_mode"),
+            ]
         ),
         split_precomputation=split_precomputation,
         split_weight_update=approx_split_weight_update,

fbgemm_gpu/codegen/embedding_backward_split_template.cu

@@ -292,8 +292,6 @@ split_embedding_backward_codegen_{{ optimizer }}_{{ wdesc }}_kernel_cta_per_row_
             }
             {% endfor %}
 
-            {{ split_precomputation }}
-
             struct SharedMemory<Vec4T<acc_type<cache_t, true>>> weight_update_buffer;
             Vec4T<acc_type<cache_t, true>>* shared_weight_update_row = weight_update_buffer.getPointer();
 
@@ -315,6 +313,9 @@ split_embedding_backward_codegen_{{ optimizer }}_{{ wdesc }}_kernel_cta_per_row_
             if (std::is_same<emb_t, uint8_t>::value && !cache_weights) {
                 qparams_template = weight_row_template.load_qparams();
             }
+
+            {{ split_precomputation }}
+
             float2 qparams_new;
             #pragma unroll kMaxVecsPerThread
             for (int32_t i = 0;
@@ -506,7 +507,6 @@ split_embedding_backward_codegen_{{ optimizer }}_{{ wdesc }}_kernel_warp_per_row
     }
     {% endfor %}
 
-    {{ split_precomputation }}
     struct SharedMemory<Vec4T<acc_type<cache_t, true>>> weight_update_buffer;
     Vec4T<acc_type<cache_t, true>>* shared_weight_update_row = weight_update_buffer.getPointer();
     auto weight_row_template = WeightRow<emb_t, cache_t, acc_type<cache_t, true>>(weights, cache_weights, D, nullptr);
@@ -526,6 +526,9 @@ split_embedding_backward_codegen_{{ optimizer }}_{{ wdesc }}_kernel_warp_per_row
     if (std::is_same<emb_t, uint8_t>::value && !cache_weights){
         qparams_template = weight_row_template.load_qparams();
     }
+
+    {{ split_precomputation }}
+
     float2 qparams_new;
     #pragma unroll kMaxVecsPerThread
     for (int32_t i = 0;

fbgemm_gpu/codegen/lookup_args.py

@@ -39,6 +39,7 @@ class OptimizerArgs(NamedTuple):
     beta1: float
     beta2: float
     weight_decay: float
+    weight_decay_mode: int
     eta: float
     momentum: float
 

fbgemm_gpu/codegen/split_embedding_codegen_lookup_invoker.template

@@ -71,6 +71,9 @@ def invoke(
             {% if "weight_decay" in args.split_function_arg_names %}
             weight_decay=optimizer_args.weight_decay,
             {% endif %}
+            {% if "weight_decay_mode" in args.split_function_arg_names %}
+            weight_decay_mode=optimizer_args.weight_decay_mode,
+            {% endif %}
             {% if "eta" in args.split_function_arg_names %}
             eta=optimizer_args.eta,
             {% endif %}
@@ -135,6 +138,9 @@ def invoke(
             {% if "weight_decay" in args.split_function_arg_names %}
             weight_decay=optimizer_args.weight_decay,
             {% endif %}
+            {% if "weight_decay_mode" in args.split_function_arg_names %}
+            weight_decay_mode=optimizer_args.weight_decay_mode,
+            {% endif %}
             {% if "eta" in args.split_function_arg_names %}
             eta=optimizer_args.eta,
             {% endif %}

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops.py

@@ -65,6 +65,11 @@ class BoundsCheckMode(enum.IntEnum):
     NONE = 3
 
 
+class WeightDecayMode(enum.IntEnum):
+    L2 = 0
+    DECOUPLE = 1
+
+
 RecordCacheMetrics: NamedTuple = NamedTuple(
     "RecordCacheMetrics",
     [("record_cache_miss_counter", bool), ("record_tablewise_cache_miss", bool)],
@@ -201,7 +206,8 @@ def __init__(  # noqa C901
         learning_rate: float = 0.01,
         eps: float = 1.0e-8,  # used by Adagrad, LAMB, and Adam
         momentum: float = 0.9,  # used by LARS-SGD
-        weight_decay: float = 0.0,  # used by LARS-SGD, LAMB, and ADAM
+        weight_decay: float = 0.0,  # used by LARS-SGD, LAMB, Adagrad, and ADAM
+        weight_decay_mode: WeightDecayMode = WeightDecayMode.L2,  # used by Adagrad
         eta: float = 0.001,  # used by LARS-SGD,
         beta1: float = 0.9,  # used by LAMB and ADAM
         beta2: float = 0.999,  # used by LAMB and ADAM
@@ -357,6 +363,7 @@ def __init__(  # noqa C901
             beta1=beta1,
             beta2=beta2,
             weight_decay=weight_decay,
+            weight_decay_mode=weight_decay_mode.value,
             eta=eta,
             momentum=momentum,
         )
@@ -493,7 +500,7 @@ def __init__(  # noqa C901
             dtype=cache_embedding_dtype,
         )
 
-        logging.debug(
+        logging.info(
             f"Using fused {optimizer} with optimizer_args={self.optimizer_args}"
         )
 

fbgemm_gpu/test/split_table_batched_embeddings_test.py

@@ -16,10 +16,11 @@
 import numpy as np
 import torch
 from fbgemm_gpu.split_table_batched_embeddings_ops import (
+    BoundsCheckMode,
     OptimType,
-    SparseType,
     RecordCacheMetrics,
-    BoundsCheckMode,
+    SparseType,
+    WeightDecayMode,
 )
 from hypothesis import HealthCheck, Verbosity, assume, given, settings
 from torch import Tensor
@@ -1409,6 +1410,7 @@ def execute_backward_optimizers_(  # noqa C901
         long_segments: bool,
         pooling_mode: split_table_batched_embeddings_ops.PoolingMode,
         use_cpu: bool,
+        weight_decay_mode: WeightDecayMode = WeightDecayMode.L2,
     ) -> None:
         # NOTE: limit (T * B * L * D) to avoid timeout for CPU version!
         assume(not use_cpu or T * B * L * D <= 2048)
@@ -1534,6 +1536,10 @@ def execute_backward_optimizers_(  # noqa C901
         if optimizer in (OptimType.EXACT_ROWWISE_ADAGRAD, OptimType.EXACT_ADAGRAD):
             optimizer_kwargs["eps"] = eps
 
+        if optimizer == OptimType.EXACT_ROWWISE_ADAGRAD:
+            optimizer_kwargs["weight_decay"] = weight_decay
+            optimizer_kwargs["weight_decay_mode"] = weight_decay_mode
+
         if optimizer in (OptimType.PARTIAL_ROWWISE_ADAM, OptimType.ADAM):
             optimizer_kwargs["eps"] = eps
             optimizer_kwargs["beta1"] = beta1
@@ -1591,25 +1597,30 @@ def execute_backward_optimizers_(  # noqa C901
                 # to_dense in GPU is non-deterministic due to atmomics used in
                 # coalescing and floating point non-associativity.
                 dense_cpu_grad = bs[t].weight.grad.cpu().to_dense()
+                if rowwise and not use_cpu and weight_decay_mode == WeightDecayMode.L2:
+                    # NOTE: CPU code path (https://fburl.com/diffusion/rte4cu6c) is not executed in unit test.
+                    dense_cpu_grad += weight_decay * bs[t].weight.cpu()
                 m1_ref = (
                     dense_cpu_grad.pow(2)
                     if not rowwise
                     else dense_cpu_grad.pow(2).mean(dim=1)
                 )
                 torch.testing.assert_allclose(
-                    m1.float().cpu(), m1_ref.float(), atol=1.0e-4, rtol=1.0e-4
+                    m1.float().index_select(dim=0, index=x[t].view(-1)).cpu(),
+                    m1_ref.float().index_select(dim=0, index=x[t].view(-1).cpu()),
+                    atol=1.0e-4,
+                    rtol=1.0e-4
                 )
                 weights_new = split_weights[t]
-                weights_ref = bs[t].weight.cpu() - lr * dense_cpu_grad / (
-                    torch.sqrt(
-                        m1_ref if not rowwise else m1_ref.view(m1_ref.numel(), 1)
-                    )
-                    + eps
-                )
+                denom = torch.sqrt(m1_ref if not rowwise else m1_ref.view(m1_ref.numel(), 1)) + eps
+                if rowwise and not use_cpu and weight_decay_mode == WeightDecayMode.DECOUPLE:
+                    weights_ref = bs[t].weight.cpu() - lr * (dense_cpu_grad / denom + weight_decay * bs[t].weight.cpu())
+                else:
+                    weights_ref = bs[t].weight.cpu() - lr * dense_cpu_grad / denom
                 # TODO: why is tolerance off here?
                 torch.testing.assert_allclose(
-                    weights_new.float().cpu(),
-                    weights_ref.float(),
+                    weights_new.index_select(dim=0, index=x[t].view(-1)).cpu(),
+                    weights_ref.index_select(dim=0, index=x[t].view(-1).cpu()),
                     atol=1.0e-2,
                     rtol=1.0e-2,
                 )
@@ -1793,6 +1804,12 @@ def test_backward_optimizers_adam(  # noqa C901
             ]
         ),
         use_cpu=st.booleans() if torch.cuda.is_available() else st.just(True),
+        weight_decay_mode=st.sampled_from(
+            [
+                WeightDecayMode.L2,
+                WeightDecayMode.DECOUPLE,
+            ]
+        ),
     )
     @settings(
         verbosity=Verbosity.verbose,
@@ -1813,9 +1830,11 @@ def test_backward_optimizers_adagrad(  # noqa C901
         long_segments: bool,
         pooling_mode: split_table_batched_embeddings_ops.PoolingMode,
         use_cpu: bool,
+        weight_decay_mode: WeightDecayMode,
     ) -> None:
         self.execute_backward_optimizers_(T, D, B, log_E, L, weighted,
-            mixed, optimizer, long_segments, pooling_mode, use_cpu)
+            mixed, optimizer, long_segments, pooling_mode, use_cpu,
+            weight_decay_mode)
 
     @given(
         T=st.integers(min_value=1, max_value=5),