[fix] The default save function for tensorflow will not be patched at this time, as this can lead to unexpected errors.

Now using de.keras.models.de_hvd_save_model to replace tf.keras.models.save_model.

Author: MoFHeka

demo/dynamic_embedding/movielens-1m-keras-with-horovod/movielens-1m-keras-with-horovod.py

@@ -451,14 +451,14 @@ def export_to_savedmodel(model, savedmodel_dir):
   #                                        proc_size=hvd.size(),
   #                                        proc_rank=hvd.rank())
 
-  # TFRA modify the Keras save function with a monkey patch.
+  # TFRA modify the Keras save function with a patch.
   # !!!! Run save_model function in all rank !!!!
-  tf.keras.models.save_model(model,
-                             savedmodel_dir,
-                             overwrite=True,
-                             include_optimizer=True,
-                             save_traces=True,
-                             options=save_options)
+  de.keras.models.de_hvd_save_model(model,
+                                    savedmodel_dir,
+                                    overwrite=True,
+                                    include_optimizer=True,
+                                    save_traces=True,
+                                    options=save_options)
 
 
 def export_for_serving(model, export_dir):
@@ -506,9 +506,9 @@ def serve(*args, **kwargs):
   #                                        proc_size=hvd.size(),
   #                                        proc_rank=hvd.rank())
 
-  # TFRA modify the Keras save function with a monkey patch.
+  # TFRA modify the Keras save function with a patch.
   # !!!! Run save_model function in all rank !!!!
-  tf.keras.models.save_model(
+  de.keras.models.de_hvd_save_model(
       model,
       export_dir,
       overwrite=True,
@@ -559,6 +559,9 @@ def train():
   # horovod callback is used to broadcast the value generated by initializer of rank0.
   hvd_opt_init_callback = de.keras.callbacks.DEHvdBroadcastGlobalVariablesCallback(
       root_rank=0)
+  ckpt_callback = de.keras.callbacks.DEHvdModelCheckpoint(
+      filepath=FLAGS.model_dir + '/weights_epoch{epoch:03d}_loss{loss:.4f}',
+      options=save_options)
   callbacks_list = [hvd_opt_init_callback, ckpt_callback]
   # The log class callback only takes effect in rank0 for convenience
   if hvd.rank() == 0:

tensorflow_recommenders_addons/dynamic_embedding/python/keras/__init__.py

@@ -1,2 +1,3 @@
 from tensorflow_recommenders_addons.dynamic_embedding.python.keras import layers
-from tensorflow_recommenders_addons.dynamic_embedding.python.keras import callbacks
+from tensorflow_recommenders_addons.dynamic_embedding.python.keras import callbacks
+from tensorflow_recommenders_addons.dynamic_embedding.python.keras import models

tensorflow_recommenders_addons/dynamic_embedding/python/keras/callbacks.py

@@ -21,6 +21,7 @@
 from tensorflow.python.keras import backend as K
 from tensorflow.python.keras import callbacks
 from tensorflow.python.keras.utils import tf_utils
+from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.util.deprecation import deprecated
@@ -108,10 +109,6 @@ def __init__(self, root_rank=0, device='', local_variables=None):
       self.register_local_var(var)
 
 
-@deprecated(
-    None, "\n!!!! Using this callback will cause a save twice error. !!!!\n"
-    "The callbacks.ModelCheckpoint for HvdAllToAllEmbedding has been deprecated, use original ModelCheckpoint instead.\n"
-    "!!!! Using this callback will cause a save twice error. !!!!\n")
 class DEHvdModelCheckpoint(callbacks.ModelCheckpoint):
 
   def __init__(self, *args, **kwargs):
@@ -129,23 +126,26 @@ def _save_de_model(self, filepath):
                         options=self._options)
     else:
       de_dir = os.path.join(filepath, "variables", "TFRADynamicEmbedding")
-      for layer in self.model.layers:
-        if hasattr(layer, "params") and isinstance(layer, HvdAllToAllEmbedding):
+      for var in self.model.variables:
+        if not hasattr(var, "params") or not isinstance(var, TrainableWrapper):
+          continue
+        if not hasattr(var.params, "_created_in_class"):
+          continue
+        de_var = var.params
+        a2a_emb = de_var._created_in_class
+        if issubclass(a2a_emb.__class__, HvdAllToAllEmbedding):
           # save Dynamic Embedding Parameters
-          logging.warning(
-              "!!!! Using this callback will cause a save twice error. !!!!\n"
-              "The callbacks.ModelCheckpoint for HvdAllToAllEmbedding has been deprecated, use original ModelCheckpoint instead.\n"
-              "!!!! Using this callback will cause a save twice error. !!!!\n")
-          layer.params.save_to_file_system(dirpath=de_dir,
-                                           proc_size=hvd.size(),
-                                           proc_rank=hvd.rank())
+          de_var.save_to_file_system(dirpath=de_dir,
+                                     proc_size=hvd.size(),
+                                     proc_rank=hvd.rank())
           # save optimizer parameters of Dynamic Embedding
-          opt_de_vars = layer.optimizer_vars.as_list() if hasattr(
-              layer.optimizer_vars, "as_list") else layer.optimizer_vars
-          for opt_de_var in opt_de_vars:
-            opt_de_var.save_to_file_system(dirpath=de_dir,
+          de_opt_vars = a2a_emb.optimizer_vars.as_list() if hasattr(
+              a2a_emb.optimizer_vars, "as_list") else a2a_emb.optimizer_vars
+          for de_opt_var in de_opt_vars:
+            de_opt_var.save_to_file_system(dirpath=de_dir,
                                            proc_size=hvd.size(),
                                            proc_rank=hvd.rank())
+    hvd.join()  # Sync for avoiding data conflict or missing rank
 
   def _save_model(self, epoch, logs):
     """Saves the model.

tensorflow_recommenders_addons/dynamic_embedding/python/keras/layers/embedding.py

@@ -550,6 +550,7 @@ def __init__(self,
     else:
       self._mpi_size = mpi_size
     super(HvdAllToAllEmbedding, self).__init__(*args, **kwargs)
+    self.params._created_in_class = self
 
   def __relocate_dense_feature__(self, ids, batch_size=None):
     """

tensorflow_recommenders_addons/dynamic_embedding/python/keras/models.py

@@ -0,0 +1,139 @@
+# Copyright 2023 The TensorFlow Recommenders-Addons Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# lint-as: python3
+
+import functools
+import os.path
+
+from tensorflow_recommenders_addons import dynamic_embedding as de
+
+try:
+  from keras.saving.saved_model import save as keras_saved_model_save
+except:
+  keras_saved_model_save = None
+from tensorflow.python.keras.saving.saved_model import save as tf_saved_model_save
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import tf_logging
+
+tf_original_save_func = tf_saved_model_save.save
+if keras_saved_model_save is not None:
+  keras_original_save_func = keras_saved_model_save.save
+
+
+def _de_keras_save_func(original_save_func,
+                        model,
+                        filepath,
+                        overwrite,
+                        include_optimizer,
+                        signatures=None,
+                        options=None,
+                        save_traces=True,
+                        *args,
+                        **kwargs):
+  """Overwrite TF Keras save function
+    Calling the TF save API for all ranks causes file conflicts, 
+    so KV files other than rank0 need to be saved by calling the underlying API separately.
+    This is a convenience function for saving HvdAllToAllEmbedding to KV files in different rank.
+  """
+  try:
+    import horovod.tensorflow as hvd
+    try:
+      hvd.rank()
+    except:
+      hvd = None
+  except:
+    hvd = None
+
+  call_original_save_func = functools.partial(
+      original_save_func,
+      model=model,
+      filepath=filepath,
+      overwrite=overwrite,
+      include_optimizer=include_optimizer,
+      signatures=signatures,
+      options=options,
+      save_traces=save_traces,
+      *args,
+      **kwargs)
+
+  def _traverse_emb_layers_and_save(hvd_rank):
+    de_dir = os.path.join(filepath, "variables", "TFRADynamicEmbedding")
+    for var in model.variables:
+      if not hasattr(var, "params"):
+        continue
+      if not hasattr(var.params, "_created_in_class"):
+        continue
+      de_var = var.params
+      a2a_emb = de_var._created_in_class
+      if issubclass(a2a_emb.__class__, de.keras.layers.HvdAllToAllEmbedding):
+        if de_var._saveable_object_creator is None:
+          if hvd_rank == 0:
+            tf_logging.warning(
+                "Please use FileSystemSaver when use HvdAllToAllEmbedding. "
+                "It will allow TFRA load KV files when Embedding tensor parallel. "
+                f"The embedding shards at each horovod rank are now temporarily stored in {de_dir}"
+            )
+        else:
+          if not isinstance(de_var.kv_creator.saver, de.FileSystemSaver):
+            # This function only serves FileSystemSaver.
+            continue
+          if hvd_rank == 0:
+            # FileSystemSaver works well at rank 0.
+            continue
+          # save Dynamic Embedding Parameters
+          de_var.save_to_file_system(dirpath=de_dir,
+                                     proc_size=hvd.size(),
+                                     proc_rank=hvd.rank())
+          # save optimizer parameters of Dynamic Embedding
+          if include_optimizer is True:
+            de_opt_vars = a2a_emb.optimizer_vars.as_list() if hasattr(
+                a2a_emb.optimizer_vars, "as_list") else a2a_emb.optimizer_vars
+            for de_opt_var in de_opt_vars:
+              de_opt_var.save_to_file_system(dirpath=de_dir,
+                                             proc_size=hvd.size(),
+                                             proc_rank=hvd.rank())
+
+  if hvd is None:
+    call_original_save_func()
+  else:
+    if hvd.rank() == 0:
+      call_original_save_func()
+    _traverse_emb_layers_and_save(hvd.rank())
+    hvd.join()  # Sync for avoiding data conflict
+
+
+def de_hvd_save_model(model,
+                      filepath,
+                      overwrite=True,
+                      include_optimizer=True,
+                      signatures=None,
+                      options=None,
+                      save_traces=True,
+                      *args,
+                      **kwargs):
+  if keras_saved_model_save is not None:
+    _save_handle = functools.partial(_de_keras_save_func,
+                                     keras_original_save_func)
+  else:
+    _save_handle = functools.partial(_de_keras_save_func, tf_original_save_func)
+  return _save_handle(model,
+                      filepath,
+                      overwrite,
+                      include_optimizer,
+                      signatures=signatures,
+                      options=options,
+                      save_traces=save_traces,
+                      *args,
+                      **kwargs)

tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/horovod_sync_train_test.py

@@ -321,9 +321,11 @@ def common_all_to_all_embedding_trainable_v2(self, base_opt, test_opt, name):
         if hvd.rank() == 0:
           if os.path.exists(save_dir):
             shutil.rmtree(save_dir)
-        hvd.broadcast(tensor=tf.constant(1),
-                      root_rank=0)  # Sync for avoiding files conflict
-        base_model.save(save_dir, options=save_options)
+        hvd.join()  # Sync for avoiding files conflict
+        # base_model.save(save_dir, options=save_options)
+        de.keras.models.de_hvd_save_model(base_model,
+                                          save_dir,
+                                          options=save_options)
         del base_model
         new_base_model = get_emb_sequential_model(
             de.keras.layers.HvdAllToAllEmbedding,
@@ -333,13 +335,11 @@ def common_all_to_all_embedding_trainable_v2(self, base_opt, test_opt, name):
             bp_v2=False,
             kv_creator=kv_creator,
             name='all2all_emb')
-        hvd.broadcast(tensor=tf.constant(1),
-                      root_rank=0)  # Sync for avoiding files conflict
+        hvd.join()  # Sync for avoiding files conflict
         new_base_model.load_weights(save_dir + '/variables/variables')
         new_a2aemb_size = new_base_model.layers[0].params.size()
         self.assertEqual(a2aemb_size, new_a2aemb_size)
-        hvd.broadcast(tensor=tf.constant(1),
-                      root_rank=0)  # Sync for avoiding files conflict
+        hvd.join()  # Sync for avoiding files conflict
 
 
 if __name__ == "__main__":