RuntimeError: `merge_call` called while defining a new graph or a tf.function.

[innat]

System information.

• Have I written custom code: Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Colab
• TensorFlow version (use command below): 2.6, 2.9

---

Describe the problem

I have code that works fine but gives the following error if I use with strategy.scope().

RuntimeError: merge_call called while defining a new graph or a tf.function. This can often happen if the function fn passed to strategy.run() contains a nested @tf.function, and the nested @tf.function contains a synchronization point, such as aggregating gradients (e.g, optimizer.apply_gradients), or if the function fn uses a control flow statement which contains a synchronization point in the body. Such behaviors are not yet supported. Instead, please avoid nested tf.functions or control flow statements that may potentially cross a synchronization boundary, for example, wrap the fn passed to strategy.run or the entire strategy.run inside a tf.function or move the control flow out of fn

Describe the expected behavior

I think, It should work.

• Do you want to contribute a PR? (yes/no): No.
• If yes, please read this page for instructions

Standalone code to reproduce the issue

The code is for gradient accumulation techniques. Here it is done by overriding the trian_step with fit method. This code works fine (as said above) without with strategy.scope(). Now, I like to use it for multi-gpu cases, and so I use strategy scope but ened up the the above mentioned error.

Gist.

Follow-up Questions

1. In the above gist, for multi-gpu training, do I need to adjust anythong for BATCH_SIZE = 32 * strategy.num_replicas_in_sync inside the train_step method? Or it will be handled auto?
2. In the above gist, I use mixed precisoin technique, and so I also wrap (as described) optimizer with LossScaleOptimizer and use optimizer.get_scaled_loss(loss) and optimizer.get_unscaled_gradients(gradients).
   But the official documentation talks about normal fit and custom loop training cases. In case of custom loop, it's suggested to wrap the optimizer and scale the loss and gradient but what about the combination of fit and custom loop (overriding train_step)? Does it sill need to wrap the optimizer and scale the loss and gradient or it will be handled by the API?

---

Others: #107 cc @chenmoneygithub @nikitamaia @bhack

[innat]

cc. @andreped
andreped/GradientAccumulator#5

[andreped]

Hello, @innat!

I have not had the time to add multi-GPU support to GradientAccumulator, but can make an attempt at it today.

However batch training + gradient accumulation + mixed preicision works seemlessly.

I have been using it for various projects already.

[innat]

Thanks for your response.
I like to extend the above gist (custom fit + override the train step func) for multi-gpu (and hopefully tpu). Also, I've updated my query on stack with bounty (HERE).

I noticed that @stefan-falk also faced similar error tensorflow/tensorflow#50454 that I reported above. He tried many ways, HERE, it may give some insight.

Regarding the mixed precison, as I said, I was wondersing if we need to call opt.get_scaled_loss and opt.get_unscaled_gradients. In official doc, it is said to do so in custom loop training only.

cc. @MrForExample

[andreped]

Regarding the mixed precison, as I said, I was wondersing if we need to call opt.get_scaled_loss and opt.get_unscaled_gradients. In official doc, it is said to do so in custom loop training only.

Hmm, that's interesting. However, can't it be argued that overloading the train_step actually introduces a custom training loop? I mean, that method could do anything at this point. But I agree, it is not so clear from the documentations. If anyone wishes to read further, see here.

---

Will start on the multi-GPU support now. Did you have a gist I could use for debugging/testing, @innat? Also note that the GradientAccumulator (without multi-GPU) also works with TPUs. But I am only able to run tests locally, as I doubt I am allowed to use multi-GPUs in a single colab session.

[innat]

Here is a gist, (also mentioned above).
For multi-gpu (and tpu), you can use kaggle env. It now provides multi-gpu env.

[bhack]

As mentioned in the other ticket Graphcore had a design as optimizer wrapper including cross replica:

https://github.com/graphcore/tensorflow/blob/r2.6/sdk-release-3.1/tensorflow/python/ipu/optimizers/gradient_accumulation_optimizer.py

/Cc @georgepaw

[SuryanarayanaY]

As the error suggests aggregating gradients inside nested tf.function which is not yet supported as per the error.

RuntimeError: merge_call called while defining a new graph or a tf.function. This can often happen if the function fn passed to strategy.run() contains a nested @tf.function, and the nested @tf.function contains a synchronization point, such as aggregating gradients (e.g, optimizer.apply_gradients), or if the function fn uses a control flow statement which contains a synchronization point in the body. Such behaviors are not yet supported. Instead, please avoid nested tf.functions or control flow statements that may potentially cross a synchronization boundary, for example, wrap the fn passed to strategy.run or the entire strategy.run inside a tf.function or move the control flow out of fn. If you are subclassing a tf.keras.Model, please avoid decorating overridden methods test_step and train_step in tf.function

Hence I tried the code in Eager mode by setting model.compile(run_eagerly=True) to check and it works fine in Eager mode.Please refer to attached log below.First the code works with model.compile(run_eagerly=True) which executes fine and the same code tested with model.compile(run_eagerly=False) and execution terminated immediately with runtime error.
This testing done on 2-GPU machine

(tf) suryanarayanay@ubuntu-20-04-test-gpu-surya:~$ python 17429_grad_accumulation_on_multi_gpu_r1.py 
2023-01-20 05:39:37.545128: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-01-20 05:39:38.049160: I tensorflow/core/util/port.cc:104] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2023-01-20 05:39:40.134530: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: :/home/suryanarayanay/miniconda3/envs/tf/lib/
2023-01-20 05:39:40.134656: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: :/home/suryanarayanay/miniconda3/envs/tf/lib/
2023-01-20 05:39:40.134680: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
2023-01-20 05:39:44.877039: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 38285 MB memory:  -> device: 0, name: NVIDIA A100-SXM4-40GB, pci bus id: 0000:00:04.0, compute capability: 8.0
2023-01-20 05:39:44.880513: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 38397 MB memory:  -> device: 1, name: NVIDIA A100-SXM4-40GB, pci bus id: 0000:00:05.0, compute capability: 8.0
run_eagerly=True in model.compile()
Epoch 1/3
WARNING:tensorflow:Using MirroredStrategy eagerly has significant overhead currently. We will be working on improving this in the future, but for now please wrap `call_for_each_replica` or `experimental_run` or `run` inside a tf.function to get the best performance.
    1/10000 [..............................] - ETA: 7:36:28 - loss: 2.3965 - accuracy: 0.0000e+00WARNING:tensorflow:Using MirroredStrategy eagerly has significant overhead currently. We will be working on improving this in the future, but for now please wrap `call_for_each_replica` or `experimental_run` or `run` inside a tf.function to get the best performance.
WARNING:tensorflow:Using MirroredStrategy eagerly has significant overhead currently. We will be working on improving this in the future, but for now please wrap `call_for_each_replica` or `experimental_run` or `run` inside a tf.function to get the best performance.
    3/10000 [..............................] - ETA: 7:06 - loss: 2.2653 - accuracy: 0.1111       WARNING:tensorflow:Using MirroredStrategy eagerly has significant overhead currently. We will be working on improving this in the future, but for now please wrap `call_for_each_replica` or `experimental_run` or `run` inside a tf.function to get the best performance.
WARNING:tensorflow:Using MirroredStrategy eagerly has significant overhead currently. We will be working on improving this in the future, but for now please wrap `call_for_each_replica` or `experimental_run` or `run` inside a tf.function to get the best performance.
    9/10000 [..............................] - ETA: 7:01 - loss: 2.3390 - accuracy: 0.14812023-01-20 05:39:51.483722: W tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:115] *** WARNING *** You are using ptxas 10.1.243, which is older than 11.1. ptxas before 11.1 is known to miscompile XLA code, leading to incorrect results or invalid-address errors.

You may not need to update to CUDA 11.1; cherry-picking the ptxas binary is often sufficient.
2023-01-20 05:39:51.485528: W tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:234] Falling back to the CUDA driver for PTX compilation; ptxas does not support CC 8.0
2023-01-20 05:39:51.485549: W tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:237] Used ptxas at ptxas
2023-01-20 05:39:51.694080: W tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:115] *** WARNING *** You are using ptxas 10.1.243, which is older than 11.1. ptxas before 11.1 is known to miscompile XLA code, leading to incorrect results or invalid-address errors.

You may not need to update to CUDA 11.1; cherry-picking the ptxas binary is often sufficient.
2023-01-20 05:39:52.020570: W tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:115] *** WARNING *** You are using ptxas 10.1.243, which is older than 11.1. ptxas before 11.1 is known to miscompile XLA code, leading to incorrect results or invalid-address errors.

You may not need to update to CUDA 11.1; cherry-picking the ptxas binary is often sufficient.
WARNING:tensorflow:5 out of the last 5 calls to <function _BaseOptimizer._update_step_xla at 0x7f6a1c0b3160> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
WARNING:tensorflow:6 out of the last 6 calls to <function _BaseOptimizer._update_step_xla at 0x7f6a1c0b3160> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
2023-01-20 05:39:52.275011: W tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:115] *** WARNING *** You are using ptxas 10.1.243, which is older than 11.1. ptxas before 11.1 is known to miscompile XLA code, leading to incorrect results or invalid-address errors.

You may not need to update to CUDA 11.1; cherry-picking the ptxas binary is often sufficient.
   48/10000 [..............................] - ETA: 13:48 - loss: 2.1689 - accuracy: 0.2639WARNING:tensorflow:5 out of the last 5 calls to <function _apply_all_reduce.<locals>._all_reduce at 0x7f6a145b11f0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
   58/10000 [..............................] - ETA: 12:41 - loss: 2.1184 - accuracy: 0.3190WARNING:tensorflow:6 out of the last 6 calls to <function _apply_all_reduce.<locals>._all_reduce at 0x7f6a145101f0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
10000/10000 [==============================] - 459s 46ms/step - loss: 0.2902 - accuracy: 0.9181
Epoch 2/3
10000/10000 [==============================] - 456s 46ms/step - loss: 0.1310 - accuracy: 0.9618
Epoch 3/3
10000/10000 [==============================] - 457s 46ms/step - loss: 0.0934 - accuracy: 0.9728
run_eagerly=False in model.compile()
Epoch 1/2
Traceback (most recent call last):
  File "/home/suryanarayanay/17429_grad_accumulation_on_multi_gpu_r1.py", line 94, in <module>
    custom_model.fit(x_train, y_train, batch_size=6, epochs=2, verbose = 1)
  File "/home/suryanarayanay/miniconda3/envs/tf/lib/python3.9/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/tmp/__autograph_generated_file6ty2k0r9.py", line 15, in tf__train_function
    retval_ = ag__.converted_call(ag__.ld(step_function), (ag__.ld(self), ag__.ld(iterator)), None, fscope)
  File "/home/suryanarayanay/17429_grad_accumulation_on_multi_gpu_r1.py", line 42, in train_step
    tf.cond(
  File "/home/suryanarayanay/17429_grad_accumulation_on_multi_gpu_r1.py", line 49, in apply_accu_gradients
    self.optimizer.apply_gradients(zip(self.gradient_accumulation, self.trainable_variables))
RuntimeError: in user code:

    File "/home/suryanarayanay/miniconda3/envs/tf/lib/python3.9/site-packages/keras/engine/training.py", line 1249, in train_function  *
        return step_function(self, iterator)
    File "/home/suryanarayanay/miniconda3/envs/tf/lib/python3.9/site-packages/keras/engine/training.py", line 1233, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/home/suryanarayanay/miniconda3/envs/tf/lib/python3.9/site-packages/keras/engine/training.py", line 1222, in run_step  **
        outputs = model.train_step(data)
    File "/home/suryanarayanay/17429_grad_accumulation_on_multi_gpu_r1.py", line 42, in train_step
        tf.cond(
    File "/home/suryanarayanay/17429_grad_accumulation_on_multi_gpu_r1.py", line 49, in apply_accu_gradients
        self.optimizer.apply_gradients(zip(self.gradient_accumulation, self.trainable_variables))
    File "/home/suryanarayanay/miniconda3/envs/tf/lib/python3.9/site-packages/keras/mixed_precision/loss_scale_optimizer.py", line 1301, in apply_gradients
        grads_and_vars = self._optimizer.aggregate_gradients(grads_and_vars)
    File "/home/suryanarayanay/miniconda3/envs/tf/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/optimizer.py", line 1105, in aggregate_gradients
        return optimizer_utils.all_reduce_sum_gradients(grads_and_vars)
    File "/home/suryanarayanay/miniconda3/envs/tf/lib/python3.9/site-packages/keras/optimizers/optimizer_v2/utils.py", line 42, in all_reduce_sum_gradients
        reduced = tf.distribute.get_replica_context().merge_call(

    RuntimeError: `merge_call` called while defining a new graph or a tf.function. This can often happen if the function `fn` passed to `strategy.run()` contains a nested `@tf.function`, and the nested `@tf.function` contains a synchronization point, such as aggregating gradients (e.g, optimizer.apply_gradients), or if the function `fn` uses a control flow statement which contains a synchronization point in the body. Such behaviors are not yet supported. Instead, please avoid nested `tf.function`s or control flow statements that may potentially cross a synchronization boundary, for example, wrap the `fn` passed to `strategy.run` or the entire `strategy.run` inside a `tf.function` or move the control flow out of `fn`. If you are subclassing a `tf.keras.Model`, please avoid decorating overridden methods `test_step` and `train_step` in `tf.function`.

(tf) suryanarayanay@ubuntu-20-04-test-gpu-surya:~$ 

@innat Whether eager mode is OK for you though it has performance issue but it seems this works fine here.

[innat]

@SuryanarayanaY
Thanks for the test. Please note, eager mode is a nice option to test the code with some circumstances. But It should not be treated as a solutions as it brings lots of performance cost (raised a ticket regarding cost of eager mode).

[innat]

[Info add]
SRC https://keras.io/api/optimizers/

... the nested @tf.function contains a synchronization point, such as aggregating gradients (e.g, optimizer.apply_gradients), ...

Optimizer.apply_gradients(
    grads_and_vars, name=None, skip_gradients_aggregation=False, **kwargs
)

• skip_gradients_aggregation: If true, gradients aggregation will not be performed inside optimizer. Usually this arg is set to True when you write custom code aggregating gradients outside the optimizer.