XLA and MultiWorkerMirroredStrategy

Antonio · 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: