tf.distribute.experimental.CommunicationOptions

Параметры для межпроцессорной коммуникации, например, All-reduce.

tf.distribute.experimental.CommunicationOptions(
    bytes_per_pack=0,
    timeout_seconds=None,
    implementation=tf.distribute.experimental.CollectiveCommunication.AUTO
)

Это можно передать в методы, такие как tf.distribute.get_replica_context().all_reduce() для оптимизации производительности коллективных операций. Обратите внимание, что это только подсказки, которые могут или не могут изменить фактическое поведение. Некоторые параметры применяются только к определённой стратегии и игнорируются другими.

Одной из распространённых оптимизаций является разделение операции all-reduce для градиентов на несколько пакетов, чтобы обновления весов могли перекрываться с all-reduce градиентов.

Примеры:

options = tf.distribute.experimental.CommunicationOptions(
    bytes_per_pack=50 * 1024 * 1024,
    timeout_seconds=120.0,
    implementation=tf.distribute.experimental.CommunicationImplementation.NCCL
)
grads = tf.distribute.get_replica_context().all_reduce(
    'sum', grads, options=options)
optimizer.apply_gradients(zip(grads, vars),
    experimental_aggregate_gradients=False)