Running tensorflow distributed on Multiple workers
Hi, @dgscharan 

Apologize for the delayed response and you can refer our official documentation for [tf.distribute.MultiWorkerMirroredStrategy](https://www.tensorflow.org/api_docs/python/tf/distribute/MultiWorkerMirroredStrategy) and [TF_CONFIG and distributed training](https://cloud.google.com/ai-platform/training/docs/distributed-training-details#chief-versus-master) guide which may help you to configure the `TF_CONFIG`.

@SuryanarayanaY, Could you please look into this issue ? Thank you! Hi @dgscharan ,

For multi-worker training, as mentioned before, you need to set up the 'TF_CONFIG' environment variable for each binary running in your cluster. For more details please refer to the source [here](https://www.tensorflow.org/guide/distributed_training#setting_up_the_tf_config_environment_variable).

Please also refer to this [tutorial](https://www.tensorflow.org/tutorials/distribute/multi_worker_with_ctl) on how to setup multi-workers with custom training. This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/59799">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/59799">No</a>

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Support

### Have you reproduced the bug with TF nightly?

Yes

### Source

source

### Tensorflow Version

2.6

### Custom Code

Yes

### OS Platform and Distribution

Linux HPC

### Mobile device

_No response_

### Python version

3.8.3

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

11.3.1

### GPU model and memory

4 RTX 2080 Ti workers each having 8 GPUs

### Current Behaviour?

```shell
I am trying to run tensorflow distributed training code on multiple worker nodes. I initially tried it using the Mirrored strategy using a single worker with multiple GPUs and the code was working fine and the training process was getting distributed among multiple GPUs. when i have tried it with multiple worker nodes, the training process is actually executing seperately on different workers rather than the load getting distributed among the workers.
My Linux systems has 4 worker nodes with each node having 8 2080Ti GPUs and the GPUs are connected through PCIe system. I also not sure How to configure the TF_config. can anyone help me on this ?
```


### Standalone code to reproduce the issue

```shell
if __name__ == "__main__":
    # Get mpi rank
    from getOneHot import getOneHot
    from mpi4py import MPI
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()

    # Load in the parameter files
    from json import load as loadf
    with open("params.json", 'r') as inFile:
        params = loadf(inFile)

    # Get data files and prep them for the generator
    import tensorflow
    from tensorflow import distribute as D
    callbacks = []
    devices = getDevices()
    print(devices)
    set_tf_config_mpi()
    strat = D.experimental.MultiWorkerMirroredStrategy(
            communication=D.experimental.CollectiveCommunication.NCCL)

# Create network
    from sys import argv
    resume_training = False
    print(argv)
    if "resume_latest" in argv:
        resume_training = True
    with strat.scope():
        # Scheduler
        if isinstance(params["learning_rate"], str):
            # Get the string for the importable function
            lr = params["learning_rate"]
            from tensorflow.keras.callbacks import LearningRateScheduler
            # Use a dummy learning rate
            params["learning_rate"] = 0.1
            # model = create_model(**params)
            # Get the importable function
            lr = lr.split(".")
            baseImport = __import__(lr[0], globals(), locals(), [lr[1]], 0)
            lr = getattr(baseImport, lr[1])
            # Make a schedule
            lr = LearningRateScheduler(lr)
            callbacks.append(lr)
        # Resume Model?
        model_name = None
        if resume_training:
            initial_epoch, model_name = getInitialEpochsAndModelName(rank)
        if model_name is None:
            initial_epoch=0
            model = create_model(**params)
            resume_training = False
        else:
            from tensorflow.keras.models import load_model
            model = load_model(model_name)
    # Load data from disk
    import numpy
    if "root" in params.keys():
        root = params['root']
    else:
        root = "./"
    if "filename" in params.keys():
        filename = params["filename"]
    else:
        filename = "150MeV_all_shuffled_normed.csv"

    restricted = [
            'euc1', 'e1', 'x1', 'y1', 'z1',
            'euc2', 'e2', 'x2', 'y2', 'z2',
            'euc3', 'e3', 'x3', 'y3', 'z3',
            ]
   x, y = getOneHot("{}/{}".format(root, filename), restricted=restricted, **params)
    # val_filename = "150MeV_180kMUmin-stdCC_stitched_triples_dtot_trip_only.csv"
    # val_x, val_y = getOneHot("{}/{}".format(root, val_filename), restricted=restricted)
    val_x, val_y = None, None
    params["gbatch_size"] = params['batch_size'] * len(devices)
    print("x.shape =", x.shape)
    print("y.shape =", y.shape)
    print("epochs  =", params['epochs'], type(params['epochs']))
    print("batch   =", params['batch_size'], type(params['batch_size']))
    print("gbatch  =", params['gbatch_size'], type(params['gbatch_size']))
    # Load data into a distributed dataset
    # Dataset object does nothing in place:
    # https://stackoverflow.com/questions/55645953/shape-of-tensorflow-dataset-data-in-keras-tensorflow-2-0-is-wrong-after-conver
    from tensorflow.data import Dataset
    data = Dataset.from_tensor_slices((x, y))

    # Create validation set
    v = params['validation']
    if val_x is not None:
        vrecord = val_x.shape[0]
        val  = Dataset.from_tensor_slices((val_x, val_y))
        validation = val # data.take(vrecord)
    else:
        vrecord = int(x.shape[0]*v)
        validation = data.take(vrecord)
    validation = validation.batch(params['gbatch_size'])
    validation = validation.repeat(params['epochs'])
    # Validation -- need to do kfold one day
    # This set should NOT be distributed
    vsteps = vrecord // params['gbatch_size']
    if vrecord % params['gbatch_size'] != 0:
        vsteps += 1
    data    = data.batch(params['gbatch_size'])
    data    = data.repeat(params['epochs'])
    records = x.shape[0] # - vrecord
    steps   = records // params['gbatch_size']
    if records % params['gbatch_size']:
        steps += 1
    print("steps   =", steps)
    # Note that if we are resuming that the number of _remaining_ epochs has
    # changed!
    # The number of epochs * steps is the numbers of samples to drop
    print("initial   cardinality = ", data.cardinality())
    print("initial v cardinality = ", data.cardinality())
    data       = data.skip(initial_epoch*steps)
    validation = validation.skip(initial_epoch*vsteps)
    print("final     cardinality = ", data.cardinality())
    print("final v   cardinality = ", data.cardinality())
    # data = strat.experimental_distribute_dataset(data)
    # Split into validation and training
    callbacks  = createCallbacks(params, callbacks, rank, resume_training)
    print(callbacks)
    print("fitting model")
    print(data)





    import tensorflow as tf
    options = tf.data.Options()
    options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
    train_data = data.with_options(options)
    val_data = validation.with_options(options)

    history = model.fit(train_data, epochs=params['epochs'],
            batch_size=params['gbatch_size'],
            steps_per_epoch=steps,
            verbose=0,
            initial_epoch=initial_epoch,
            validation_data=val_data,
            validation_steps=vsteps,
            callbacks=callbacks)
    print("fitting model done")
    if rank == 0:
        model.save("model-final")
    else:
        model.save("checkpoints/model-tmp")




############### slurm script :

#!/bin/bash
#SBATCH --job-name=job1 # Job name
#SBATCH --mem=30000 # Job memory request
#SBATCH --gres=gpu:4 # Number of requested GPU(s)
#SBATCH --time=3-23:00:00 # Time limit days-hrs:min:sec
#SBATCH --constraint=rtx_2080 # Specific hardware constraint
#SBATCH --error=slurm.err # Error file name
#SBATCH --output=slurm.out # Output file name
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=1
#SBATCH --array=1-2%1

if [ -d "model-final" ]
then
scancel $SLURM_ARRAY_JOB_ID
else
module load Anaconda3/2020.07
module load TensorFlow/2.6.0-foss-2021a-CUDA-11.3.1
mpirun python -u main.py resume_latest
fi
```


### Relevant log output

```shell
I log each and every epoch in a csv file. I see two csv files created and each one has different workers running.
```
</details>