Руководство по флагам 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 (межчипового межсоединения), что создает больше возможностей для перекрывающегося выполнения. 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 флага следует использовать совместно для активации асинхронных операций сбора всех данных в версии 5e. 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		 Эти два флага следует использовать совместно для активации асинхронных операций reduce в версии 5e. 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		 Этот флаг предназначен для операций сбора данных с высокой задержкой (т.е. небольшого размера). Его включение запускает определенные оптимизации, которые могут сократить время выполнения таких операций. Обычно он используется в задачах вывода результатов. 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		 Этот флаг предназначен для операций сбора данных с высокой задержкой (т.е. небольшого размера). Его включение запускает определенные оптимизации, которые могут сократить время выполнения операций сокращения с высокой задержкой. Обычно он используется в задачах вывода результатов. 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		 Этот флаг предназначен для операций сбора данных с высокой задержкой (т.е. небольшого размера). Его включение запускает определенные оптимизации, которые могут сократить время выполнения для операций коллективной перестановки с высокой задержкой. Обычно он используется в задачах вывода результатов. 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		 Этот флаг предназначен для операций сбора данных с высокой задержкой (т.е. небольшого размера). Его включение запускает определенные оптимизации, которые могут сократить время выполнения для операций сбора данных со всеми участниками с высокой задержкой. Обычно он используется в задачах вывода результатов. 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_gpu_enable_analytical_sol_latency_estimator		 Этот флаг включает аналитический оценщик, который максимизирует перекрытие вычислительных и коммуникационных процессов на графических процессорах. xla_gpu_enable_analytical_sol_latency_estimator=true xla_gpu_enable_analytical_sol_latency_estimator=false true/false

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

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

Флаги TPU XLA

Флаг Тип Примечания
xla_tpu_enable_data_parallel_all_reduce_opt Логическое значение (истина/ложь) Оптимизация для увеличения возможностей пересечения потоков данных в сетях центров обработки данных (DCN) с использованием алгоритмов all-reduce для параллельного сегментирования данных.
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 Логическое значение (истина/ложь) Конвейеры обработки данных (в настоящее время мегамасштабные системы обработки данных) осуществляются в обратном порядке через циклы сканирования.

Флаги GPU XLA

Уровень оптимизации -O1 включает расширенные проходы компилятора для повышения производительности GPU, включая несколько категорий флагов, перечисленных ниже: конвейеризация параллельных вычислений ( xla_gpu_enable_pipelined_all_gather , xla_gpu_enable_pipelined_all_reduce , xla_gpu_enable_pipelined_reduce_scatter ), развертывание цикла while ( xla_gpu_enable_while_loop_double_buffering ), планирование с сокрытием задержки ( xla_gpu_enable_latency_hiding_scheduler ) и оценка задержки SOL на Hopper/Blackwell ( xla_gpu_enable_analytical_sol_latency_estimator ). Подробности см. в разделе «Уровни усилий GPU» .

Флаг Тип Примечания
xla_gpu_enable_latency_hiding_scheduler Логическое значение (истина/ложь) Этот флаг позволяет планировщикам, скрывающим задержки, эффективно перекрывать асинхронную связь с вычислениями. Значение по умолчанию — False.
xla_gpu_enable_analytical_sol_latency_estimator Логическое значение (истина/ложь) Позволяет принимать решения по планированию, специфичные для каждой платформы, что, в свою очередь, улучшает взаимодействие между вычислительными ресурсами и обменом данными. Значение по умолчанию — true.
xla_gpu_analytical_latency_estimator_options Структурированная строка Настраивает параметры для xla_gpu_enable_analytical_sol_latency_estimator . Настройте их, установив nic_speed_gbps=$NIC_SPEED,nccl_op_launch_us=$LAUNCH_OVERHEAD,chunk_prep_us=$CHUNK_PREP,rtt_us=$RTT,chunk_size_bytes=$CHUNK_SIZE,gpus_per_node=$GPUS_PER_NODE . Значение по умолчанию зависит от обнаруженной платформы.
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 следует установить их достаточно высокими, чтобы объединить как минимум AllGather / ReduceScatter с весом, равным весу слоя 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-catter.
xla_gpu_enable_pipelined_all_reduce Логическое значение (истина/ложь) Включить конвейерную обработку всех инструкций reduce.
xla_gpu_enable_pipelined_host_offloading Логическое значение (истина/ложь) Включить конвейерную обработку инструкций по разгрузке вычислений на хосте.
xla_gpu_enable_while_loop_double_buffering Логическое значение (истина/ложь) Включить двойную буферизацию для цикла while.
xla_gpu_enable_all_gather_combine_by_dim Логическое значение (истина/ложь) Объединять все операции сбора данных с одинаковым измерением сбора или независимо от их измерения.
xla_gpu_enable_reduce_scatter_combine_by_dim Логическое значение (истина/ложь) Объединяйте операции уменьшения рассеяния с одинаковым измерением или независимо от его измерения.