Тюнинг флагов Megascale XLA

MegascaleXLA (MXLA) координирует обмен данными между срезами TPU по сетям центров обработки данных (DCN). В этом руководстве представлены рекомендации и лучшие практики по настройке и оптимизации флагов MegascaleXLA для пользователей, работающих с многосрезовыми рабочими нагрузками TPU. Для большинства случаев были установлены разумные значения по умолчанию.

Размеры коллективных буферов

For small collectives, it is generally more efficient to package small TPU send/receives and network send/receives into larger ones to reduce latency overhead. For larger collectives, it may be advantageous to use a greater number of smaller intermediate buffers. A number of tools exist which allow for altering this behavior. TPU DMA reads may be coalesced or split for some collectives by passing,

--megascale_target_dma_size=<value>

Эта функция попытается запланировать операции чтения DMA, размер которых будет максимально близок к целевому. По умолчанию используется значение 8 МБ. Изменять это значение рекомендуется только в том случае, если трассировка профилирования показывает значительную задержку между моментом инициирования DMA и началом первой сетевой передачи для конкретной группы.

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

--megascale_sparse_reduction_threshold=<value>

По умолчанию это значение равно 256 КБ. Если профилировщик показывает значительную задержку при операции редукции, может быть полезно увеличить это значение, чтобы использовать меньшее количество редьюсеров, что приведет к уменьшению числа больших сетевых передач.

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

--megascale_max_reduction_shard_size=<value>

Это разделит буфер на большее количество меньших фрагментов, чем требуется для коллективного выполнения. По умолчанию это значение равно 8 МБ. Должно быть так, что megascale_max_reduction_shard_size >> megascale_sparse_reduction_threshold поскольку эти значения оказывают противоположное воздействие.

Для передачи данных «один к одному» (например, при collective-permute ) передачу данных на хосте можно разделить на части, используя следующие методы:

--megascale_chunk_size=<value>

Если значение не равно нулю, то каждая передача данных «один к одному» на хосте делится на несколько фрагментов размером megascale_chunk_size байт (по умолчанию: 8 МБ) для организации конвейерной передачи данных через DMA и сеть. Все остальные параметры коллективного управления используют megascale_target_dma_size , megascale_sparse_reduction_threshold и megascale_max_reduction_shard_size для управления сегментированием и разбиением на фрагменты.

Управление памятью и передача данных без копирования.

Передача данных по сетям мегамасштаба спроектирована таким образом, чтобы не требовалось копирование: области памяти хоста закрепляются и предварительно отображаются для прямого доступа к TPU и обратно. Если предварительно отображенная память исчерпана во время выполнения, система переключается на отображение по требованию или динамическое копирование памяти, что серьезно влияет на пропускную способность.

Предварительно отображенная область памяти 

--megascale_grpc_premap_memory_bytes=17179869184  # Default: 16 GB (16LL << 30)

Управляет размером области закрепленной памяти хоста, выделяемой при инициализации для передач DMA.

Диагностические симптомы

  1. MapDmaBuffer в стационарном режиме : В окне просмотра трассировки XProf найдите MapDmaBuffer . Эти вызовы должны происходить только во время запуска. Если они продолжаются во время итераций обучения, происходит динамическое переназначение.
  2. Megascale: Memory Copy : Если в XProf в дорожке «Транспорт связи» отображаются события Megascale: Memory Copy , это означает, что буферы памяти копируются, а не передаются с помощью DMA с нулевым копированием.

Разрешение

Увеличьте значение --megascale_grpc_premap_memory_bytes (например, до 24 ГБ или 32 ГБ, в зависимости от доступной памяти хоста на вашем экземпляре виртуальной машины TPU) и перезапустите задание.