Repository: tensorflow
Issue: Inconsistency in SyncBatchNormalization/BatchNormalization(synchronized=True) Results during Distributed Training on CPUs
Labels: stat:awaiting tensorflower,type:bug,comp:dist-strat,TF 2.11

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

Yes

### Source

source

### Tensorflow Version

tf 2.11.0

### Custom Code

Yes

### OS Platform and Distribution

Linux Ubuntu 20.04.5 LTS

### Mobile device

_No response_

### Python version

3.8.10

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

11.7

### GPU model and memory

_No response_

### Current Behaviour?

I am reporting an issue encountered during the distributed training of a model with `SyncBatchNormalization` layers on CPUs. The issue presents as a relatively large inconsistency between the results produced when training the same model on a single device versus multiple devices.

We've prepared a reproduction script demonstrating this issue. We set the weights of the `SyncBatchNormalization`/`BatchNormalization` layer to reproduce this bug. After conducting a one-step training, the predictions from the identical model trained on a single CPU and two CPUs exhibited relatively large differences (measured by L-infinity distance, Linf = 0.00074467063). As a comparison, we build another model with exactly the same architecture but removing the `SyncBatchNormalization` layer, and the prediction difference from this model trained on a single CPU and two CPUs is Linf = 1.0131771e-09.

Another comparison experiment is on GPUs. We trained the same model containing `SyncBatchNormalization` layer on a single GPU and two GPUs. The Linf of model’s prediction result is 2.5331974e-07. It’s much smaller than 1CPU vs 2CPU (Linf = 0.00074467063). This is weird because I would expect the CPU executions to be more deterministic than GPUs, and I would expect larger inconsistencies in the GPU runs if the inconsistencies are caused by variance.

To ensure a controlled environment, we used the same training data for one-step training on both setups, and then evaluated the model. To guarantee that the distributed trainings are expected to produce the same results given different number of devices, we use `MirroredStrategy`, keep the global batch size the same, and use `tf.keras.losses.Reduction.AUTO` reduction. To make the difference more apparent given a limited amount of training data, we deliberately chose a relatively high learning rate (lr=10).

It's noteworthy that `SyncBatchNormalization` has been deprecated as of TensorFlow version 2.12. Therefore, we also conducted the same experiment using TensorFlow's nightly build (2.13.0-dev20230428), replacing `SyncBatchNormalization()` with `BatchNormalization(synchronized=True)`. This experiment still manifested the same inconsistency. The prediction from the model containing batchnorm layer trained on a single CPU and two CPUs exhibits a substantial difference of Linf = 0.00010111928, while for the model removing batchnorm layer it’s Linf = 6.212488e-36.


### Standalone code to reproduce the issue

```shell
For TF 2.11.0 CPU experiments, the model contains `SyncBatchNormalization` layer:
https://colab.research.google.com/drive/11EseXxq_uHweY7omt7JCr1mc-42Kf1H5?usp=sharing

For TF 2.11.0 GPU experiments, the model contains `SyncBatchNormalization` layer (fix seed to generate same inputs):
https://colab.research.google.com/drive/1m_gCWzb_OKVTYAMMs8YbdKUYM3aINNPV?usp=sharing 

For TF nightly 2.13.0-dev20230428, the model contains `BatchNormalization(synchronized=True)` layer:
https://colab.research.google.com/drive/1N-aXPcfckVb8fPDSmghMlzBOFEGq6RzM?usp=sharing
```


### Relevant log output

For TF 2.11.0:

```shell
contain syncbatchnorm:  True
1CPU vs 2CPU: 0.00074467063

contain syncbatchnorm:  False
1CPU vs 2CPU: 1.0131771e-09

contain syncbatchnorm:  True
1GPU vs 2GPU: 2.5331974e-07
```

For TF nightly 2.13.0-dev20230428:

```shell
contain syncbatchnorm:  True
1CPU vs 2CPU: 0.00010111928

contain syncbatchnorm:  False
1CPU vs 2CPU: 6.212488e-36
```

</details>