Repository: sql-machine-learning/elasticdl

Issue:
Support tf.function when model has ElasticDL embedding layer
Currently, ElasticDL embedding layer looks up embedding vectors from Redis and use `tf.GradientTape.watch(embedding_vectors)` to tell gradient tape to record gradients of embedding vectors. However, in a function decorated by `tf.function`, `tape.watch` does not work.

We can use trainable variables instead. In a function decorated by `tf.function`, TensorFlow supports assigning values to variables, and gradient tape can automatically record gradients of trainable variables. So, we can create a TensorFlow variable called `batch_embedding` in advance. Every time worker does forward-pass, embedding layer assigns the embedding vectors to the `batch_embedding`.

One detail is worthy mention. Currently, ElasticDL supports the situation that an embedding layer is called more than once in one forward-pass. Please refer to this model as an example:

https://github.com/sql-machine-learning/elasticdl/blob/db0a1de92f3550878745feaea90c2f9fa0a2454b/elasticdl/python/tests/embedding_test_module.py#L7-L41

However, it is not simple to support this feature under `tf.function` because we plan to create TensorFlow variable in advance. We don't think this is an important feature. Thus we plan to not use `tf.function` if an embedding layer is called more than once.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/common/model_helper.py b/elasticdl/python/common/model_helper.py
--- a/elasticdl/python/common/model_helper.py
+++ b/elasticdl/python/common/model_helper.py
@@ -160,7 +160,7 @@ def get_non_embedding_trainable_vars(model, embedding_layers):
     """
     embedding_items = []
     for layer in embedding_layers:
-        embedding_items.extend([bet for bet, _ in layer.bet_ids_pair])
+        embedding_items.extend(layer.trainable_variables)
     return [
         var for var in model.trainable_variables if var not in embedding_items
     ]
diff --git a/elasticdl/python/elasticdl/layers/embedding.py b/elasticdl/python/elasticdl/layers/embedding.py
--- a/elasticdl/python/elasticdl/layers/embedding.py
+++ b/elasticdl/python/elasticdl/layers/embedding.py
@@ -1,9 +1,15 @@
+import collections
+
 import numpy as np
 import tensorflow as tf
 from tensorflow.python.keras.utils import tf_utils
 
 from elasticdl.python.master.embedding_service import EmbeddingService
 
+EmbeddingAndIds = collections.namedtuple(
+    "EmbeddingAndIds", ["batch_embedding", "batch_ids"]
+)
+
 
 class Embedding(tf.keras.layers.Layer):
     """
@@ -51,7 +57,29 @@ def __init__(
         self.embedding_service_endpoint = embedding_service_endpoint
         self.tape = None
         self.lookup_func = None
-        self.bet_ids_pair = []
+
+        self._embedding_and_ids_eagerly = []
+
+        # BET's shape and ids' shape in `self._embedding_and_ids_graph` have
+        # `None` dimension. This is because they have different shapes in
+        # different iterations.
+        # `tf.Variable` requires initial value if shape has `None` dimension.
+        self._embedding_and_ids_graph = [
+            EmbeddingAndIds(
+                batch_embedding=tf.Variable(
+                    initial_value=tf.zeros((1, self.output_dim)),
+                    shape=tf.TensorShape((None, self.output_dim)),
+                    dtype=tf.float32,
+                    trainable=True,
+                ),
+                batch_ids=tf.Variable(
+                    initial_value=tf.zeros((1, 1), dtype=tf.int64),
+                    shape=tf.TensorShape(None),
+                    dtype=tf.int64,
+                    trainable=False,
+                ),
+            )
+        ]
 
     @tf_utils.shape_type_conversion
     def compute_output_shape(self, input_shape):
@@ -141,6 +169,22 @@ def lookup_embedding(self, unique_ids):
         embedding_vectors = np.concatenate(embedding_vectors, axis=0)
         return embedding_vectors.reshape((len(keys), self.output_dim))
 
+    def _record_gradients(self, batch_embedding, ids):
+        if tf.executing_eagerly():
+            self.tape.watch(batch_embedding)
+            self._embedding_and_ids_eagerly.append(
+                EmbeddingAndIds(batch_embedding, ids)
+            )
+        else:
+            # In graph mode, assigning tensors to trainable variables is
+            # allowed and tape can record the gradients of trainable
+            # variables automatically.
+            embedding_and_ids = self._embedding_and_ids_graph[0]
+            embedding_and_ids.batch_embedding.assign(batch_embedding)
+            embedding_and_ids.batch_ids.assign(ids)
+            batch_embedding = embedding_and_ids.batch_embedding
+        return batch_embedding
+
     def call(self, input):
         if isinstance(input, tf.SparseTensor):
             return self._sparse_input_call(input)
@@ -148,19 +192,17 @@ def call(self, input):
         ids = tf.convert_to_tensor(input, name="embedding_ids")
         flat_ids = tf.reshape(ids, [-1])
         unique_ids, idx = tf.unique(flat_ids)
-        batch_embedding_tensor = tf.py_function(
+        # Gradient for `batch_embedding` is SparseTensor here due to
+        # `tf.gather` op. `tf.gather` accesses tensor slices, resulting in
+        # sparse tensor gradient.
+        batch_embedding = tf.py_function(
             self.lookup_embedding, inp=[unique_ids], Tout=tf.float32
         )
+        # TODO: use tf.cond rather than python if statement
         if self.tape:
-            # tape.watch works with eager mode only.
-            # Gradient for embeddings is SparseTensor here due to tf.gather op.
-            # tf.gather accesses tensor slices, resulting in sparse tensor
-            # gradient.
-            if not tf.executing_eagerly():
-                raise RuntimeError("tape.watch only works with eager mode")
-            self.tape.watch(batch_embedding_tensor)
-            self.bet_ids_pair.append((batch_embedding_tensor, flat_ids))
-        outputs = tf.gather(batch_embedding_tensor, idx)
+            batch_embedding = self._record_gradients(batch_embedding, flat_ids)
+
+        outputs = tf.gather(batch_embedding, idx)
         # tf.reshape does not support shape with None. Replace None with -1.
         if ids.get_shape().rank == 2:
             output_shape = (-1, ids.get_shape()[1], self.output_dim)
@@ -176,27 +218,33 @@ def _sparse_input_call(self, sparse_input):
                 "combiner must set sum, mean or sqrtn for sparse input"
             )
         unique_ids, idx = tf.unique(sparse_input.values)
-        embeddings = tf.py_function(
+        # Gradient for `batch_embedding` is dense tensor.
+        batch_embedding = tf.py_function(
             self.lookup_embedding, inp=[unique_ids], Tout=tf.float32
         )
+        # TODO: use tf.cond rather than python if statement
         if self.tape:
-            # tape.watch works with eager mode only
-            # gradient for embeddings is dense tensor for sparse_input_call
-            if not tf.executing_eagerly():
-                raise RuntimeError("tape.watch only works with eager mode")
-            self.tape.watch(embeddings)
-            self.bet_ids_pair.append((embeddings, unique_ids))
+            batch_embedding = self._record_gradients(
+                batch_embedding, unique_ids
+            )
+
         segment_ids = sparse_input.indices[:, 0]
         if segment_ids.dtype != tf.int32:
             segment_ids = tf.cast(segment_ids, tf.int32)
 
         if self.combiner == "sum":
-            embeddings = tf.sparse.segment_sum(embeddings, idx, segment_ids)
+            batch_embedding = tf.sparse.segment_sum(
+                batch_embedding, idx, segment_ids
+            )
         elif self.combiner == "mean":
-            embeddings = tf.sparse.segment_mean(embeddings, idx, segment_ids)
+            batch_embedding = tf.sparse.segment_mean(
+                batch_embedding, idx, segment_ids
+            )
         elif self.combiner == "sqrtn":
-            embeddings = tf.sparse.segment_sqrt_n(embeddings, idx, segment_ids)
-        return embeddings
+            batch_embedding = tf.sparse.segment_sqrt_n(
+                batch_embedding, idx, segment_ids
+            )
+        return batch_embedding
 
     def compute_mask(self, inputs, mask=None):
         if isinstance(input, tf.SparseTensor):
@@ -206,7 +254,7 @@ def compute_mask(self, inputs, mask=None):
         return tf.math.not_equal(inputs, 0)
 
     def reset(self):
-        self.bet_ids_pair = []
+        self._embedding_and_ids_eagerly = []
         self.tape = None
 
     def set_tape(self, tape):
@@ -214,3 +262,12 @@ def set_tape(self, tape):
 
     def set_endpoint(self, endpoint):
         self.embedding_service_endpoint = endpoint
+
+    @property
+    def embedding_and_ids(self):
+        """
+        Return bet and ids pairs.
+        """
+        if self._embedding_and_ids_eagerly:
+            return self._embedding_and_ids_eagerly
+        return self._embedding_and_ids_graph
diff --git a/elasticdl/python/worker/worker.py b/elasticdl/python/worker/worker.py
--- a/elasticdl/python/worker/worker.py
+++ b/elasticdl/python/worker/worker.py
@@ -110,6 +110,7 @@ def __init__(
     def set_model(self, model_inst):
         """Set model instance to worker."""
         self._model = model_inst
+        self._train_eagerly = False
         self._init_embedding_layer()
         self._var_created = self._model.built
         self._non_embed_vars = []
@@ -125,6 +126,9 @@ def _init_embedding_layer(self):
         self._embedding_layers = find_layer(self._model, Embedding)
         for layer in self._embedding_layers:
             layer.set_endpoint(self._embedding_service_endpoint)
+        self._need_embedding_layer_check = (
+            True if self._embedding_layers else False
+        )
 
     def _set_tape_for_embedding(self, tape):
         for layer in self._embedding_layers:
@@ -204,7 +208,7 @@ def report_gradient(self, grads):
             # `bet_id_pair`.
             bet_number = 0
             for layer in self._embedding_layers:
-                bet_number += len(layer.bet_ids_pair)
+                bet_number += len(layer.embedding_and_ids)
             if len(edl_embedding_grads) != bet_number:
                 raise ValueError(
                     "elasticdl.layers.embedding related gradient number %d "
@@ -216,7 +220,7 @@ def report_gradient(self, grads):
             for layer in self._embedding_layers:
                 g_values = None
                 g_indices = None
-                for _, ids in layer.bet_ids_pair:
+                for _, ids in layer.embedding_and_ids:
                     grad = edl_embedding_grads[grad_accum_iter]
                     grad_accum_iter += 1
                     # ElasticDL embedding layer with Sparse Gradients
@@ -268,14 +272,39 @@ def report_prediction_outputs(self, predictions):
             )
         return True
 
-    def _create_variable_and_report(self, features):
-        # Use model.call to create variables, then report to ps
-        _ = self._model.call(features)
+    def _run_model_call_before_training(self, features):
+        """Call `self._model.call` before training for two things:
+            * Create variables and report to ps if not created.
+            * Check whether there is an embedding layer that is called
+              more than once during one forward-pass.
+        """
+        if self._embedding_layers:
+            with tf.GradientTape() as tape:
+                self._set_tape_for_embedding(tape)
+                _ = self._model.call(features)
+        else:
+            _ = self._model.call(features)
         self._non_embed_vars = get_non_embedding_trainable_vars(
             self._model, self._embedding_layers
         )
-        self.report_variable()
-        self._var_created = True
+
+        if not self._var_created:
+            self.report_variable()
+            self._var_created = True
+
+        if self._need_embedding_layer_check:
+            self._train_eagerly = False
+            for layer in self._embedding_layers:
+                if len(layer.embedding_and_ids) > 1:
+                    self._train_eagerly = True
+                    logger.warning(
+                        "ElasticDL embedding layer %s is called more than "
+                        "once, this will make the training process unable "
+                        "to accelerate with tf.function." % (layer.name)
+                    )
+            self._need_embedding_layer_check = False
+
+        self._reset_embedding()
 
     def get_trainable_items(self):
         """
@@ -286,7 +315,12 @@ def get_trainable_items(self):
         bets = []
         if self._embedding_layers:
             for layer in self._embedding_layers:
-                bets.extend([i for (i, _) in layer.bet_ids_pair])
+                bets.extend(
+                    [
+                        batch_embedding
+                        for (batch_embedding, _) in layer.embedding_and_ids
+                    ]
+                )
         return self._non_embed_vars + bets
 
     def training_process(self, features, labels):
@@ -294,7 +328,7 @@ def training_process(self, features, labels):
         training for models with elasticdl.layers.embedding does not
         support tf.function decorator
         """
-        if self._embedding_layers:
+        if self._train_eagerly:
             return self.training_process_eagerly(features, labels)
         else:
             return self.training_process_with_acceleration(features, labels)
@@ -343,8 +377,8 @@ def _run_prediction_task(self, features):
     def _process_minibatch(
         self, task_type, features, labels, min_model_version
     ):
-        if not self._var_created:
-            self._create_variable_and_report(features)
+        if self._need_embedding_layer_check or not self._var_created:
+            self._run_model_call_before_training(features)
         for _ in range(self._max_minibatch_retry_num):
             if task_type == elasticdl_pb2.EVALUATION:
                 if min_model_version == -1:


Test patch:
diff --git a/elasticdl/python/tests/example_test.py b/elasticdl/python/tests/example_test.py
--- a/elasticdl/python/tests/example_test.py
+++ b/elasticdl/python/tests/example_test.py
@@ -124,7 +124,8 @@ def distributed_train_and_evaluate(
         )
 
         if dataset == "imagenet":
-            shards = {create_imagenet_recordio_file(16, feature_shape): 16}
+            batch_size = 8
+            shards = {create_imagenet_recordio_file(8, feature_shape): 8}
         elif dataset == "frappe":
             shards = {create_frappe_recordio_file(16, feature_shape, 5383): 16}
         else:
diff --git a/elasticdl/python/tests/layer_test.py b/elasticdl/python/tests/layer_test.py
--- a/elasticdl/python/tests/layer_test.py
+++ b/elasticdl/python/tests/layer_test.py
@@ -89,8 +89,8 @@ def create_embedding_layer(
 
 
 @tf.function
-def layer_call(layer, inputs):
-    return layer.call(inputs)
+def module_call(module, inputs):
+    return module.call(inputs)
 
 
 def get_correct_values_for_sparse_test(indices, values, dense_shape, combiner):
@@ -115,6 +115,15 @@ def get_correct_values_for_sparse_test(indices, values, dense_shape, combiner):
 
 
 class EmbeddingLayerTest(unittest.TestCase):
+    def _run_forward_pass_and_compare(
+        self, call_fns, correct_values, output_dim
+    ):
+        for call_fn in call_fns:
+            values = call_fn()
+            values = values.numpy().reshape(-1, output_dim)
+            for v, correct_v in zip(values, correct_values):
+                self.assertTrue((np.absolute(v - correct_v) < 0.00001).all())
+
     def test_embedding_layer(self):
         output_dim = 8
         embedding_size = 16
@@ -125,16 +134,13 @@ def test_embedding_layer(self):
         self.assertEqual(output_shape, input_shape + (output_dim,))
 
         ids = [0, 1, 3, 8, 3, 2, 3]
-        values = layer.call(ids)
-        values = values.numpy()
-        for index, idx in enumerate(ids):
-            correct_value = np.array([idx] * output_dim, dtype=np.float32)
-            self.assertTrue((values[index] == correct_value).all())
-
-        results = layer_call(layer, ids)
-        results = results.numpy()
-        for index, idx in enumerate(ids):
-            self.assertTrue((values[index] == results[index]).all())
+        call_fns = [lambda: layer.call(ids), lambda: module_call(layer, ids)]
+        correct_values = np.array(
+            [np.array([idx] * output_dim, dtype=np.float32) for idx in ids]
+        )
+        self._run_forward_pass_and_compare(
+            call_fns, correct_values, output_dim
+        )
 
         # Keras model without/with input_layer
         model_without_input_layer = tf.keras.models.Sequential([layer])
@@ -145,11 +151,14 @@ def test_embedding_layer(self):
         )
         models = [model_without_input_layer, model_with_input_layer]
         for model in models:
-            outputs = model.call(tf.constant([ids, ids]))
-            values = outputs.numpy()[1]
-            for index, idx in enumerate(ids):
-                correct_value = np.array([idx] * output_dim, dtype=np.float32)
-                self.assertTrue((values[index] == correct_value).all())
+            inputs = tf.constant([ids])
+            call_fns = [
+                lambda: model.call(inputs),
+                lambda: module_call(model, inputs),
+            ]
+            self._run_forward_pass_and_compare(
+                call_fns, correct_values, output_dim
+            )
 
     def test_embedding_layer_with_input_length(self):
         output_dim = 8
@@ -160,11 +169,16 @@ def test_embedding_layer_with_input_length(self):
         )
         ids = [[0, 1, 3, 8], [5, 3, 2, 3]]
         flatten_ids = ids[0] + ids[1]
-        values = layer.call(ids)
-        values = values.numpy().reshape(-1, output_dim)
-        for index, idx in enumerate(flatten_ids):
-            correct_value = np.array([idx] * output_dim, dtype=np.float32)
-            self.assertTrue((values[index] == correct_value).all())
+        call_fns = [lambda: layer.call(ids), lambda: module_call(layer, ids)]
+        correct_values = np.array(
+            [
+                np.array([idx] * output_dim, dtype=np.float32)
+                for idx in flatten_ids
+            ]
+        )
+        self._run_forward_pass_and_compare(
+            call_fns, correct_values, output_dim
+        )
 
     def test_embedding_layer_with_sparse_input(self):
         output_dim = 8
@@ -181,14 +195,22 @@ def test_embedding_layer_with_sparse_input(self):
             inputs = tf.SparseTensor(
                 indices=indices, values=values, dense_shape=dense_shape
             )
-            outputs = layer.call(inputs)
-            outputs = outputs.numpy()
-            correct_values = get_correct_values_for_sparse_test(
+            call_fns = [
+                lambda: layer.call(inputs),
+                lambda: module_call(layer, inputs),
+            ]
+            correct_value_single_line = get_correct_values_for_sparse_test(
                 indices, values, dense_shape, combiner
             )
-            place = 8 if combiner == "sum" else 5
-            for n, v in enumerate(correct_values):
-                self.assertAlmostEqual(outputs[n][0], v, place)
+            correct_values = np.array(
+                [
+                    np.array([idx] * output_dim, dtype=np.float32)
+                    for idx in correct_value_single_line
+                ]
+            )
+            self._run_forward_pass_and_compare(
+                call_fns, correct_values, output_dim
+            )
 
     def test_embedding_layer_with_mask_zero(self):
         output_dim = 8
@@ -231,25 +253,44 @@ def test_embedding_layer_gradient(self):
         multiply_tensor = tf.reshape(multiply_tensor, [2, 3, 8])
 
         for inputs in inputs_list:
-            with tf.GradientTape() as tape:
-                layer.set_tape(tape)
-                output = layer.call(inputs)
-                output = output * multiply_tensor
-            bet = layer.bet_ids_pair[0][0]
-            grads = tape.gradient(output, bet)
-            self.assertTrue((grads.values.numpy() == multiply_values).all())
-            self.assertTrue(
-                (
-                    layer.bet_ids_pair[0][1].numpy()
-                    == inputs.numpy().reshape(-1)
-                ).all()
-            )
-            layer.reset()
-
-        for inputs in inputs_list:
-            with tf.GradientTape() as tape:
-                layer.set_tape(tape)
-                self.assertRaises(RuntimeError, layer_call, layer, inputs)
+            call_fns = [
+                lambda: layer.call(inputs),
+                lambda: module_call(layer, inputs),
+            ]
+            for call_fn in call_fns:
+                with tf.GradientTape() as tape:
+                    layer.set_tape(tape)
+                    output = module_call(layer, inputs)
+                    output = output * multiply_tensor
+                batch_embedding = layer.embedding_and_ids[0].batch_embedding
+                grads = tape.gradient(output, batch_embedding)
+                self.assertTrue(
+                    (grads.values.numpy() == multiply_values).all()
+                )
+                self.assertTrue(
+                    (
+                        layer.embedding_and_ids[0].batch_ids.numpy()
+                        == inputs.numpy().reshape(-1)
+                    ).all()
+                )
+                layer.reset()
+
+        # Test when an embedding layer is called more than once.
+        # When an embedding layer is called more than once, `tf.function` is
+        # not supported.
+        correct_ids_list = [[0, 1, 3, 1, 2, 0], [0, 10, 3, 11, 2, 1]]
+        with tf.GradientTape() as tape:
+            layer.set_tape(tape)
+            for inputs in inputs_list:
+                _ = layer.call(inputs)
+            self.assertTrue(len(layer.embedding_and_ids) == len(inputs_list))
+            for i, correct_ids in enumerate(correct_ids_list):
+                self.assertTrue(
+                    (
+                        layer.embedding_and_ids[i].batch_ids.numpy()
+                        == correct_ids
+                    ).all()
+                )
 
     def test_embedding_layer_gradient_with_sparse_inputs(self):
         output_dim = 8
@@ -258,6 +299,10 @@ def test_embedding_layer_gradient_with_sparse_inputs(self):
         indices = [[0, 0], [1, 1], [1, 2], [2, 1], [3, 1], [3, 3], [3, 5]]
         values = [1, 1, 3, 2, 0, 2, 6]
         dense_shape = [4, 6]
+        inputs = tf.SparseTensor(
+            indices=indices, values=values, dense_shape=dense_shape
+        )
+        inputs = tf.dtypes.cast(inputs, tf.int64)
 
         multiply_values = np.ndarray(shape=(4, output_dim), dtype=np.float32)
         for i in range(4):
@@ -284,27 +329,52 @@ def test_embedding_layer_gradient_with_sparse_inputs(self):
             layer = create_embedding_layer(
                 embedding_size, output_dim, combiner=combiner
             )
-            inputs = tf.SparseTensor(
-                indices=indices, values=values, dense_shape=dense_shape
+            call_fns = [
+                lambda: layer.call(inputs),
+                lambda: module_call(layer, inputs),
+            ]
+            for call_fn in call_fns:
+                with tf.GradientTape() as tape:
+                    layer.set_tape(tape)
+                    output = module_call(layer, inputs)
+                    output = output * multiply_tensor
+                batch_embedding = layer.embedding_and_ids[0].batch_embedding
+                grads = tape.gradient(output, batch_embedding)
+                grads = grads.numpy()
+                place = 8 if combiner == "sum" else 5
+                for n, v in enumerate(correct_grads[combiner]):
+                    self.assertAlmostEqual(grads[n][0], v, place)
+                self.assertTrue(
+                    (
+                        np.absolute(
+                            layer.embedding_and_ids[0].batch_ids.numpy()
+                            - np.array([1, 3, 2, 0, 6])
+                        )
+                        < 0.00001
+                    ).all()
+                )
+                layer.reset()
+
+        # Test when an embedding layer is called more than once.
+        # When an embedding layer is called more than once, `tf.function` is
+        # not supported.
+        correct_ids = [1, 3, 2, 0, 6]
+        for combiner in combiners:
+            layer = create_embedding_layer(
+                embedding_size, output_dim, combiner=combiner
             )
             with tf.GradientTape() as tape:
                 layer.set_tape(tape)
-                output = layer.call(inputs)
-                output = output * multiply_tensor
-            bet = layer.bet_ids_pair[0][0]
-            grads = tape.gradient(output, bet)
-            grads = grads.numpy()
-            place = 8 if combiner == "sum" else 5
-            for n, v in enumerate(correct_grads[combiner]):
-                self.assertAlmostEqual(grads[n][0], v, place)
-            self.assertTrue(
-                (
-                    layer.bet_ids_pair[0][1].numpy()
-                    - np.array([1, 3, 2, 0, 6])
-                    < 0.00001
-                ).all()
-            )
-            layer.reset()
+                _ = layer.call(inputs)
+                _ = layer.call(inputs)
+                self.assertTrue(len(layer.embedding_and_ids) == 2)
+                for i in range(2):
+                    self.assertTrue(
+                        (
+                            layer.embedding_and_ids[i].batch_ids.numpy()
+                            == correct_ids
+                        ).all()
+                    )
 
 
 if __name__ == "__main__":
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
@@ -121,7 +121,10 @@ def _train_edl_embedding_with_optimizer_wrapper(
         embed_items = []
         for layer in embed_layers:
             embed_items.extend(
-                [(bet, layer.name, ids) for bet, ids in layer.bet_ids_pair]
+                [
+                    (bet, layer.name, ids)
+                    for bet, ids in layer.embedding_and_ids
+                ]
             )
 
         grads = tape.gradient(
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
@@ -6,6 +6,7 @@
 import tensorflow as tf
 
 from elasticdl.python.common.constants import JobType
+from elasticdl.python.elasticdl.layers.embedding import EmbeddingAndIds
 from elasticdl.python.master.embedding_service import EmbeddingService
 from elasticdl.python.master.servicer import MasterServicer
 from elasticdl.python.tests.in_process_master import InProcessMaster
@@ -62,15 +63,12 @@ def custom_model():
 class MockEdlEmbedding:
     def __init__(self, name):
         self._name = name
-        self.bet_ids_pair = []
+        self.embedding_and_ids = []
 
     @property
     def name(self):
         return self._name
 
-    def add_bet_and_ids(self, bet, ids):
-        self.append((bet, ids))
-
 
 class ReportBETGradientTest(unittest.TestCase):
     def _create_master_and_worker(
@@ -119,15 +117,15 @@ def test_report_bet_gradients_worker_to_master(self):
         )
 
         layer1 = MockEdlEmbedding(layer_names[0])
-        layer1.bet_ids_pair = [
-            (None, tf.constant([1, 2])),
-            (None, tf.constant([2, 3])),
+        layer1.embedding_and_ids = [
+            EmbeddingAndIds(None, tf.constant([1, 2])),
+            EmbeddingAndIds(None, tf.constant([2, 3])),
         ]
 
         layer2 = MockEdlEmbedding(layer_names[1])
-        layer2.bet_ids_pair = [
-            (None, tf.constant([3, 1])),
-            (None, tf.constant([3, 4])),
+        layer2.embedding_and_ids = [
+            EmbeddingAndIds(None, tf.constant([3, 1])),
+            EmbeddingAndIds(None, tf.constant([3, 4])),
         ]
 
         edlembed_grads = [
@@ -169,14 +167,20 @@ def test_report_bet_gradients_worker_to_master(self):
             layer1.name: tf.IndexedSlices(
                 tf.concat(edlembed_grads[:2], axis=0),
                 tf.concat(
-                    [layer1.bet_ids_pair[0][1], layer1.bet_ids_pair[1][1]],
+                    [
+                        layer1.embedding_and_ids[0].batch_ids,
+                        layer1.embedding_and_ids[1].batch_ids,
+                    ],
                     axis=0,
                 ),
             ),
             layer2.name: tf.IndexedSlices(
                 tf.concat(edlembed_grads[2:], axis=0),
                 tf.concat(
-                    [layer2.bet_ids_pair[0][1], layer2.bet_ids_pair[1][1]],
+                    [
+                        layer2.embedding_and_ids[0].batch_ids,
+                        layer2.embedding_and_ids[1].batch_ids,
+                    ],
                     axis=0,
                 ),
             ),
@@ -307,8 +311,8 @@ def test_report_bet_gradients_master_to_service(self):
     def test_get_trainable_variable(self):
         master, worker = self._create_master_and_worker()
         layer = MockEdlEmbedding("test")
-        layer.bet_ids_pair = [
-            (tf.Variable([1, 2, 3], name="test_bet"), [1, 2, 3])
+        layer.embedding_and_ids = [
+            EmbeddingAndIds(tf.Variable([1, 2, 3], name="test_bet"), [1, 2, 3])
         ]
         worker._embedding_layers = [layer]
         train_vars = worker.get_trainable_items()
diff --git a/elasticdl/python/tests/worker_test.py b/elasticdl/python/tests/worker_test.py
--- a/elasticdl/python/tests/worker_test.py
+++ b/elasticdl/python/tests/worker_test.py
@@ -1,23 +1,47 @@
+import copy
 import os
 import tempfile
 import unittest
 from contextlib import closing
 
+import mock
 import numpy as np
 import recordio
 import tensorflow as tf
+from tensorflow.keras.layers import Concatenate, Dense, Flatten
 
 from elasticdl.proto import elasticdl_pb2
 from elasticdl.python.common.constants import JobType
+from elasticdl.python.elasticdl.layers.embedding import Embedding
+from elasticdl.python.master.embedding_service import EmbeddingService
 from elasticdl.python.master.evaluation_service import EvaluationService
 from elasticdl.python.master.servicer import MasterServicer
 from elasticdl.python.master.task_dispatcher import _TaskDispatcher
 from elasticdl.python.tests.in_process_master import InProcessMaster
+from elasticdl.python.tests.mock_kv_store import MockKvStore
 from elasticdl.python.worker.worker import Worker
 
 _model_zoo_path = os.path.dirname(os.path.realpath(__file__))
 
 
+class CustomModel(tf.keras.Model):
+    def __init__(self, output_dim=16):
+        super(CustomModel, self).__init__(name="CustomModel")
+        self.output_dim = output_dim
+        self.embedding_1 = Embedding(output_dim)
+        self.embedding_2 = Embedding(output_dim)
+        self.concat = Concatenate()
+        self.dense = Dense(1)
+        self.flatten = Flatten()
+
+    def call(self, inputs, training=False):
+        f1 = self.embedding_1(inputs["f1"])
+        f2 = self.embedding_2(inputs["f2"])
+        x = self.concat([f1, f2])
+        x = self.dense(x)
+        return self.flatten(x)
+
+
 def create_recordio_file(size):
     temp_file = tempfile.NamedTemporaryFile(delete=False)
     with closing(recordio.Writer(temp_file.name)) as f:
@@ -131,6 +155,111 @@ def test_embedding_layer(self):
         )
         self.assertTrue(len(worker._embedding_layers) == 2)
 
+    def test_train_acceleration_with_embedding(self):
+        kv_store = MockKvStore()
+        model_inst = CustomModel()
+        master = MasterServicer(
+            2,
+            2,
+            tf.optimizers.SGD(0.1),
+            None,
+            init_var=model_inst.trainable_variables,
+            checkpoint_filename_for_init=None,
+            checkpoint_service=None,
+            evaluation_service=None,
+        )
+        worker = Worker(
+            1,
+            JobType.TRAINING_ONLY,
+            32,
+            _model_zoo_path,
+            model_def="embedding_test_module.EdlEmbeddingModel",
+            channel=None,
+        )
+        worker._stub = InProcessMaster(master)
+
+        inputs_list = [
+            {
+                "f1": tf.constant([[0], [1], [2]], tf.int64),
+                "f2": tf.constant([[2], [1], [0]], tf.int64),
+            },
+            {
+                "f1": tf.constant([[3], [4], [3]], tf.int64),
+                "f2": tf.constant([[2], [1], [0]], tf.int64),
+            },
+        ]
+        labels_list = [[0, 1, 0], [1, 1, 0]]
+        input_dim = 5
+        embedding_dim = 16
+        worker.set_model(model_inst)
+
+        # initialize kv store
+        for layer in model_inst.layers:
+            if isinstance(layer, Embedding):
+                name = layer.name
+                keys = [Embedding.get_key([name, i]) for i in range(input_dim)]
+                values = [
+                    np.random.rand(embedding_dim).astype(np.float32)
+                    for i in range(input_dim)
+                ]
+                kv_store.update(keys, values)
+
+        with mock.patch.object(
+            EmbeddingService, "lookup_embedding", kv_store.lookup
+        ), mock.patch.object(
+            EmbeddingService, "update_embedding", kv_store.update
+        ):
+            worker._init_embedding_layer()
+            worker._run_model_call_before_training(inputs_list[0])
+
+            # run training process without tf.function
+            correct_grads = []
+            correct_ids_list = []
+            for features, labels in zip(inputs_list, labels_list):
+                loss, grads = worker.training_process_eagerly(features, labels)
+                correct_grads.append(grads)
+                ids = {}
+                for layer in worker._embedding_layers:
+                    ids[layer.name] = layer.embedding_and_ids[0].batch_ids
+                correct_ids_list.append(ids)
+                worker._reset_embedding()
+
+            # run training process with tf.function
+            test_grads = []
+            test_ids_list = []
+            for features, labels in zip(inputs_list, labels_list):
+                self.assertFalse(worker._train_eagerly)
+                loss, grads = worker.training_process(features, labels)
+                test_grads.append(grads)
+                ids = {}
+                for layer in worker._embedding_layers:
+                    ids[layer.name] = copy.deepcopy(
+                        layer.embedding_and_ids[0].batch_ids
+                    )
+                test_ids_list.append(ids)
+                worker._reset_embedding()
+
+        # compare the gradients
+        for test_g, correct_g in zip(test_grads, correct_grads):
+            for g1, g2 in zip(test_g, correct_g):
+                if isinstance(g1, tf.IndexedSlices):
+                    self.assertTrue(
+                        (g1.values - g2.values < 0.0001).numpy().all()
+                    )
+                    self.assertTrue(
+                        tf.equal(g1.indices, g2.indices).numpy().all()
+                    )
+                else:
+                    self.assertTrue((g1 - g2 < 0.0001).numpy().all())
+
+        for test_ids, correct_ids in zip(correct_ids_list, test_ids_list):
+            for layer_name in correct_ids.keys():
+                self.assertTrue(
+                    tf.equal(test_ids[layer_name], correct_ids[layer_name])
+                    .numpy()
+                    .all()
+                )
+
 
 if __name__ == "__main__":
     unittest.main()