Repository: sql-machine-learning/elasticdl

Issue:
[async SGD] Do not copy edl_embedding_gradients/indexed_grads/tmp(other grads) to self._*, but use directly for update_model



Discussion:
Multiple async ReportGradient calls will override each other. Also no need to do the extra copy.

Resolving patch:
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
@@ -200,25 +200,21 @@ def _update_edl_embedding_table(self, name_var_list):
                 embedding_service_endpoint=self._embedding_service_endpoint,
             )
 
-    def _update_model(self):
+    def _update_model(self, grads, indexed_grads, edl_embedding_gradients):
         grad_var = []
 
         # (grad, var) pairs excluding keras Embedding layer and
         # ElasticDL Embedding layer
-        for k in self._gradient_sum:
-            if not self._use_async:
-                self._gradient_sum[k] = (
-                    self._gradient_sum[k] / self._grad_to_wait
-                )
-            grad_var.append((self._gradient_sum[k], self._model[k]))
+        for k in grads:
+            grad_var.append((grads[k], self._model[k]))
 
         # (grad, var) pair of Keras Embedding layer
-        for k in self._gradient_sum_indexed:
-            grad_var.append((self._gradient_sum_indexed[k], self._model[k]))
+        for k in indexed_grads:
+            grad_var.append((indexed_grads[k], self._model[k]))
 
         # (grad, var) pair of ElasticDL Embedding layer
-        if self._edl_embedding_gradients:
-            for layer_name, grads in self._edl_embedding_gradients.items():
+        if edl_embedding_gradients:
+            for layer_name, grads in edl_embedding_gradients.items():
                 grad_var.append((grads, layer_name))
 
         self._opt.apply_gradients(grad_var)
@@ -414,16 +410,27 @@ def _process_gradients(
                 else:
                     self._gradient_sum[k] = v
             self._grad_n += 1
-        else:
-            # TODO: do not accumulate gradients but apply directly.
-            pass
 
-        # staleness-aware learning rate modulation
+        need_to_update_model = self._use_async
+        if not self._use_async and self._grad_n >= self._grad_to_wait:
+            need_to_update_model = True
+            # get gradient average for sync SGD
+            for k in self._gradient_sum:
+                self._gradient_sum[k] = (
+                    self._gradient_sum[k] / self._grad_to_wait
+                )
+            edl_embedding_gradients = self._edl_embedding_gradients
+            indexed_grads = self._gradient_sum_indexed
+            grads = self._gradient_sum
+        if need_to_update_model:
+            self._update_optimizer(request_version)
+            self._update_model(grads, indexed_grads, edl_embedding_gradients)
+
+    def _update_optimizer(self, request_version):
         if self._lr_modulation:
+            # staleness-aware learning rate modulation
             staleness = max(1, self._version - request_version)
             self._lr_modulation.set_multiplier(1.0 / staleness)
-        if self._use_async or self._grad_n >= self._grad_to_wait:
-            self._update_model()
 
     def ReportTaskResult(self, request, _):
         if request.err_message:


Test patch:
diff --git a/elasticdl/python/tests/report_gradients_of_bet_test.py b/elasticdl/python/tests/report_gradients_of_bet_test.py
--- a/elasticdl/python/tests/report_gradients_of_bet_test.py
+++ b/elasticdl/python/tests/report_gradients_of_bet_test.py
@@ -283,7 +283,11 @@ def test_report_bet_gradients_master_to_service(self):
         ):
             with master._lock:
                 assert master._lock.locked()
-                master._update_model()
+                master._update_model(
+                    master._gradient_sum,
+                    master._gradient_sum_indexed,
+                    master._edl_embedding_gradients,
+                )
 
         expected_embedding_table = {
             layer_names[0]: np.array(