tf.distribute.experimental.PreemptionCheckpointHandler

Обработчик прерываний и ошибок для синхронного обучения.

tf.distribute.experimental.PreemptionCheckpointHandler(
    cluster_resolver,
    checkpoint_or_checkpoint_manager,
    checkpoint_dir=None,
    termination_config=None
)

Примечание: Этот API поддерживает использование только с tf.distribute.MultiWorkerMirroredStrategy и tf.distribute.TPUStrategy.

PreemptionCheckpointHandler координирует работу всех рабочих узлов, чтобы сохранить контрольную точку при получении сигнала о прерывании. Он также помогает точно распространять сообщения об ошибках приложения по кластеру. При создании объекта PreemptionCheckpointHandler он восстанавливает значения из последнего файла контрольной точки, если такой файл существует.

Сразу после инициализации объект начинает отслеживать сигналы о завершении для всех участников кластера. При получении сигнала в следующий раз, когда рабочий узел выполняет PreemptionCheckpointHandler.run, PreemptionCheckpointHandler синхронизирует работу всех рабочих узлов для сохранения контрольной точки. Затем, если через tf.distribute.experimental.TerminationConfig настроен exit_fn, он будет вызван. В противном случае процесс просто завершится, а платформа должна будет перезапустить его позже.

Примечание: Пользователям tf.distribute.MultiWorkerMirroredStrategy, которые решили настроить собственный exit_fn в tf.distribute.experimental.TerminationConfig, рекомендуется включить sys.exit(CODE_OR_MESSAGE) в exit_fn, чтобы после перезапуска все рабочие узлы могли правильно инициализировать службы связи. Пользователи tf.distribute.TPUStrategy, которые не хотят перезапускать кластер, но предпочитают перезапуск в рамках текущего процесса (то есть оставить координатор работающим и повторно выполнить шаги подключения к кластеру, инициализации системы TPU и создания объекта TPUStrategy), могут настроить exit_fn как пустую операцию.

Для пользователей tf.distribute.MultiWorkerMirroredStrategy основной API — PreemptionCheckpointHandler.run:

strategy = tf.distribute.MultiWorkerMirroredStrategy()

trained_epoch = tf.Variable(initial_value=tf.constant(0, dtype=tf.dtypes.int64), name='epoch')
step_in_epoch = tf.Variable(initial_value=tf.constant(0, dtype=tf.dtypes.int64), name='step_in_epoch')

with strategy.scope():
  dataset, model, optimizer = ...

  checkpoint = tf.train.Checkpoint(optimizer=optimizer,
                                   model=model,
                                   trained_epoch=trained_epoch,
                                   step_in_epoch=step_in_epoch)

  preemption_checkpoint_handler = tf.distribute.experimental.PreemptionCheckpointHandler(cluster_resolver, checkpoint, checkpoint_dir)

while trained_epoch.numpy() < NUM_EPOCH:

  while step_in_epoch.numpy() < STEPS_PER_EPOCH:

    # distributed_train_function contains a call to strategy.run.
    loss += preemption_checkpoint_handler.run(distributed_train_function, args=(next(iterator),))
    # For users of MultiWorkerMirroredStrategy, usually
    # STEPS_PER_TRAIN_FUNCTION = 1.
    step_in_epoch.assign_add(STEPS_PER_TRAIN_FUNCTION)
    ...

  epoch.assign_add(1)
  step_in_epoch.assign(0)

Для пользователей tf.distribute.TPUStrategy основными API являются PreemptionCheckpointHandler.run и PreemptionCheckpointHandler.watch_preemption_scope:

strategy = tf.distribute.TPUStrategy(tpu_cluster_resolver)

# Rest of TPU init omitted, see documentation for TPUSTrategy.

with preemption_checkpoint_handler.watch_preemption_scope():
  while trained_epoch.numpy() < NUM_EPOCH:

    while step_in_epoch.numpy() < STEPS_PER_EPOCH:

      # distributed_train_function contains a call to strategy.run.
      loss += preemption_checkpoint_handler.run(distributed_train_function, args=(next(iterator),))

      # For users of TPUStrategy, usually STEPS_PER_TRAIN_FUNCTION >> 1 since
      # clustering multiple steps within a tf.function amortizes the overhead
      # of launching a multi-device function on TPU Pod.
      step_in_epoch.assign_add(STEPS_PER_TRAIN_FUNCTION)
      ...

    epoch.assign_add(1)
    step_in_epoch.assign(0)

Не все прерывания сопровождаются заблаговременным уведомлением, которое позволяет PreemptionCheckpointHandler обработать их, например прерывания, вызванные аппаратным сбоем. Если пользователь самостоятельно сохраняет контрольные точки для таких случаев за пределами PreemptionCheckpointHandler и использует tf.train.CheckpointManager, передайте его как аргумент checkpoint_or_checkpoint_manager в PreemptionCheckpointHandler. Если пользователь не использует tf.train.CheckpointManager, а работает напрямую с tf.train.Checkpoint, рекомендуется сохранять контрольные точки в каталоге, переданном как аргумент checkpoint_dir. Таким образом, при запуске программы PreemptionCheckpointHandler сможет восстановить последнюю контрольную точку из каталога, независимо от того, была ли она сохранена самим пользователем или PreemptionCheckpointHandler до прерывания.

Примечание о платформе:

PreemptionCheckpointHandler может обрабатывать только такие завершения, о которых сообщается заранее. Сейчас API распознает сигналы о завершении для CPU, GPU и TPU в Google Borg, а также для CPU и GPU на Google Cloud Platform. В таких случаях PreemptionCheckpointHandler автоматически выберет подходящий механизм обнаружения уведомлений о прерывании или обслуживании. Пользователи других платформ могут настроить мониторинг обнаружения с помощью tf.distribute.experimental.TerminationConfig. Здесь также можно настроить поведение при выходе и длительность периода ожидания.

Методы

run

Просмотреть исходный код

run(
    distributed_train_function, *args, **kwargs
)

Выполняет функцию обучения с обработкой ошибок и прерываний.

Эта функция обрабатывает сигнал о прерывании от любого узла кластера: сохраняет ход обучения и корректно завершает работу. Она также распространяет на все рабочие узлы сообщения об ошибках программы, возникших во время выполнения distributed_train_function, чтобы они могли вызвать ту же ошибку.

Аргумент distributed_train_function должен быть распределенной функцией обучения (то есть содержать вызов tf.distribute.Strategy.run). Пользователям tf.distribute.MultiWorkerMirroredStrategy рекомендуется передавать в PreemptionCheckpointHandler.run одношаговый distributed_train_function, чтобы контрольная точка успела сохраниться в случае отправки сигнала о прерывании или уведомления о техническом обслуживании.

Помимо обработки прерываний и ошибок, PreemptionCheckpointHandler.run(distributed_train_function, *args, **kwargs) работает так же, как distributed_train_function(*args, **kwargs), и возвращает тот же результат. distributed_train_function может возвращать результат или не возвращать его. Ниже приведен сокращенный пример:

@tf.function
def distributed_train_step(iterator):
  # A distributed single-step training function.

  def step_fn(inputs):
    # A per-replica single-step training function.
    x, y = inputs
    ...
    return loss

  per_replica_losses = strategy.run(step_fn, args=(next(iterator),))
  return strategy.reduce(
      tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)

for epoch in range(preemption_handler.total_run_calls // STEPS_PER_EPOCH,
                   EPOCHS_TO_RUN):
  iterator = iter(multi_worker_dataset)
  total_loss = 0.0
  num_batches = 0

  for step in range(preemption_handler.total_run_calls % STEPS_PER_EPOCH,
                    STEPS_PER_EPOCH):
    total_loss += preemption_handler.run(distributed_train_step)
    num_batches += 1

  train_loss = total_loss / num_batches
  print('Epoch: %d, train_loss: %f.' %(epoch.numpy(), train_loss))

  train_accuracy.reset_states()

save_checkpoint_if_preempted

Просмотреть исходный код

save_checkpoint_if_preempted(
    *args, **kwargs
)

Сохраняет контрольную точку, если поступил сигнал о прерывании.

Это альтернативный API для PreemptionCheckpointHandler.run и PreemptionCheckpointHandler.watch_preemption_scope. Этот метод работает как с tf.distribute.MultiWorkerMirroredStrategy, так и с tf.distribute.TPUStrategy. Однако при использовании с TPUStrategy этот метод добавит точку синхронизации между рабочими узлами и координатором, что может повлиять на производительность. Если это вызывает опасения, вместо него используйте сочетание PreemptionCheckpointHandler.watch_preemption_scope и PreemptionCheckpointHandler.run.

strategy = tf.distribute.TPUStrategy(tpu_cluster_resolver)
# initialization omitted

with strategy.scope():
  # Save in the checkpoint.
  trained_step = tf.Variable(initial_value=tf.constant(0, dtype=tf.dtypes.int64), name='trained_step', aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA)

  checkpoint_manager = tf.train.CheckpointManager(checkpoint, directory, max_to_keep=1)
  preemption_handler = tf.distribute.experimental.PreemptionCheckpointHandler(cluster_resolver, checkpoint_manager)

while trained_step.numpy() < NUM_STEPS:
  # Train STEPS_IN_FUNCTION steps at once.
  train_multi_step_function()
  trained_step.assign_add(STEPS_IN_FUNCTION)
  preemption_handler.save_checkpoint_if_preempted()

watch_preemption_scope

Просмотреть исходный код

@tf_contextlib.contextmanager
watch_preemption_scope()

Синхронизирует ошибки и, возможно, сохраняет контрольную точку для использования с TPUStrategy.

Примечание: При использовании с tf.distribute.MultiWorkerMirroredStrategy этот API не требуется.

Пример использования:

with preemption_checkpoint_handler.watch_preemption_scope():
  while trained_step.numpy() < NUM_STEPS:

    # distributed_train_function contains a call to strategy.run.
    loss += preemption_checkpoint_handler.run(distributed_train_function, args=(next(iterator),))
    trained_step.assign_add(STEPS_PER_TRAIN_FUNCTION)

В этом сценарии PreemptionCheckpointHandler.run отметит получение сигнала о прерывании, а watch_preemption_scope обработает этот сигнал, сохранив контрольную точку, а затем либо завершит работу для перезапуска, либо выполнит переданный пользователем exit_fn из tf.distribute.experimental.TerminationConfig. Если во время выполнения операций и функций в области действия блока сигнал о прерывании не поступит, watch_preemption_scope обеспечит завершение всех асинхронных операций и функций при выходе, а также вызовет исключения, если асинхронное выполнение завершится ошибкой.