Problem about distributed training with XLA compiling.

[sseung0703]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow):
  • custom layer and custom training step
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04):
  • I have tested on Windows 10, Ubuntu 16.04, and Ubuntu 18.04.
• TensorFlow installed from (source or binary):
  • both
• TensorFlow version (use command below):
  • I have tried TF 2.4, 2.5 distributed version and source installed 2.4
• Python version:
  • 3.7
• Bazel version (if compiling from source):
  • 3.5.0
• GCC/Compiler version (if compiling from source):
  • 7.5
• CUDA/cuDNN version:
  • 10.1 and 11.0
• GPU model and memory:
  • 1080ti x4

Describe the current behavior
When I train my model on multi-gpu with XLA compiling below error is occurred.

Training starts
Traceback (most recent call last):
  File "FFP_/train_w_pruning.py", line 76, in <module>
    train_step(*data)
  File "/home/cvip/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 787, in __call__
    result = self._call(*args, **kwds)
  File "/home/cvip/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 854, in _call
    filtered_flat_args, self._concrete_stateful_fn.captured_inputs)  # pylint: disable=protected-access
  File "/home/cvip/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 1920, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager))
  File "/home/cvip/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 561, in call
    ctx=ctx)
  File "/home/cvip/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/execute.py", line 60, in quick_execute
    inputs, attrs, num_outputs)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Trying to access resource ResNet/conv/kernel/replica_1_879 located in device /job:localhost/replica:0/task:0/device:GPU:0 [Op:__inference_train_step_dist_88943]

Describe the expected behavior
I want to compile my multi-gpu code but it seems unavailable.

Standalone code to reproduce the issue
https://github.com/sseung0703/TF2-multi-gpu-training

[amahendrakar]

@sseung0703,
The code provided in the GitHub repo is fairly complex hence it would be difficult for us to pinpoint the issue. Could you please get the example down to the simplest possible repro? That will allow us to determine the source of the issue easily. Thanks!

[sseung0703]

@amahendrakar
Sorry for the complex example. I write a code as simple as possible as below.
It requires three codes in my repo, i.e., nets/ResNet.py, nets/tcl.py, and op_utils.py
It is still complex, but I think they should be checked because the problem may be in custom layers or a training step.

import os, argparse
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

import numpy as np
import tensorflow as tf
tf.debugging.set_log_device_placement(False)

import op_utils
from nets import ResNet

parser = argparse.ArgumentParser(description='')

parser.add_argument("--gpu_id", default= [0], type=int, nargs = '+')
parser.add_argument("--compile", default=True, action = 'store_true')
parser.add_argument("--learning_rate", default=1e-1, type=float)
parser.add_argument("--weight_decay", default=5e-4, type=float)
parser.add_argument("--batch_size", default=128, type=int)
args = parser.parse_args()

if __name__ == '__main__':
    gpus = tf.config.list_physical_devices('GPU')
    tf.config.set_visible_devices([tf.config.list_physical_devices('GPU')[i] for i in args.gpu_id], 'GPU')
    for gpu_id in args.gpu_id:
        tf.config.experimental.set_memory_growth(gpus[gpu_id], True)
    devices = ['/gpu:{}'.format(i) for i in args.gpu_id]
    strategy = tf.distribute.MirroredStrategy(devices, cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())

    with strategy.scope():
        options = tf.data.Options()
        options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATA

        def inference(image, label):
            image = tf.cast(image, tf.float32)
            image = (image-np.array([113.9,123.0,125.3]))/np.array([66.7,62.1,63.0])
            return image, label

        from tensorflow.keras.datasets.cifar100 import load_data
        (train_images, train_labels), (test_images, test_labels) = load_data()

        test_ds = tf.data.Dataset.from_tensor_slices((test_images, test_labels))
        test_ds = test_ds.map(inference, num_parallel_calls=tf.data.experimental.AUTOTUNE)
        test_ds = test_ds.batch(args.batch_size)
        test_ds = test_ds.with_options(options)
        test_ds = test_ds.prefetch(tf.data.experimental.AUTOTUNE)

        model = ResNet.Model(num_layers = 56, num_class = 100, name = 'ResNet', trainable = True)

        train_step, train_loss, train_accuracy, optimizer = op_utils.Optimizer(args, model, strategy)

        for step, data in enumerate(test_ds):
            train_step(*data)
            print('check')

[amahendrakar]

@sseung0703,
I was able to run the code without any issues on TF v2.4 and CUDA 11/cuDNN 8. Please check the below screenshot for reference.

Could you please create a new virtual environment and check if you are facing the same error in that as well? Thanks!

[sseung0703]

@amahendrakar
Did you check the above code on multi-GPU?

[amahendrakar]

@sseung0703,
Yes, I ran the code on a machine with two Tesla V100s. Thanks!