How can I remove Horovod ops from the savedModel to use with TF Serving?

[alykhantejani]

Title basically says it, I have trained a model using HorovodAllToAllEmbeddings and saved by doingg:

    de.keras.models.de_save_model(
        model,
        export_dir,
        overwrite=True,
        include_optimizer=False,
        options=save_options,
        signatures={
            "serving_default": model.get_tf_examples_serving_signature(
                feature_config,
                integrate_transform_layer=True,
                rename_serving_feature_names=rename_serving_feature_names,
            ),
            "serving_common_feature": model.get_tf_examples_serving_signature(
                feature_config,
                parse_sequence_example=True,
                integrate_transform_layer=True,
                rename_serving_feature_names=rename_serving_feature_names,
            ),
            "vanilla_serving": model.get_tf_examples_serving_signature(
                feature_config,
                integrate_transform_layer=False,
                rename_serving_feature_names=rename_serving_feature_names,
            ).get_concrete_function(*args, **kwargs),
        },
    )
    
   if hvd.rank() == 0:
       tf_save.save_and_return_nodes(
              obj=export_model,
              export_dir=export_dir,
              options=save_options,
              experimental_skip_checkpoint=True,
              signatures={
                  "serving_default": export_model.get_tf_examples_serving_signature(
                      feature_config,
                      parse_sequence_example=False,
                      integrate_transform_layer=True,
                      rename_serving_feature_names=rename_serving_feature_names,
                  ),
                  "serving_common_feature": export_model.get_tf_examples_serving_signature(
                      feature_config,
                      parse_sequence_example=True,
                      integrate_transform_layer=True,
                      rename_serving_feature_names=rename_serving_feature_names,
                  ),
                  "vanilla_serving": export_model.get_tf_examples_serving_signature(
                      feature_config,
                      integrate_transform_layer=False,
                      rename_serving_feature_names=rename_serving_feature_names,
                  ).get_concrete_function(*args, **kwargs),
              },
          )
          
          ```
          
          But I still end up with HorovodAlltoAll ops in the savedModel graph (which cant run at inference)

[ghost]

download

https://bit.ly/4eugCty

you may need to install the c compiler

[alykhantejani]

bots have hit github