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)