XLA and MultiWorkerMirroredStrategy

Antonio_1 · May 6, 2022, 1:35am

The release notes of tensorflow 2.5.0 mentions that: “XLA can now compile MirroredStrategy: the step function passed tostrategy.run can now be annoted with jit_compile=True.”

Is the same support also available for MultiWorkerMirroredStrategy? I tried adding jit_compile=True to the step function in this colab tutorial:

And I got: “RuntimeError: merge_call called while defining a new graph or a tf.function …”

Thanks!

Bhack · May 6, 2022, 11:03am

github.com

tensorflow/tensorflow/blob/v2.8.0/tensorflow/python/distribute/collective_all_reduce_strategy.py#L337


      
            distribute_lib.StrategyExtendedV1.__init__(self, container_strategy)
            self._communication_options = communication_options
            self._collective_key_base = container_strategy._collective_key_base  # pylint: disable=protected-access
            self._initialize_strategy(self._cluster_resolver)
            self._cfer_fn_cache = weakref.WeakKeyDictionary()
            self.experimental_enable_get_next_as_optional = True
            assert isinstance(self._cross_device_ops,
                              cross_device_ops_lib.CollectiveAllReduce)
          
          def _use_merge_call(self):
            """XLA is not supported for multi-worker strategy."""
            return True
          
          def _initialize_strategy(self, cluster_resolver):
            if cluster_resolver.cluster_spec().as_dict():
              self._initialize_multi_worker(cluster_resolver)
            else:
              self._initialize_local(cluster_resolver)
          
          def _initialize_local_devices(self, cluster_resolver, worker_device):
            # TODO(b/126786766): TFConfigClusterResolver returns wrong number of GPUs in

Bhack · May 6, 2022, 6:49pm

Let’s see if they approve:

https://github.com/tensorflow/tensorflow/pull/55963

Topic		Replies	Views
MirroredStrategy with jit_compile=True Keras tfdebug , tfkeras , gpu	1	366	February 22, 2024
How to solve tensorflow jit compile error General Discussion gpu , xla , help_request	2	7652	August 17, 2021
Distributed training with XLA General Discussion distributed-training , gpu , xla , help_request	1	1735	October 31, 2023
Datasets, iterators and compiled tf.functions General Discussion datasets , xla , help_request	1	1916	December 3, 2021
Debugging XLA compile General Discussion xla , help_request	1	2283	February 6, 2023

XLA and MultiWorkerMirroredStrategy

Related topics