Repository: sql-machine-learning/elasticdl

Issue:
Support async SGD in  Optimizer Wrapper



Discussion:

Resolving patch:
diff --git a/elasticdl/python/master/optimizer_wrapper.py b/elasticdl/python/master/optimizer_wrapper.py
--- a/elasticdl/python/master/optimizer_wrapper.py
+++ b/elasticdl/python/master/optimizer_wrapper.py
@@ -1,6 +1,8 @@
 """Optimizer Wrapper for ElasticDL"""
 
 
+import threading
+
 import numpy as np
 import tensorflow as tf
 from tensorflow.keras.optimizers import (
@@ -118,15 +120,19 @@ def __init__(
                 {layer name: `embedding_dim`} where layer name is the
                 name of ElasticDL embedding layer and `embedding_dim`
                 is the output dimension of corresponding embedding layer.
-            use_async: A python bool. True if using asynchronous updates.
+            use_async: A python bool. True if using asynchronous updates. When
+                using asynchronoues updates, `OptimizerWrapper` is thread-safe
+                for non-embedding variables and is not thread-safe for
+                embedding table.
         """
         self._opt = opt
         self._kv_store_endpoint = kv_store_endpoint
         self._embed_dims = embedding_dims
         self._use_async = use_async
         self._slot_initial_value = {}
-        self._embed_variables = {}
-        self._slot_variables = {}
+
+        self._tls = threading.local()
+        self._init_thread_local()
 
         # "-" in slot name is not supported
         if isinstance(opt, SGD):
@@ -164,8 +170,10 @@ def __init__(
         for slot in self._allowed_slot_names:
             self._slot_initial_value.setdefault(slot, 0.0)
 
-        # record unique ids of gradients
-        self._unique_ids_all_layers = {}
+    def _init_thread_local(self):
+        self._tls._unique_ids_all_layers = {}
+        self._tls._embed_variables = {}
+        self._tls._slot_variables = {}
 
     def apply_gradients(self, grads_and_vars):
         """Update variable values.
@@ -174,13 +182,9 @@ def apply_gradients(self, grads_and_vars):
             grads_and_vars: A list of (gradient, variable) pairs.
 
         """
-
-        # TODO (yunjian.lmh): support `use_async=True`
-        if self._use_async:
-            raise NotImplementedError(
-                "`use_async=True` in Optimizer Wrapper is not "
-                "supported now."
-            )
+        # TODO (#1255): Discuss whether `OptimizerWrapper` needs a lock after
+        # implementing PS.
+        self._init_thread_local()
 
         grads_and_vars = list(grads_and_vars)
 
@@ -212,6 +216,8 @@ def apply_gradients(self, grads_and_vars):
 
         self._report_to_kv_store()
 
+        self._delete_variables()
+
     def _lookup_embeddings_and_slots(self, grads_and_vars):
         """Look up embedding vectors and slot values form kv store.
 
@@ -293,19 +299,19 @@ def _generate_lookup_keys(self, grads_and_vars):
         embed_key_index = {}
         slot_keys = []
         slot_key_index = {}
-        self._unique_ids_all_layers = {}
+        self._tls._unique_ids_all_layers = {}
 
         # generate keys
         for it, (grad, layer_name) in enumerate(grads_and_vars):
             # de-duplicate gradient's indices
             unique_ids, indices = tf.unique(grad.indices)
             unique_ids = unique_ids.numpy()
-            if layer_name in self._unique_ids_all_layers:
+            if layer_name in self._tls._unique_ids_all_layers:
                 # TODO: support grads_and_vars with duplicated layer name
                 logger.warning(
                     "grads_and_vars has duplicated layer name %s." % layer_name
                 )
-            self._unique_ids_all_layers[layer_name] = unique_ids
+            self._tls._unique_ids_all_layers[layer_name] = unique_ids
             grad_new = tf.IndexedSlices(grad.values, indices)
             grads_and_vars[it] = (grad_new, layer_name)
 
@@ -365,11 +371,13 @@ def _set_slot_values_to_variables(self, values):
 
     def _get_slot_variable(self, layer_name, slot_name):
         """Get the variable for specified slot."""
-        return self._slot_variables.get(layer_name, {}).get(slot_name, None)
+        return self._tls._slot_variables.get(layer_name, {}).get(
+            slot_name, None
+        )
 
     def _get_embedding_variable(self, layer_name):
         """Get the variable for the specified ElasticDL embedding layer."""
-        return self._embed_variables.get(layer_name, None)
+        return self._tls._embed_variables.get(layer_name, None)
 
     # TODO: refactor _create_slot_variable and _create_embedding_variable
     # into one function
@@ -381,7 +389,7 @@ def _create_embedding_variable(self, layer_name, initial_value):
         # this number may differ between different iterations
         shape = tf.TensorShape((None, dim))
 
-        if self._embed_variables.get(layer_name, None) is not None:
+        if self._tls._embed_variables.get(layer_name, None) is not None:
             raise RuntimeError(
                 "Embedding variable with layer name=%s has already be "
                 "created." % (layer_name)
@@ -389,12 +397,12 @@ def _create_embedding_variable(self, layer_name, initial_value):
 
         embed_var = tf.Variable(
             initial_value,
-            name=layer_name,
+            name=layer_name + str(threading.get_ident()),
             shape=shape,
             dtype=tf.float32,
             trainable=False,
         )
-        self._embed_variables[layer_name] = embed_var
+        self._tls._embed_variables[layer_name] = embed_var
         return embed_var
 
     def _create_slot_variable(self, layer_name, slot_name, initial_value):
@@ -405,7 +413,9 @@ def _create_slot_variable(self, layer_name, slot_name, initial_value):
         # this number may differ between different iterations
         shape = tf.TensorShape((None, dim))
 
-        slot_variables_dict = self._slot_variables.setdefault(layer_name, {})
+        slot_variables_dict = self._tls._slot_variables.setdefault(
+            layer_name, {}
+        )
         if slot_variables_dict.get(slot_name, None) is not None:
             raise RuntimeError(
                 "Slot variable with (layer name=%s, slot name=%s) has "
@@ -459,7 +469,7 @@ def _report_to_kv_store(self):
         """Report updated embedding vectors and slots to kv store."""
         keys = []
         values = []
-        for layer, ids in self._unique_ids_all_layers.items():
+        for layer, ids in self._tls._unique_ids_all_layers.items():
             value = self._get_embedding_variable(layer).numpy()
             for id, v in zip(ids, value):
                 keys.append(Embedding.get_key([layer, id]))
@@ -475,6 +485,26 @@ def _report_to_kv_store(self):
             keys, values, self._kv_store_endpoint
         )
 
+    def _delete_variables(self):
+        # Slot variable access in optimizer requires corresponding embedding
+        # variable information. Delete slot variables first.
+        for layer_name, slots in self._tls._slot_variables.items():
+            embed_var = self._get_embedding_variable(layer_name)
+            embed_var_key = _var_key(embed_var)
+            del self._opt._slots[embed_var_key]
+            for _, var in slots.items():
+                self._opt._weights.remove(var)
+        for key in list(self._tls._slot_variables.keys()):
+            del self._tls._slot_variables[key]
+
+        # Delete variables in embed_variables.
+        for key in list(self._tls._embed_variables.keys()):
+            del self._tls._embed_variables[key]
+
+        # Delete variables in unique_ids_all_layers.
+        for key in list(self._tls._unique_ids_all_layers.keys()):
+            del self._tls._unique_ids_all_layers[key]
+
     @property
     def allowed_slot_names(self):
         return self._allowed_slot_names


Test patch:
diff --git a/elasticdl/python/tests/optimizer_wrapper_test.py b/elasticdl/python/tests/optimizer_wrapper_test.py
--- a/elasticdl/python/tests/optimizer_wrapper_test.py
+++ b/elasticdl/python/tests/optimizer_wrapper_test.py
@@ -1,7 +1,9 @@
 import copy
 import os
 import random
+import time
 import unittest
+from concurrent.futures import ThreadPoolExecutor
 
 import mock
 import numpy as np
@@ -269,7 +271,7 @@ def test_generate_lookup_keys(self):
 
         for layer, ids in zip(layers, ids_list):
             self.assertTrue(
-                (opt_wrapper._unique_ids_all_layers[layer] == ids).all()
+                (opt_wrapper._tls._unique_ids_all_layers[layer] == ids).all()
             )
 
     def test_parse_lookup_values(self):
@@ -334,7 +336,7 @@ def test_lookup(self):
 
         for ids, layer in zip(ids_list, layers):
             self.assertTrue(
-                (opt_wrapper._unique_ids_all_layers[layer] == ids).all()
+                (opt_wrapper._tls._unique_ids_all_layers[layer] == ids).all()
             )
 
             values, _ = mock_kv_store.lookup(
@@ -392,17 +394,20 @@ def test_set_slot_values_to_variables(self):
                 self.assertTrue(
                     (
                         slots_dict[slot].numpy()
-                        == opt_wrapper._slot_variables[layer][slot].numpy()
+                        == opt_wrapper._tls._slot_variables[layer][
+                            slot
+                        ].numpy()
                     ).all()
                 )
 
                 slots_dict[slot].assign(tf.ones((10, 4)))
                 self.assertTrue(
                     np.isclose(
-                        opt_wrapper._slot_variables[layer][slot].numpy(), 1.0
+                        opt_wrapper._tls._slot_variables[layer][slot].numpy(),
+                        1.0,
                     ).all()
                 )
-                opt_wrapper._slot_variables[layer][slot].assign(
+                opt_wrapper._tls._slot_variables[layer][slot].assign(
                     -tf.ones((10, 4))
                 )
                 self.assertTrue(
@@ -413,7 +418,7 @@ def test_set_slot_values_to_variables(self):
         opt_wrapper._set_slot_values_to_variables(slot_values_new)
         self.assertTrue(
             np.isclose(
-                opt_wrapper._slot_variables["test-1"]["m"].numpy(), 0.0
+                opt_wrapper._tls._slot_variables["test-1"]["m"].numpy(), 0.0
             ).all()
         )
 
@@ -443,14 +448,14 @@ def test_set_embedding_values_to_variables(self):
         for i, layer in enumerate(layers):
             self.assertTrue(
                 (
-                    opt_wrapper._embed_variables[layer].numpy()
+                    opt_wrapper._tls._embed_variables[layer].numpy()
                     == embedding_values[layer]
                 ).all()
             )
             self.assertTrue(
                 (
                     grads_and_vars[i][1].numpy()
-                    == opt_wrapper._embed_variables[layer].numpy()
+                    == opt_wrapper._tls._embed_variables[layer].numpy()
                 ).all()
             )
 
@@ -461,7 +466,7 @@ def test_set_embedding_values_to_variables(self):
         )
         self.assertTrue(
             np.isclose(
-                opt_wrapper._embed_variables["test-1"].numpy(), 0.0
+                opt_wrapper._tls._embed_variables["test-1"].numpy(), 0.0
             ).all()
         )
 
@@ -470,16 +475,16 @@ def test_report_to_kv_store(self):
         opt_wrapper = OptimizerWrapper(opt, None, {})
 
         ids_list = [[1, 5], [10]]
-        opt_wrapper._unique_ids_all_layers = {
+        opt_wrapper._tls._unique_ids_all_layers = {
             "test_1": np.array(ids_list[0]),
             "test_2": np.array(ids_list[1]),
         }
         t = np.array([1.0, 1.0, 1.0])
-        opt_wrapper._embed_variables = {
+        opt_wrapper._tls._embed_variables = {
             "test_1": tf.Variable([t, t * 5]),
             "test_2": tf.Variable([t * 10]),
         }
-        opt_wrapper._slot_variables = {
+        opt_wrapper._tls._slot_variables = {
             "test_1": {"momentum": tf.Variable([t / 10.0, t / 2.0])},
             "test_2": {"momentum": tf.Variable([t])},
         }
@@ -620,6 +625,188 @@ def test_correctness(self):
                     opt, X_sparse_one_iter, Y_sparse_one_iter, seed, **kargs
                 )
 
+    def _test_async_correctness(
+        self,
+        grads_and_vars_batches,
+        embed_values,
+        expected_non_embed_values,
+        expected_embed_values=None,
+    ):
+        """Checks the correctness of async OptimizerWrapper. This function
+        creates many threads and these threads call
+        `OptimizerWrapper.apply_gradients` simultaneously.
+
+        Args:
+            grads_and_vars_batches: A python list of `grads_and_vars`. Every
+                thread takes a `grads_and_vars` and calls `apply_gradients`.
+            embed_values: A python dictionary of
+                `(layer_name, embedding table)`.
+            expected_non_embed_values: A python list of expected non-embdding
+                values after applying gradients.
+            expected_embed_values: A python dictionary of expected embedding
+                values after applying gradients. None means no need to check
+                embedding values.
+        """
+        thread_num = len(grads_and_vars_batches)
+        embed_dims = {}
+        embed_var_n = len(embed_values)
+        mock_kv_store = MockKvStore()
+        for layer, values in embed_values.items():
+            embed_dims[layer] = values.shape[1]
+            input_dim = values.shape[0]
+
+            keys = [
+                Embedding.get_key([layer, idx]) for idx in range(input_dim)
+            ]
+            mock_kv_store.update(keys, values)
+
+        opt = SGD(0.1)
+        opt_wrapper = OptimizerWrapper(opt, None, embed_dims, True)
+
+        with mock.patch.object(
+            EmbeddingService, "lookup_embedding", mock_kv_store.lookup
+        ), mock.patch.object(
+            EmbeddingService, "update_embedding", mock_kv_store.update
+        ):
+            # call optimizer_wrapper.apply_gradients asynchronously
+            def _apply_gradients(opt_wrapper, grads_and_vars):
+                # sleep 1s to wait that all threads are in this method call
+                time.sleep(1)
+                opt_wrapper.apply_gradients(grads_and_vars)
+
+            executor = ThreadPoolExecutor(max_workers=thread_num)
+            tasks = [
+                executor.submit(_apply_gradients, opt_wrapper, grads_and_vars)
+                for grads_and_vars in grads_and_vars_batches
+            ]
+            _ = [task.result() for task in tasks]
+
+            # check updated results of non-embedding variables
+            non_embed_vars = [
+                var for grad, var in grads_and_vars_batches[0][:-embed_var_n]
+            ]
+            for var, expected_value in zip(
+                non_embed_vars, expected_non_embed_values
+            ):
+                self.assertTrue(np.isclose(var.numpy(), expected_value).all())
+
+            # `expected_embed_values=None` means that no need to check
+            # embedding table
+            if not expected_embed_values:
+                return
+            # check updated results of embedding table
+            for layer, expected_values in expected_embed_values.items():
+                keys = [
+                    Embedding.get_key([layer, idx]) for idx in range(input_dim)
+                ]
+                raw_value, _ = mock_kv_store.lookup(keys)
+                value = np.concatenate(raw_value).reshape(input_dim, -1)
+
+                self.assertTrue(
+                    any(
+                        [
+                            np.isclose(value, expected).all()
+                            for expected in expected_values
+                        ]
+                    )
+                )
+
+    def test_async_correctness(self):
+        """Tests the correctness of async updates in `OptimizerWrapper`.
+
+        Testing the correctness is not simple because OptimizerWrapper is not
+        thread-safe for embedding table. This test case lists all the possible
+        results when `thread_number=2` and test the correctness. This test case
+        also tests that OptimizerWrapper does not raise Error with a large
+        thread number(8).
+        """
+        max_thread_num = 8
+        input_dim = 4
+        output_dim = 3
+        non_embed_vars = [
+            tf.Variable([0.0] * output_dim),
+            tf.Variable([1.0] * output_dim),
+        ]
+        non_embed_vars_copy = copy.deepcopy(non_embed_vars)
+        non_embed_grads_batches = [
+            [
+                tf.constant([i + 1] * output_dim, dtype=tf.float32),
+                tf.constant([-i - 1] * output_dim, dtype=tf.float32),
+            ]
+            for i in range(max_thread_num)
+        ]
+        embed_shape = (input_dim, output_dim)
+        embed_value_count = output_dim * input_dim
+        embed_layers = ["embed_1", "embed_2"]
+        embed_values = {
+            embed_layers[0]: np.arange(
+                embed_value_count, dtype=np.float32
+            ).reshape(embed_shape),
+            embed_layers[1]: np.arange(
+                embed_value_count, dtype=np.float32
+            ).reshape(embed_shape),
+        }
+
+        embed_grads_batches = [
+            [
+                tf.IndexedSlices(
+                    tf.reshape(
+                        tf.constant([i + 1.0] * embed_value_count), embed_shape
+                    ),
+                    tf.constant(list(range(input_dim))),
+                ),
+                tf.IndexedSlices(
+                    tf.reshape(
+                        tf.constant([-i - 1.0] * embed_value_count),
+                        embed_shape,
+                    ),
+                    tf.constant(list(range(input_dim))),
+                ),
+            ]
+            for i in range(max_thread_num)
+        ]
+
+        # thread number = 2
+        expected_non_embed_values = [[-0.3, -0.3, -0.3], [1.3, 1.3, 1.3]]
+        expected_embed_values = {
+            embed_layers[0]: [
+                (np.arange(12) - 0.1).reshape(embed_shape),
+                (np.arange(12) - 0.2).reshape(embed_shape),
+                (np.arange(12) - 0.3).reshape(embed_shape),
+            ],
+            embed_layers[1]: [
+                (np.arange(12) + 0.1).reshape(embed_shape),
+                (np.arange(12) + 0.2).reshape(embed_shape),
+                (np.arange(12) + 0.3).reshape(embed_shape),
+            ],
+        }
+        grads_and_vars_batches = [
+            list(zip(non_embed_grads_batches[i], non_embed_vars))
+            + list(zip(embed_grads_batches[i], embed_layers))
+            for i in range(2)
+        ]
+
+        self._test_async_correctness(
+            grads_and_vars_batches,
+            embed_values,
+            expected_non_embed_values,
+            expected_embed_values,
+        )
+
+        # thread number = 8
+        grads_sum = max_thread_num * (max_thread_num + 1) / 2 / 10.0
+        expected_non_embed_values = [[-grads_sum] * 3, [1 + grads_sum] * 3]
+        grads_and_vars_batches = [
+            list(zip(non_embed_grads_batches[i], non_embed_vars_copy))
+            + list(zip(embed_grads_batches[i], embed_layers))
+            for i in range(max_thread_num)
+        ]
+        # Do not check updating results of embedding table when `thread_num>2`.
+        # Because there are too many possible results.
+        self._test_async_correctness(
+            grads_and_vars_batches, embed_values, expected_non_embed_values
+        )
+
 
 if __name__ == "__main__":
     unittest.main()