Руководство по флагам XLA

Это руководство предлагает тщательно подобранный набор ключевых флагов XLA, помогающих пользователям эффективно ориентироваться и использовать возможности XLA. В следующих разделах подробно описаны флаги, которые могут существенно влиять на производительность выполнения и использование памяти. Если после включения флага возникнут какие-либо проблемы, например, сбои, рекомендуется вернуться к настройкам по умолчанию и создать запрос на GitHub.

Флаги корректности

Флаг Описание Значения по умолчанию Предлагаемые значения Ценности кандидата
xla_mosaic_on_device_checks Этот флаг включает проверки на устройстве для кодогенерации Mosaic. В настоящее время поддерживаются проверки на границах памяти, то есть при обращении к памяти, выходящей за границы, компиляция/выполнение обнаружит это. xla_mosaic_on_device_checks=bounds xla_mosaic_on_device_checks=bounds xla_mosaic_on_device_checks=bounds

Флаги производительности

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

Флаг Описание Значения по умолчанию Предлагаемые значения Ценности кандидата
Конвейеризация
1. xla_should_allow_loop_variant_parameter_in_chain
2. xla_should_add_loop_invariant_op_in_chain
3. xla_tpu_enable_ici_ag_pipelining		 Эти 3 флага следует использовать совместно, чтобы включить коллективную конвейеризацию операций сбора данных ICI (Interchip-Interconnect), что создает больше возможностей для перекрывающегося выполнения. 1. xla_should_allow_loop_variant_parameter_in_chain=kDisabled
2. xla_should_add_loop_invariant_op_in_chain=kDisabled
3. xla_tpu_enable_ici_ag_pipelining=false		 1. xla_should_allow_loop_variant_parameter_in_chain=kEnabled
2. xla_should_add_loop_invariant_op_in_chain=kEnabled
3. xla_tpu_enable_ici_ag_pipelining=true		 1. xla_should_allow_loop_variant_parameter_in_chain=kDisabled/kEnabled/kAuto
2. xla_should_add_loop_invariant_op_in_chain=kDisabled/kEnabled/kAuto
3. xla_tpu_enable_ici_ag_pipelining=true/false
v5e/Асинхронный
xla_enable_async_all_gather
xla_tpu_enable_async_collective_fusion
xla_tpu_enable_async_collective_fusion_fuse_all_gather		 Эти 3 флага следует использовать совместно для активации асинхронных операций сбора всех данных на v5e. xla_enable_async_all_gather=kAuto
xla_tpu_enable_async_collective_fusion=true
xla_tpu_enable_async_collective_fusion_fuse_all_gather=true		 xla_enable_async_all_gather=kAuto
xla_tpu_enable_async_collective_fusion=true
xla_tpu_enable_async_collective_fusion_fuse_all_gather=true		 xla_enable_async_all_gather=kDisabled/kEnabled/kAuto
xla_tpu_enable_async_collective_fusion=true/false
xla_tpu_enable_async_collective_fusion_fuse_all_gather=true/false
v5e/Асинхронный
xla_tpu_enable_async_collective_fusion
xla_tpu_enable_async_collective_fusion_fuse_all_reduce		 Эти 2 флага следует использовать совместно для активации асинхронных операций all-reduce на v5e. xla_tpu_enable_async_collective_fusion=true
xla_tpu_enable_async_collective_fusion_fuse_all_reduce=false		 xla_tpu_enable_async_collective_fusion=true
xla_tpu_enable_async_collective_fusion_fuse_all_reduce=true		 xla_tpu_enable_async_collective_fusion=true/false
xla_tpu_enable_async_collective_fusion_fuse_all_reduce=true/false
Асинхронный
xla_tpu_enable_async_all_to_all		 Этот флаг включает асинхронную связь «все со всеми». xla_tpu_enable_async_all_to_all=false xla_tpu_enable_async_all_to_all=true xla_tpu_enable_async_all_to_all=true/false
Ограниченная задержка
xla_all_gather_latency_bound_threshold_in_bytes		 Этот флаг предназначен для операций all-gather с ограниченной задержкой (т.е. для небольших по размеру) (all-gather). Включение этого флага запускает определённые оптимизации, которые могут сократить время выполнения all-gather с ограниченной задержкой. Обычно он используется в рабочих нагрузках вывода. xla_all_gather_latency_bound_threshold_in_bytes=-1
(который не включен)		 4~16Mb(ie 4~16 * 1024 * 1024) [0, 9223372036854775807]
Ограниченная задержка
xla_all_reduce_latency_bound_threshold_in_bytes		 Этот флаг предназначен для операций all-gather с ограниченной задержкой (т.е. для небольших по размеру). Включение этого флага запускает определённые оптимизации, которые могут сократить время выполнения операций all-reduce с ограниченной задержкой. Обычно он используется в рабочих нагрузках вывода. xla_all_reduce_latency_bound_threshold_in_bytes=-1
(который не включен)		 4~16Mb(ie 4~16 * 1024 * 1024) [0, 9223372036854775807]
Ограниченная задержка
xla_collective_permute_latency_bound_threshold_in_bytes		 Этот флаг предназначен для операций all-gather с ограниченной задержкой (т.е. небольших по размеру). Включение этого флага запускает определённые оптимизации, которые могут сократить время выполнения операций collective-permutes с ограниченной задержкой. Обычно он используется в рабочих нагрузках вывода. xla_collective_permute_latency_bound_threshold_in_bytes=-1
(который не включен)		 4~16Mb(ie 4~16 * 1024 * 1024) [0, 9223372036854775807]
Ограниченная задержка
xla_all_to_all_latency_bound_threshold_in_bytes		 Этот флаг предназначен для операций all-gather с ограниченной задержкой (т.е. для небольших по размеру) (all-gather). Включение этого флага запускает определённые оптимизации, которые могут сократить время выполнения all-to-all операций с ограниченной задержкой. Обычно он используется в рабочих нагрузках вывода. xla_all_to_all_latency_bound_threshold_in_bytes=-1
(который не включен)		 4~16Mb(ie 4~16 * 1024 * 1024) [0, 9223372036854775807]
xla_enable_async_collective_permute Переписывает все операции коллективной перестановки в их асинхронные варианты. При установке значения auto XLA может автоматически включать асинхронную коллективную перестановку на основе других конфигураций или условий. xla_enable_async_collective_permute=kAuto xla_enable_async_collective_permute=kAuto xla_enable_async_collective_permute=kAuto/kEnabled/kDisabled

Центрированные на вычислениях
xla_tpu_enable_dot_strength_reduction | Этот флаг переписывает не требующие больших вычислительных затрат точки как операции умножения и сокращения. | xla_tpu_enable_dot_strength_reduction=true | xla_tpu_enable_dot_strength_reduction=true | xla_tpu_enable_dot_strength_reduction=true/false

xla_tpu_dot_dot_fusion | Этот флаг включает слияние «точка-точка», которое объединяет операцию «производитель-точка» с операцией «потребитель-точка». При этом выходные данные «производителя-точки» не отображаются в медленной/основной памяти, что приводит к уменьшению объёма памяти. | xla_tpu_dot_dot_fusion=true | xla_tpu_dot_dot_fusion=true | xla_tpu_dot_dot_fusion=true/false
xla_jf_enable_multi_output_fusion | Этот флаг включает слияния, которые объединяют несколько потребителей (т. е. полученное слияние будет иметь несколько выходов) | xla_jf_enable_multi_output_fusion=true | xla_jf_enable_multi_output_fusion=true | xla_jf_enable_multi_output_fusion=true/false
xla_tpu_scoped_vmem_limit_kib | Этот флаг задаёт объём оперативной памяти VMEM, доступной для локального использования на каждую операцию (в килобайтах). Остальная часть VMEM используется как буферное пространство. | xla_tpu_scoped_vmem_limit_kib=16384 | xla_tpu_scoped_vmem_limit_kib=16384 | xla_tpu_scoped_vmem_limit_kib=[4096, VMEM size of the architecture - 1024]
xla_tpu_async_copy_bandwidth_scaling_factor | Масштабирует эффективную пропускную способность для асинхронных копий. Используется при принятии решений о предварительной выборке и определении того, какие тензоры следует разместить в VMEM. | xla_tpu_async_copy_bandwidth_scaling_factor=1 | xla_tpu_async_copy_bandwidth_scaling_factor=1 | xla_tpu_async_copy_bandwidth_scaling_factor=(0, 1]
xla_msa_enable_cross_program_prefetch_freeing | Включает оптимизацию освобождения для буферов с предварительной выборкой между программами. | xla_msa_enable_cross_program_prefetch_freeing=enabled | xla_msa_enable_cross_program_prefetch_freeing=enabled | xla_msa_enable_cross_program_prefetch_freeing=enabled/disabled
xla_tpu_msa_inefficient_use_to_copy_ratio | Отношение количества используемых байтов к количеству копируемых байтов для заданного сайта распределения, ниже которого сайт считается неэффективным. Это значение используется при принятии решений о размещении VMEM. Значение 0 будет считать все сайты эффективными, а значение 1 потребует, чтобы количество используемых байтов на сайте было не менее количества байтов асинхронного копирования. | xla_tpu_msa_inefficient_use_to_copy_ratio=0.5 | xla_tpu_msa_inefficient_use_to_copy_ratio=0.5 | xla_tpu_msa_inefficient_use_to_copy_ratio=[0, 1]

Флаги памяти

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

Флаг Описание Значения по умолчанию Предлагаемые значения Ценности кандидата
Планировщик
xla_latency_hiding_scheduler_rerun		 Этот параметр регулирует поведение планировщика, скрывающего задержку. Он работает, постепенно уменьшая лимит памяти, выделенный для планирования, при каждом «перезапуске» процесса. xla_latency_hiding_scheduler_rerun=1 xla_latency_hiding_scheduler_rerun=5 0~10(it doesn't make much sense beyond 10 reruns)
Слияние
xla_tpu_rwb_fusion		 Этот флаг включает слияния типа «reduce+broadcast» и может уменьшить использование памяти. xla_tpu_rwb_fusion=true xla_tpu_rwb_fusion=false xla_tpu_rwb_fusion=true/false
Планировщик
xla_memory_scheduler		 Этот флаг определяет алгоритм, который будет использовать планировщик памяти для минимизации потребления памяти. Использование более сложного алгоритма может обеспечить меньшее потребление памяти, но за счёт увеличения времени компиляции. xla_memory_scheduler=kDefault xla_memory_scheduler=kBrkga xla_memory_scheduler=kDefault/kList/kDfs/kPostOrder/kBrkga
Планировщик
xla_tpu_enable_latency_hiding_scheduler		 Этот флаг включает планировщик, скрывающий задержку, что позволяет выполнять асинхронные коллективные операции вместо синхронных. Отключение этого флага снижает потребление памяти за счёт потери прироста производительности, обеспечиваемого этими асинхронными операциями. xla_tpu_enable_latency_hiding_scheduler=true xla_tpu_enable_latency_hiding_scheduler=false xla_tpu_enable_latency_hiding_scheduler=true/false
СПМД
xla_jf_spmd_threshold_for_windowed_einsum_mib		 Этот флаг устанавливает нижний порог минимального размера точки для запуска коллективного выполнения. Установка большего значения позволит сэкономить память за счёт потери возможности выполнения коллективного выполнения. xla_jf_spmd_threshold_for_windowed_einsum_mib=-1 10Mb~1Gb (ie 10*1024*1024 ~ 1024*1024*1024) [0, 9223372036854775807]

Другие часто используемые флаги

Флаг Тип Примечания
xla_dump_to Строка (путь к файлу) Папка, в которой будут размещены файлы HLO предварительной оптимизации и другие артефакты (см. Инструменты XLA ).

Флаги TPU XLA

Флаг Тип Примечания
xla_tpu_enable_data_parallel_all_reduce_opt Булевое значение (истина/ложь) Оптимизация для увеличения возможностей перекрытия для DCN (сетей центров обработки данных) сокращает использование параллельного шардинга данных.
xla_tpu_data_parallel_opt_different_sized_ops Булевое значение (истина/ложь) Обеспечивает конвейеризацию параллельных операций с данными в нескольких итерациях, даже если размер их выходных данных не соответствует объёму, который можно сохранить в стековых переменных. Может увеличить нагрузку на память.
xla_tpu_spmd_rng_bit_generator_unsafe Булевое значение (истина/ложь) Следует ли запускать RngBitGenerator HLO разделенным способом, что небезопасно, если ожидаются детерминированные результаты с разным шардингом на разных частях вычислений.
xla_tpu_megacore_fusion_allow_ags Булевое значение (истина/ложь) Позволяет объединять все-сборы со свертками/всеми-свертками.
xla_tpu_enable_ag_backward_pipelining Булевое значение (истина/ложь) Конвейеры всех сборов (в настоящее время мегамасштабных всех сборов) в обратном направлении через циклы сканирования.

Флаги XLA графического процессора

Флаг Тип Примечания
xla_gpu_enable_latency_hiding_scheduler Булевое значение (истина/ложь) Этот флаг позволяет планировщикам, скрывающим задержку, эффективно накладывать асинхронную коммуникацию на вычисления. Значение по умолчанию — False.
xla_gpu_enable_triton_gemm Булевое значение (истина/ложь) Используйте матричное умножение на основе Triton.
xla_gpu_enable_command_buffer Список CommandBufferCmdType Какой тип
команды должны быть записаны в буферы команд.
xla_gpu_all_reduce_combine_threshold_bytes Целое число (байты) Эти флаги определяют, когда следует объединять несколько небольших AllGather / ReduceScatter / AllReduce в один большой AllGather / ReduceScatter / AllReduce, чтобы сократить время, затрачиваемое на взаимодействие между устройствами. Например, для пороговых значений AllGather / ReduceScatter в рабочей нагрузке на базе Transformer рекомендуется установить их достаточно высокими, чтобы объединить вес хотя бы одного слоя Transformer. По умолчанию значение combine_threshold_bytes равно 256.
xla_gpu_all_gather_combine_threshold_bytes Целое число (байты) См. xla_gpu_all_reduce_combine_threshold_bytes выше.
xla_gpu_reduce_scatter_combine_threshold_bytes Целое число (байты) См. xla_gpu_all_reduce_combine_threshold_bytes выше.
xla_gpu_enable_pipelined_all_gather Булевое значение (истина/ложь) Включить конвейеризацию инструкций по сбору всех данных.
xla_gpu_enable_pipelined_reduce_scatter Булевое значение (истина/ложь) Включить конвейеризацию инструкций reduce-scatter.
xla_gpu_enable_pipelined_all_reduce Булевое значение (истина/ложь) Включить конвейеризацию всех инструкций по сокращению.
xla_gpu_enable_while_loop_double_buffering Булевое значение (истина/ложь) Включить двойную буферизацию для цикла while.
xla_gpu_enable_all_gather_combine_by_dim Булевое значение (истина/ложь) Объедините операции по сбору всех предметов с одинаковым измерением сбора или независимо от их измерения.
xla_gpu_enable_reduce_scatter_combine_by_dim Булевое значение (истина/ложь) Объедините операции уменьшения-разброса с одинаковым измерением или независимо от их измерения.