Using TensorFlow on Clusters/Supercomputers.
We are presently amidst an effort to refresh the documentation for distributed training. We plan to include guidance on setting up workflows on platforms like Google Cloud, along with clearer recommendations about when to use one distribution strategy vs. another.

In the meantime, the easy answer is that "yes," TF support training on multiple network-linked machines with CPUs, GPUs, or TPUs. TF can also leverage the dedicated high-speed interconnects of TPUs and GPUs (e.g., NCCL). If you're planning to train on multiple machines without GPUs/TPUs, you may consider using `ParameterServerStrategy`. Here is the relevant section of the docs for setting up a cluster:
https://www.tensorflow.org/tutorials/distribute/parameter_server_training#real_clusters

Please let me know if you have any further questions or topics you would like the docs to cover. Hello @jszaday 

Sorry for the late reply, I'm currently on a holiday. Thank you so much for your response. 

I'm currently having trouble understanding a few more things from the docs I hope if you are able to clear me up. 

I'm not able to understand how different distribution strategies work. For example - if each device (GPU or a separate computer system node) creates a copy of the model and calculates a specific subset of gradient which master worker collects and forms a union gradient (or full model gradient) and applies it to the model weights, then why batches are segmented? Because few parts of the documentation say - the batches are divided for each device and then averaged. I'm not able to understand these two things going simultaneously here, because it's impossible to calculate gradient of each portion of the model with a batch of data which is different for each device. 


This is because I've plans to implement a certain special optimizer which requires variable batch size, where the order of dataset is important and distribution and averaging wouldn't work - but partial gradient calculation at each worker might work. 


Thank you so much. Hopefully you might clear my this doubt as well. 
 TF's distribution strategies primarily employ data parallelism, wherein all the devices run a mirrored copy of the model and train on a local subset (or _shard_) of an unordered dataset. For example, if one distributes a dataset of 128 samples across eight devices, each device would train on 16 (random) samples. The cross-device weight update aggregation at the end of each training step ensures that all the devices have identical weights at the start of the next training step. Thus, one can imagine that each device shares any insights it learns from its samples with all the other devices during aggregation. Since all the devices' copies of the models are always kept in sync and, by aggregation, learn from all the samples in the dataset, the numerical results of this scheme are, theoretically, identical to those from training on a single device (for unordered datasets). 

That's all to say that your dataset may not be eligible for data parallelism if the model/optimizer must consume its samples in a particular order. Instead of processing samples in parallel, each device would have to wait on the weight updates from the previous device, inducing serialization.

Data parallelism is a special kind of model parallelism. One can imagine other model-parallel schemes that preserve the ordering of a dataset. For example, one may shard the individual samples of a dataset rather than the dataset itself; then, all the devices in the system may process one sample at a time, in parallel. TensorFlow's experimental API for model parallelism is called DTensor. Here, it's worth noting that variable batch sizes and order-sensitive datasets stray from the conventional path for good distributed performance in TF. Hello @jszaday 

Sorry for one more stupid question. I wish to know:

Suppose a dataset of 1024 instances, I use batch size = 1024 using Keras API like this:
`keras.fit(X, Y, batch_size = 1024, epochs = 1)`

on number of GPUs = 2.

Does that mean, each GPU would calculate the model weights of 512 instances and average them in `tf.distribute.MirroredStrategy()` in Keras to get the final model weight if trained on batch_size = 1024?

Sorry if I'm misunderstanding anything here... There are no stupid questions--I had to double-check on this myself.

The answer is that it depends on how the dataset is sharded.

When using distributed datasets:
- Each device will run a batch of `global_batch_size // strategy.num_replicas_in_sync` instances; see [this doc](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy#experimental_distribute_dataset).
- You should not specify `batch_size` as an argument to `model.fit`.

Otherwise, the semantics are less clear; I'd recommend building a dataset to be safe:

```
ds = tf.data.Dataset.from_tensor_slices((X, Y)).batch(1024)
ds = strategy.experimental_distribute_dataset(ds)
model.fit(ds, epochs=1, steps_per_epoch=1)
``` Hello @jszaday 

Thank you so much for the clarity on it.

See the problem is that - the specific problem we are working on is a special case of non-convex optimization over very powerful hardware. The optimizer is specific to two things - 1. Variable batch sizes 2. Distribution over the hardware.

While your advice of using `tf.data.Dataset` to load the dataset is helpful, it doesn't give full control over the dataset. For some reason, I'd need to stick to batch_size = 1024 parameter on Keras API. I wish to know if I'm using the same model.fit() under strategy.scope() block, with batch_size parameter to train a model, will Multi-training strategy would work? We'll use some sort of Python generators to load dataset instances in some specific order.

Second, to deal with variable batch sizes, I guess, I need to put an upper bound to the optimizer batch size and mask the non required batches of data instances with 0 and train them on GPUs distributed (is it possible, from my experience - Custom Optimizers use `update_step` function doesn't has feature to deal with batch indices of the tensor).

Thank you Apologies for the delayed response.

The snippet I provided earlier is essentially what Keras does internally when it's given (X, Y) as inputs; it creates a dataset through the strategy, and batches it using the `batch_size` argument to `model.fit`. Subsequently, the `batch_size` argument will always represents the global batch size, and each device will run a shard of `batch_size // num_replicas` instances. ([Relevant link](https://github.com/keras-team/keras/blob/master/keras/engine/data_adapter.py#L1323).) Thank you so much @jszaday , that clears up my confusion.

My second concern is - batch-level editing. I want to mask a certain batch number of the inputs passed to the `model.fit()` to 0 using a custom optimizer. Is there any way to edit certain batches to 0 while keeping the rest of the gradients? As far I understand now since the batches are distributed among the devices, it becomes much more difficult to control i.e. assign grads = 0 to i-th entry of batch and keep the rest.

Thank you. Rather than trying to correlate input batches with masks, there may be a way to pass the mask as an input to the model and have it "skip" past all the other layers to be applied immediately before the outputs. However, I would recommend routing further inquiries into masking with `model.fit` to the [keras-team/keras repo](https://github.com/keras-team/keras/issues); I am unable to offer specific advice on this matter. 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.
<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Documentation Feature Request

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

No

### Source

binary

### Tensorflow Version

v2.9.0-18-gd8ce9f9c301 2.9.1

### Custom Code

No

### OS Platform and Distribution

Windows 11

### Mobile device

_No response_

### Python version

3.9.5

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

Hello Folks!

There is a very little idea about the type of hardware/software that can be used with Distributed training on TensorFlow. I've access to some supercomputer instances (pardon me for ignorance, this isn't the field where I usually work, so bit daunting for me), the problem is to train some very large model on it.

The documentation featuring distribution strategy: https://www.tensorflow.org/guide/distributed_training doesn't gives much idea about the devices that are supported on TensorFlow: CPU clusters/GPU clusters or TPU Clusters. I hope it supports multiple machines with no GPU/TPU access.

Second, it doesn't give much idea about configuration of machines for the use either - I suppose if they are connected on Ethernet or Kubernetes (or something like that), they should show up in `mirrored_strategy = tf.distribute.MirroredStrategy()` as a list? Right? So, basically no third party installation is needed in the cluster nodes or anything like that, one ethernet cable should suffice? (I guess).

Kindly update the docs with better clear answers to the above questions for newbies like me. Thank You.

### Standalone code to reproduce the issue

```shell
`mirrored_strategy = tf.distribute.MirroredStrategy()`
```


### Relevant log output

_No response_</details>