Repository: sql-machine-learning/elasticdl

Issue:
Support staleness-aware async-SGD
According to [design doc](https://github.com/sql-machine-learning/elasticdl/blob/develop/elasticdl/doc/async_sgd_design.md#change-in-ps)

To support Staleness-aware asychronous SGD, PS need to modulate the learning rate in the optimizer with the staleness value. PS may have multiple threads running concurrently for model updates with the same optimizer instance. Thus, we cannot modify the learning rate in the optimizer instance. We may modify the learning rate as a callable method, and use a thread local storage threading.local() to store the staleness. The callable method uses the staleness value to modulate the learning rate. The optimizer will call this callable method when it reads the learning rate hyperparameter.





Discussion:

Resolving patch:
diff --git a/elasticdl/python/master/lr_modulation.py b/elasticdl/python/master/lr_modulation.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/master/lr_modulation.py
@@ -0,0 +1,43 @@
+import threading
+
+
+class LearningRateModulation:
+    """
+    Modify learning rate with a multiplier.
+    Support concurrent usage by using thread local storage.
+    Arguments
+      learning rate: can be a value or a callable.
+    """
+
+    def __init__(self, learning_rate):
+        self._learning_rate = learning_rate
+        self._tls = threading.local()
+        self._tls.multiplier = 1
+
+    def set_multiplier(self, multiplier):
+        self._tls.multiplier = multiplier
+
+    def get_learning_rate(self):
+        lr = self._learning_rate
+        if callable(lr):
+            lr = lr()
+        lr *= self._tls.multiplier
+        return lr
+
+
+def add_lr_modulation_to_optimizer(optimizer):
+    """
+    Add lr modulation feature in optimizer
+    Argument:
+      optimizer: the optimizer to add lr modulation feature
+    Return:
+      LearningRateModulation instance
+    """
+    # Get learning rate from optimizer
+    learning_rate = optimizer._hyper["learning_rate"]
+
+    # Replace the learning rate in optimizer with a calllable
+    lr_modulation = LearningRateModulation(learning_rate)
+    optimizer.learning_rate = lr_modulation.get_learning_rate
+
+    return lr_modulation
diff --git a/elasticdl/python/master/servicer.py b/elasticdl/python/master/servicer.py
--- a/elasticdl/python/master/servicer.py
+++ b/elasticdl/python/master/servicer.py
@@ -16,6 +16,9 @@
 from elasticdl.python.elasticdl.layers.embedding import Embedding
 from elasticdl.python.master.checkpoint_service import CheckpointService
 from elasticdl.python.master.embedding_service import EmbeddingService
+from elasticdl.python.master.lr_modulation import (
+    add_lr_modulation_to_optimizer,
+)
 from elasticdl.python.master.optimizer_wrapper import OptimizerWrapper
 
 
@@ -48,6 +51,7 @@ def __init__(
         self._grad_n = 0
         self._minibatch_size = minibatch_size
         self._use_async = use_async
+        self._lr_staleness_modulation = lr_staleness_modulation
 
         # A <string, tf.ResourceVariable> map. We use tf.ResourceVariable
         # instead ndarray to avoid copying and conversion when calling
@@ -76,6 +80,12 @@ def set_model_var(self, name, value):
             value, name=MasterServicer.var_name_encode(name)
         )
 
+    def _modulate_lr_if_needed(self, opt):
+        if self._use_async and self._lr_staleness_modulation:
+            self._lr_modulation = add_lr_modulation_to_optimizer(opt)
+        else:
+            self._lr_modulation = None
+
     def _init_model_from_var_list(self, var_list):
         for var in var_list:
             self.set_model_var(var.name, var.numpy())
@@ -104,6 +114,7 @@ def _init_optimizer(
     ):
         # `embedding_service_endpoint` is not None means ElasticDL embedding
         # layers are used
+        self._modulate_lr_if_needed(opt)
         if embedding_service_endpoint:
             return OptimizerWrapper(
                 opt, embedding_service_endpoint, embedding_dims, use_async
@@ -392,6 +403,10 @@ def ReportGradient(self, request, _):
                     self._gradient_sum[k] = v
 
             self._grad_n += 1
+            # staleness-aware learning rate modulation
+            if self._lr_modulation:
+                staleness = max(1, self._version - request.model_version)
+                self._lr_modulation.set_multiplier(1.0 / staleness)
             if self._use_async or self._grad_n >= self._grad_to_wait:
                 self._update_model()
                 self._update_evaluation()


Test patch:
diff --git a/elasticdl/python/tests/staleness_aware_test.py b/elasticdl/python/tests/staleness_aware_test.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/tests/staleness_aware_test.py
@@ -0,0 +1,83 @@
+import time
+import unittest
+from concurrent.futures import ThreadPoolExecutor
+
+import tensorflow as tf
+
+from elasticdl.python.master.lr_modulation import (
+    add_lr_modulation_to_optimizer,
+)
+
+
+class LearningRateTest(unittest.TestCase):
+    @staticmethod
+    def get_lr(lr_modulation, opt, multiplier):
+        lr_modulation.set_multiplier(multiplier)
+        # sleep 1s to wait that all threads are in this method call
+        time.sleep(1)
+        return opt.learning_rate
+
+    @staticmethod
+    def apply_gradients_with_modulation(
+        lr_modulation, opt, multiplier, variables, grads
+    ):
+        grads_and_vars = zip(grads, variables)
+        lr_modulation.set_multiplier(multiplier)
+        # sleep 1s to wait that all threads are in this method call
+        time.sleep(1)
+        opt.apply_gradients(grads_and_vars)
+        return [v.numpy() for v in variables]
+
+    def test_lr_modulation(self):
+        lr = 0.1
+        multipliers = [1, 0.5, 0.1, 0.01]
+        counts = len(multipliers)
+        opt = tf.optimizers.SGD(lr)
+        lr_modulation = add_lr_modulation_to_optimizer(opt)
+
+        executor = ThreadPoolExecutor(max_workers=counts)
+        tasks = [
+            executor.submit(self.get_lr, lr_modulation, opt, m)
+            for m in multipliers
+        ]
+        results = [tasks[i].result() for i in range(counts)]
+        for i in range(counts):
+            self.assertAlmostEqual(results[i], lr * multipliers[i])
+
+        variables = []
+        grads = []
+        original_values = [1.2, 0.8]
+        grad_values = [0.2, 0.1]
+
+        for i in range(counts):
+            variables.append([tf.Variable(v) for v in original_values])
+            grads.append([tf.convert_to_tensor(g) for g in grad_values])
+
+        tasks = [
+            executor.submit(
+                self.apply_gradients_with_modulation,
+                lr_modulation,
+                opt,
+                multipliers[i],
+                variables[i],
+                grads[i],
+            )
+            for i in range(counts)
+        ]
+        results = [tasks[i].result() for i in range(counts)]
+        place = 5
+        for i in range(0, counts):
+            i_diff = [
+                original_values[j] - results[i][j]
+                for j in range(len(original_values))
+            ]
+            for j in range(len(original_values)):
+                # variable value change ratio equals the learning rate ratio
+                # for SGD without momentum
+                self.assertAlmostEqual(
+                    i_diff[j], grad_values[j] * lr * multipliers[i], place
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()