Guía sobre las marcas de XLA

En esta guía, se ofrece una selección de marcas clave de XLA para ayudar a los usuarios a navegar y utilizar las capacidades de XLA de manera eficaz. En las siguientes secciones, se detallan las marcas que pueden afectar significativamente el rendimiento del tiempo de ejecución y el uso de la memoria. Si surge algún problema, como fallas, después de habilitar una marca, se recomienda volver a la configuración predeterminada y crear un problema de GitHub.

Alertas de rendimiento

Las siguientes marcas son fundamentales para mejorar el rendimiento del tiempo de ejecución. Experimentar con estos parámetros de configuración puede generar mejoras considerables en el rendimiento.

Marcar Descripción Valores predeterminados Valores sugeridos Valores de candidatos
Canalización
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		 Estas 3 marcas deben usarse en conjunto para habilitar la canalización colectiva de las operaciones de recopilación de todos de ICI(interconexión entre chips), lo que crea más oportunidades para la ejecución superpuesta. 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/Async
xla_enable_async_all_gather
xla_tpu_enable_async_collective_fusion
xla_tpu_enable_async_collective_fusion_fuse_all_gather		 Estas 3 marcas deben usarse en conjunto para activar las operaciones de recopilación de todos asíncronas en 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/Async
xla_tpu_enable_async_collective_fusion
xla_tpu_enable_async_collective_fusion_fuse_all_reduce		 Estas 2 marcas se deben usar en conjunto para activar las operaciones de reducción total asíncronas en 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
Async
xla_tpu_enable_async_all_to_all		 Esta marca habilita la comunicación asíncrona de todos con todos. 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
Vinculado a la latencia
xla_all_gather_latency_bound_threshold_in_bytes		 Esta marca está diseñada para operaciones de recopilación de todos (es decir, de tamaño pequeño) vinculadas a la latencia. Habilitar esta opción activa optimizaciones específicas que pueden reducir el tiempo de ejecución de las operaciones all-gather vinculadas a la latencia. Por lo general, se usa en cargas de trabajo de inferencia. xla_all_gather_latency_bound_threshold_in_bytes=-1
(que no está habilitado)		 4~16Mb(i.e. 4~16 * 1024 * 1024) [0, 9223372036854775807]
Vinculado a la latencia
xla_all_reduce_latency_bound_threshold_in_bytes		 Esta marca está diseñada para operaciones de recopilación de todos (es decir, de tamaño pequeño) vinculadas a la latencia. Habilitar este parámetro activa optimizaciones específicas que pueden reducir el tiempo de ejecución de las operaciones de reducción total vinculadas a la latencia. Por lo general, se usa en cargas de trabajo de inferencia. xla_all_reduce_latency_bound_threshold_in_bytes=-1
(que no está habilitado)		 4~16Mb(i.e. 4~16 * 1024 * 1024) [0, 9223372036854775807]
Vinculado a la latencia
xla_collective_permute_latency_bound_threshold_in_bytes		 Esta marca está diseñada para operaciones de recopilación de todos (es decir, de tamaño pequeño) vinculadas a la latencia. Habilitar esta opción activa optimizaciones específicas que pueden reducir el tiempo de ejecución de las permutaciones colectivas vinculadas a la latencia. Por lo general, se usa en cargas de trabajo de inferencia. xla_collective_permute_latency_bound_threshold_in_bytes=-1
(que no está habilitado)		 4~16Mb(i.e. 4~16 * 1024 * 1024) [0, 9223372036854775807]
Vinculado a la latencia
xla_all_to_all_latency_bound_threshold_in_bytes		 Esta marca está diseñada para operaciones de recopilación de todos (es decir, de tamaño pequeño) vinculadas a la latencia. Habilitar este parámetro activa optimizaciones específicas que pueden reducir el tiempo de ejecución para la operación de todos a todos vinculada a la latencia. Por lo general, se usa en cargas de trabajo de inferencia. xla_all_to_all_latency_bound_threshold_in_bytes=-1
(que no está habilitado)		 4~16Mb(i.e. 4~16 * 1024 * 1024) [0, 9223372036854775807]
xla_enable_async_collective_permute Vuelve a escribir todas las operaciones de collective-permute en sus variantes asíncronas. Cuando se configura en auto, XLA puede activar la comunicación colectiva asíncrona automáticamente en función de otras configuraciones o condiciones. xla_enable_async_collective_permute=kAuto xla_enable_async_collective_permute=kAuto xla_enable_async_collective_permute=kAuto/kEnabled/kDisabled

Marcas de memoria

Las marcas que se indican a continuación se proporcionan para abordar los problemas relacionados con la HBM. Estos solo se deben ajustar si encuentras errores de "memoria insuficiente" de HBM durante la compilación del modelo. En todos los demás casos, se recomiendan los valores predeterminados, ya que modificarlos podría afectar negativamente el rendimiento.

Marcar Descripción Valores predeterminados Valores sugeridos Valores de candidatos
Programador
xla_latency_hiding_scheduler_rerun		 Este parámetro de configuración ajusta el comportamiento del programador de ocultamiento de latencia. Funciona reduciendo de forma incremental el límite de memoria asignado para la programación con cada "nueva ejecución" del proceso. xla_latency_hiding_scheduler_rerun=1 xla_latency_hiding_scheduler_rerun=5 0~10(it doesn’t make much sense beyond 10 reruns)
Fusion
xla_tpu_rwb_fusion		 Esta marca habilita las fusiones de tipo reducción y transmisión, y puede disminuir el uso de memoria. xla_tpu_rwb_fusion=true xla_tpu_rwb_fusion=false xla_tpu_rwb_fusion=true/false
Programador
xla_memory_scheduler		 Esta marca especifica el algoritmo que usará el programador de memoria para minimizar el consumo de memoria. Usar un algoritmo más avanzado podría generar una programación que consuma menos memoria, pero a costa de un tiempo de compilación más largo. xla_memory_scheduler=kDefault xla_memory_scheduler=kBrkga xla_memory_scheduler=kDefault/kList/kDfs/kPostOrder/kBrkga
Programador
xla_tpu_enable_latency_hiding_scheduler		 Esta marca habilita el programador de ocultamiento de latencia, que nos permite realizar operaciones colectivas asíncronas en lugar de síncronas. Inhabilitarlo reduce el uso de memoria, pero se pierden las mejoras de rendimiento de estas operaciones asíncronas. xla_tpu_enable_latency_hiding_scheduler=true xla_tpu_enable_latency_hiding_scheduler=false xla_tpu_enable_latency_hiding_scheduler=true/false
SPMD
xla_jf_spmd_threshold_for_windowed_einsum_mib		 Esta marca establece el umbral inferior del tamaño mínimo del punto para activar la multiplicación de matrices colectiva. Si se establece en un valor más alto, se ahorraría memoria a costa de perder oportunidades para realizar la multiplicación de matrices colectiva. xla_jf_spmd_threshold_for_windowed_einsum_mib=-1 10Mb~1Gb (i.e. 10*1024*1024 ~ 1024*1024*1024) [0, 9223372036854775807]

Otras marcas de uso común

Marcar Tipo Notas
xla_dump_to Cadena (ruta de acceso al archivo) Es la carpeta en la que se colocarán los archivos HLO previos a la optimización y otros artefactos (consulta Herramientas de XLA).

Marcas de XLA de TPU

Marcar Tipo Notas
xla_tpu_enable_data_parallel_all_reduce_opt Booleano (verdadero o falso) Se optimizó para aumentar las oportunidades de superposición de las operaciones all-reduce de la DCN (red del centro de datos) que se usan para el sharding paralelo de datos.
xla_tpu_data_parallel_opt_different_sized_ops Booleano (verdadero o falso) Permite la canalización de operaciones paralelas de datos en varias iteraciones, incluso si los tamaños de salida no coinciden con lo que se puede guardar en su lugar en las variables apiladas. Puede aumentar la presión de la memoria.
xla_tpu_spmd_rng_bit_generator_unsafe Booleano (verdadero o falso) Indica si se debe ejecutar el HLO de RngBitGenerator de forma particionada, lo que no es seguro si se esperan resultados determinísticos con diferentes fragmentaciones en diferentes partes del cálculo.
xla_tpu_megacore_fusion_allow_ags Booleano (verdadero o falso) Permite fusionar todas las recopilaciones con convoluciones o todas las reducciones.
xla_tpu_enable_ag_backward_pipelining Booleano (verdadero o falso) Las canalizaciones recopilan todos los datos (actualmente, recopilan todos los datos a gran escala) hacia atrás a través de los bucles de análisis.

Marcas de XLA para GPU

Marcar Tipo Notas
xla_gpu_enable_latency_hiding_scheduler Booleano (verdadero o falso) Esta marca permite que los programadores de ocultamiento de latencia superpongan la comunicación asíncrona con el procesamiento de manera eficiente. El valor predeterminado es False.
xla_gpu_enable_triton_gemm Booleano (verdadero o falso) Usa la multiplicación de matrices basada en Triton.
xla_gpu_graph_level Marca (0-3) Es la marca heredada para establecer el nivel del grafo de la GPU. Usa xla_gpu_enable_command_buffer en casos de uso nuevos. 0 = Desactivado; 1 = Captura fusiones y memcpys; 2 = Captura gemms; 3 = Captura convoluciones.
xla_gpu_all_reduce_combine_threshold_bytes Número entero (bytes) Estos parámetros ajustan el momento en que se deben combinar varias operaciones AllGather / ReduceScatter / AllReduce pequeñas en una sola operación AllGather / ReduceScatter / AllReduce grande para reducir el tiempo dedicado a la comunicación entre dispositivos. Por ejemplo, para los umbrales de AllGather / ReduceScatter en una carga de trabajo basada en Transformer, considera ajustarlos lo suficientemente altos como para combinar al menos el peso de la capa de Transformer AllGather / ReduceScatter. De forma predeterminada, combine_threshold_bytes se establece en 256.
xla_gpu_all_gather_combine_threshold_bytes Número entero (bytes) Consulta xla_gpu_all_reduce_combine_threshold_bytes más arriba.
xla_gpu_reduce_scatter_combine_threshold_bytes Número entero (bytes) Consulta xla_gpu_all_reduce_combine_threshold_bytes más arriba.
xla_gpu_enable_pipelined_all_gather Booleano (verdadero o falso) Habilita la canalización de instrucciones de all-gather.
xla_gpu_enable_pipelined_reduce_scatter Booleano (verdadero o falso) Habilita la canalización de instrucciones de reducción y dispersión.
xla_gpu_enable_pipelined_all_reduce Booleano (verdadero o falso) Habilita la canalización de todas las instrucciones de reducción.
xla_gpu_enable_while_loop_double_buffering Booleano (verdadero o falso) Habilita el almacenamiento intermedio doble para el bucle while.
xla_gpu_enable_all_gather_combine_by_dim Booleano (verdadero o falso) Combina todas las operaciones de recopilación con la misma dimensión de recopilación o independientemente de su dimensión.
xla_gpu_enable_reduce_scatter_combine_by_dim Booleano (verdadero o falso) Combina operaciones de reducción y dispersión con la misma dimensión o independientemente de su dimensión.