Глядя на пример, описанный здесь: https://github.com/tensorflow/docs/blob/r1.15/site/en/guide/distribute_strategy.ipynb, кажется, что в функциях train_step() и step_fn() оптимизатор минимизирует потери без какой-либо связи между репликами. Я ожидаю, что градиенты должны быть агрегированы, а затем применяются градиенты. На основании документации о experimental_run_v2(): https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy#experimental_run_v2 указано, что
Запускать fn для каждой реплики с заданными аргументами.
Выполняется ops, указанные fn для каждой реплики. Если аргументы arg или kwargs имеют значения «на реплику», например, полученные с помощью «распределенного набора данных», то при выполнении fn для конкретной реплики он будет выполняться с компонентом тех значений «на реплику», которые соответствуют эта реплика.
Итак, похоже, что step_fn() запускается индивидуально на каждой реплике.