Multiple GPU training with a DeepMac model gets stuck

Prerequisites

Please answer the following questions for yourself before submitting an issue.

• [y] I am using the latest TensorFlow Model Garden release and TensorFlow 2.
• [y] I am reporting the issue to the correct repository. (Model Garden official or research directory)
• [y] I checked to make sure that this issue has not already been filed.

1. The entire URL of the file you are using

https://github.com/tensorflow/models/tree/745823259762c44a3f2e80cafa445af80c67309a/research/object_detection

2. Describe the bug

Multiple GPU training with a DeepMac model gets stuck, nothing trains, it seems that threading manager simply waits (see error log below)

However, multiple GPU training with other models such as EfficientDet and MobileNet works fine.
Also, single GPU training with the DeepMac Model also works.

3. Steps to reproduce

I run a script very similar to model_main_tf2.py, but that also builds tfrecords and adjusts model pipeline config at runtime. As noted above, training with other models with multiple gpus works, training DeepMac with single gpu also works.

The multi-gpu strategy is this:

    strategy = tf.compat.v2.distribute.MirroredStrategy(
        devices=device_list, cross_device_ops=tf.distribute.HierarchicalCopyAllReduce()
    )

4. Expected behavior

Training to work as usual

5. Additional context

Interupting the hanged program says this is the line that hangs:

https://github.com/tensorflow/models/blob/c626177d6ba65a211b8f791d612cdcf8b9c0fe7e/research/object_detection/model_lib_v2.py#L685

Here is the readout once I ctrl+c the program:

losses_dict = _dist_train_step(train_input_iter)                      
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/util/traceback_utils.py", line 150, in error_handler                 
    return fn(*args, **kwargs)                                                                                                                                        
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/def_function.py", line 915, in __call__
    result = self._call(*args, **kwds)                                                                                                                                
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/def_function.py", line 963, in _call                                                               self._initialize(args, kwds, add_initializers_to=initializers)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/def_function.py", line 785, in _initialize           
    self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/function.py", line 2480, in _get_concrete_function_internal_garbage_collected
    graph_function, _ = self._maybe_define_function(args, kwargs)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/function.py", line 2711, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)                                                                                                        
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/function.py", line 2627, in _create_graph_function
    func_graph_module.func_graph_from_py_func(        
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/framework/func_graph.py", line 1141, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/def_function.py", line 677, in wrapped_fn 
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)                       
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/framework/func_graph.py", line 1116, in autograph_handler         
    return autograph.converted_call(                                               
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/impl/api.py", line 439, in converted_call                  
    result = converted_f(*effective_args, **kwargs)
  File "/tmp/__autograph_generated_fileom5bzqf0.py", line 41, in tf___dist_train_step                                           
    ag__.if_stmt((ag__.ld(num_steps_per_iteration) > 1), if_body, else_body, get_state_1, set_state_1, (), 0)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/operators/control_flow.py", line 1341, in if_stmt     
    _py_if_stmt(cond, body, orelse)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/operators/control_flow.py", line 1394, in _py_if_stmt
    return body() if cond else orelse()                                            
  File "/tmp/__autograph_generated_fileom5bzqf0.py", line 36, in if_body
    ag__.for_stmt(ag__.converted_call(ag__.ld(tf).range, ((ag__.ld(num_steps_per_iteration) - 1),), None, fscope), None, loop_body, get_state, set_state, (), {'iterat
e_names': '_'})                                                                                                                                                       
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/operators/control_flow.py", line 418, in for_stmt
    _tf_range_for_stmt(iter_, extra_test, body, get_state, set_state,                                                                                                 
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/operators/control_flow.py", line 654, in _tf_range_for_stmt
    _tf_while_stmt(                                                                
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/operators/control_flow.py", line 1277, in _tf_while_stmt
    final_loop_vars = control_flow_ops.while_loop(         
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/ops/control_flow_ops.py", line 2705, in while_loop
    return while_v2.while_loop(
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/ops/while_v2.py", line 213, in while_loop
    body_graph = func_graph_module.func_graph_from_py_func(
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/framework/func_graph.py", line 1141, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/ops/while_v2.py", line 198, in wrapped_body
    outputs = body(
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/operators/control_flow.py", line 1248, in aug_body
    body()
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/operators/control_flow.py", line 630, in aug_body
    body(iterate)
  File "/tmp/__autograph_generated_fileom5bzqf0.py", line 34, in loop_body
    ag__.converted_call(ag__.ld(_sample_and_train), (ag__.ld(strategy), ag__.ld(train_step_fn), ag__.ld(data_iterator)), None, fscope)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/impl/api.py", line 441, in converted_call
    result = converted_f(*effective_args)
  File "/tmp/__autograph_generated_fileerjyjreo.py", line 28, in tf___sample_and_train
    ag__.if_stmt(ag__.converted_call(ag__.ld(hasattr), (ag__.ld(tf).distribute.Strategy, 'run'), None, fscope), if_body, else_body, get_state, set_state, ('per_replica_losses_dict',), 1)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/operators/control_flow.py", line 1341, in if_stmt
    _py_if_stmt(cond, body, orelse)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/operators/control_flow.py", line 1394, in _py_if_stmt
    return body() if cond else orelse()
  File "/tmp/__autograph_generated_fileerjyjreo.py", line 22, in if_body
    per_replica_losses_dict = ag__.converted_call(ag__.ld(strategy).run, (ag__.ld(train_step_fn),), dict(args=(ag__.ld(features), ag__.ld(labels))), fscope)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/impl/api.py", line 377, in converted_call
    return _call_unconverted(f, args, kwargs, options)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/impl/api.py", line 458, in _call_unconverted
    return f(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/distribute/distribute_lib.py", line 1312, in run
    return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/distribute/distribute_lib.py", line 2888, in call_for_each_replica
    return self._call_for_each_replica(fn, args, kwargs)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/distribute/mirrored_strategy.py", line 676, in _call_for_each_replica
    return mirrored_run.call_for_each_replica(
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/distribute/mirrored_run.py", line 101, in call_for_each_replica
    return _call_for_each_replica(strategy, fn, args, kwargs)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/distribute/mirrored_run.py", line 283, in _call_for_each_replica
    coord.join(threads)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/training/coordinator.py", line 385, in join
    six.reraise(*self._exc_info_to_raise)
  File "/usr/lib/python3/dist-packages/six.py", line 703, in reraise
    raise value
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/training/coordinator.py", line 293, in stop_on_exception
    yield
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/distribute/mirrored_run.py", line 225, in _call_for_each_replica
    t.has_paused.wait()
  File "/usr/lib/python3.8/threading.py", line 558, in wait
    signaled = self._cond.wait(timeout)
  File "/usr/lib/python3.8/threading.py", line 302, in wait
    waiter.acquire()
KeyboardInterrupt

6. System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 20.04.4 LTS
• Mobile device name if the issue happens on a mobile device:N/A
• TensorFlow installed from (source or binary): Docker Container
• TensorFlow version (use command below): v2.9.0-18-gd8ce9f9c301 2.9.1
• Python version: 3.8.10
• Bazel version (if compiling from source):N/A
• GCC/Compiler version (if compiling from source):N/A
• CUDA/cuDNN version: cuda-11.2
• GPU model and memory: 6 NVIDIA TITAN RTX, 24220MiB each