<h3 dir="auto">System information</h3>
<ul dir="auto">
<li><strong>Have I written custom code (as opposed to using a stock example script provided in TensorFlow)</strong>:<br>
I modified the tensorflow/models repo utility function to use CollectiveAllReduceStrategy instead of MirroredStrategy. See the one-line change at: <a href="https://gist.github.com/nvcastet/60b8c0c66da4cf2949e38fc790208a1c#file-distribution_utils-py-L47">https://gist.github.com/nvcastet/60b8c0c66da4cf2949e38fc790208a1c#file-distribution_utils-py-L47</a></li>
<li><strong>OS Platform and Distribution (e.g., Linux Ubuntu 16.04)</strong>:<br>
Ubuntu 16.04.4 ppc64le</li>
<li><strong>TensorFlow installed from (source or binary)</strong>:<br>
source</li>
<li><strong>TensorFlow version (use command below)</strong>:<br>
1.11</li>
<li><strong>Python version</strong>: Python 2.7.15</li>
<li><strong>CUDA/cuDNN version</strong>: Cuda 9.2 CuDNN 7.2.1</li>
<li><strong>GPU model and memory</strong>: V100 16GB</li>
<li><strong>Exact command to reproduce</strong>: ~/models/official/mnist$ python mnist.py --num_gpus 2</li>
</ul>
<h3 dir="auto">Describe the problem</h3>
<p dir="auto">Using the MNIST model script from the tensorflow/models repo after modifying the distribution_utils.py utility file to use CollectiveAllReduceStrategy (see modification above).<br>
<code class="notranslate">python mnist.py --num_gpus 2</code> crashes with</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="2018-09-17 19:28:36.648966: W tensorflow/core/framework/op_kernel.cc:1273] OP_REQUIRES failed at collective_ops.cc:210 : Internal: Second consumer arrived for key 10003:0:0:0:0:1
         [[{{node conv2d_1/bias/replica_1/Initializer/CollectiveBcastRecv}} = CollectiveBcastRecv[T=DT_FLOAT, _class=[&quot;loc:@conv2d_1/bias/replica_1/Assign&quot;], group_key=1, group_size=2, instance_key=10003, shape=[64], _device=&quot;/job:localhost/replica:0/task:0/device:GPU:1&quot;]()]]
         [[{{node GroupCrossDeviceControlEdges_0/group_deps_27/_2}} = _Recv[client_terminated=false, recv_device=&quot;/job:localhost/replica:0/task:0/device:GPU:0&quot;, send_device=&quot;/job:localhost/replica:0/task:0/device:GPU:1&quot;, send_device_incarnation=1, tensor_name=&quot;edge_238_GroupCrossDeviceControlEdges_0/group_deps_27&quot;, tensor_type=DT_FLOAT, _device=&quot;/job:localhost/replica:0/task:0/device:GPU:0&quot;]()]]"><pre class="notranslate"><code class="notranslate">2018-09-17 19:28:36.648966: W tensorflow/core/framework/op_kernel.cc:1273] OP_REQUIRES failed at collective_ops.cc:210 : Internal: Second consumer arrived for key 10003:0:0:0:0:1
         [[{{node conv2d_1/bias/replica_1/Initializer/CollectiveBcastRecv}} = CollectiveBcastRecv[T=DT_FLOAT, _class=["loc:@conv2d_1/bias/replica_1/Assign"], group_key=1, group_size=2, instance_key=10003, shape=[64], _device="/job:localhost/replica:0/task:0/device:GPU:1"]()]]
         [[{{node GroupCrossDeviceControlEdges_0/group_deps_27/_2}} = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:1", send_device_incarnation=1, tensor_name="edge_238_GroupCrossDeviceControlEdges_0/group_deps_27", tensor_type=DT_FLOAT, _device="/job:localhost/replica:0/task:0/device:GPU:0"]()]]
</code></pre></div>
<p dir="auto">Full log of the run with stack trace: <a href="https://github.com/tensorflow/tensorflow/files/2390053/run.log">run.log</a></p>
<h3 dir="auto">System information</h3>
<ul dir="auto">
<li>
<p dir="auto"><strong>Have I written custom code (as opposed to using a stock example script provided in TensorFlow)</strong>:<br>
I slightly updated mnist.py example so that it uses CollectiveAllReduceStrategy. Updated version is <a href="https://github.com/dmitrievanthony/models/blob/ignite/official/mnist/mnist.py#L208">here</a>.</p>
</li>
<li>
<p dir="auto"><strong>OS Platform and Distribution (e.g., Linux Ubuntu 16.04)</strong>:<br>
Mac OS 10.13.3</p>
</li>
<li>
<p dir="auto"><strong>TensorFlow installed from (source or binary)</strong>:<br>
Binary</p>
</li>
<li>
<p dir="auto"><strong>TensorFlow version (use command below)</strong>:<br>
1.11.0-dev20180913</p>
</li>
<li>
<p dir="auto"><strong>Python version</strong>:<br>
3.6.3</p>
</li>
<li>
<p dir="auto"><strong>Exact command to reproduce</strong>:<br>
See <a href="https://github.com/dmitrievanthony/models/blob/ignite/official/mnist/mnist.py#L208">updated example</a>.</p>
</li>
</ul>
<h3 dir="auto">Describe the problem</h3>
<p dir="auto">Hi,</p>
<p dir="auto">I'm trying to update <code class="notranslate">mnist</code> model from <a href="https://github.com/tensorflow/models/tree/master/official/">official repository</a> so that it uses CollectiveAllReduceStrategy as it's shown in <a href="https://github.com/tensorflow/ecosystem/blob/master/distribution_strategy/keras_model_to_estimator_client.py">keras_model_to_estimator_client.py</a>. Updated example you can find <a href="https://github.com/dmitrievanthony/models/blob/ignite/official/mnist/mnist.py#L208">here</a>. Unfortunately, it fails on <code class="notranslate">deepcopy</code> of run config.</p>
<h3 dir="auto">Source code / logs</h3>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Traceback (most recent call last):
  File &quot;mnist.py&quot;, line 286, in &lt;module&gt;
    absl_app.run(main)
  File &quot;/Users/antondmitriev/anaconda3/lib/python3.6/site-packages/absl/app.py&quot;, line 274, in run
    _run_main(main, args)
  File &quot;/Users/antondmitriev/anaconda3/lib/python3.6/site-packages/absl/app.py&quot;, line 238, in _run_main
    sys.exit(main(argv))
  File &quot;mnist.py&quot;, line 280, in main
    run_mnist(flags.FLAGS)
  File &quot;mnist.py&quot;, line 226, in run_mnist
    'data_format': data_format,
  File &quot;/Users/antondmitriev/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py&quot;, line 190, in __init__
    model_dir)
  File &quot;/Users/antondmitriev/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py&quot;, line 1591, in maybe_overwrite_model_dir_and_session_config
    config = run_config.RunConfig.replace(config, session_config=session_config)
  File &quot;/Users/antondmitriev/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/run_config.py&quot;, line 849, in replace
    copy.deepcopy(self),
  File &quot;/Users/antondmitriev/anaconda3/lib/python3.6/copy.py&quot;, line 180, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File &quot;/Users/antondmitriev/anaconda3/lib/python3.6/copy.py&quot;, line 280, in _reconstruct
    state = deepcopy(state, memo)
  File &quot;/Users/antondmitriev/anaconda3/lib/python3.6/copy.py&quot;, line 150, in deepcopy
    y = copier(x, memo)
  File &quot;/Users/antondmitriev/anaconda3/lib/python3.6/copy.py&quot;, line 240, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File &quot;/Users/antondmitriev/anaconda3/lib/python3.6/copy.py&quot;, line 180, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File &quot;/Users/antondmitriev/anaconda3/lib/python3.6/copy.py&quot;, line 280, in _reconstruct
    state = deepcopy(state, memo)
  File &quot;/Users/antondmitriev/anaconda3/lib/python3.6/copy.py&quot;, line 150, in deepcopy
    y = copier(x, memo)
  File &quot;/Users/antondmitriev/anaconda3/lib/python3.6/copy.py&quot;, line 240, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File &quot;/Users/antondmitriev/anaconda3/lib/python3.6/copy.py&quot;, line 180, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File &quot;/Users/antondmitriev/anaconda3/lib/python3.6/copy.py&quot;, line 280, in _reconstruct
    state = deepcopy(state, memo)
  File &quot;/Users/antondmitriev/anaconda3/lib/python3.6/copy.py&quot;, line 150, in deepcopy
    y = copier(x, memo)
  File &quot;/Users/antondmitriev/anaconda3/lib/python3.6/copy.py&quot;, line 240, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File &quot;/Users/antondmitriev/anaconda3/lib/python3.6/copy.py&quot;, line 169, in deepcopy
    rv = reductor(4)
TypeError: can't pickle _thread._local objects"><pre class="notranslate"><code class="notranslate">Traceback (most recent call last):
  File "mnist.py", line 286, in &lt;module&gt;
    absl_app.run(main)
  File "/Users/antondmitriev/anaconda3/lib/python3.6/site-packages/absl/app.py", line 274, in run
    _run_main(main, args)
  File "/Users/antondmitriev/anaconda3/lib/python3.6/site-packages/absl/app.py", line 238, in _run_main
    sys.exit(main(argv))
  File "mnist.py", line 280, in main
    run_mnist(flags.FLAGS)
  File "mnist.py", line 226, in run_mnist
    'data_format': data_format,
  File "/Users/antondmitriev/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 190, in __init__
    model_dir)
  File "/Users/antondmitriev/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 1591, in maybe_overwrite_model_dir_and_session_config
    config = run_config.RunConfig.replace(config, session_config=session_config)
  File "/Users/antondmitriev/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/run_config.py", line 849, in replace
    copy.deepcopy(self),
  File "/Users/antondmitriev/anaconda3/lib/python3.6/copy.py", line 180, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/Users/antondmitriev/anaconda3/lib/python3.6/copy.py", line 280, in _reconstruct
    state = deepcopy(state, memo)
  File "/Users/antondmitriev/anaconda3/lib/python3.6/copy.py", line 150, in deepcopy
    y = copier(x, memo)
  File "/Users/antondmitriev/anaconda3/lib/python3.6/copy.py", line 240, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/Users/antondmitriev/anaconda3/lib/python3.6/copy.py", line 180, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/Users/antondmitriev/anaconda3/lib/python3.6/copy.py", line 280, in _reconstruct
    state = deepcopy(state, memo)
  File "/Users/antondmitriev/anaconda3/lib/python3.6/copy.py", line 150, in deepcopy
    y = copier(x, memo)
  File "/Users/antondmitriev/anaconda3/lib/python3.6/copy.py", line 240, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/Users/antondmitriev/anaconda3/lib/python3.6/copy.py", line 180, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/Users/antondmitriev/anaconda3/lib/python3.6/copy.py", line 280, in _reconstruct
    state = deepcopy(state, memo)
  File "/Users/antondmitriev/anaconda3/lib/python3.6/copy.py", line 150, in deepcopy
    y = copier(x, memo)
  File "/Users/antondmitriev/anaconda3/lib/python3.6/copy.py", line 240, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/Users/antondmitriev/anaconda3/lib/python3.6/copy.py", line 169, in deepcopy
    rv = reductor(4)
TypeError: can't pickle _thread._local objects
</code></pre></div>