Руководство пользователя по изоляции HLO

В этом документе объясняется, как установить и использовать API и CLI изоляции HLO. Инструменты изоляции HLO помогают разработчикам и исследователям изолировать, проверять и отлаживать несоответствия числовых значений и проблемы стабильности в скомпилированных модулях HLO.

Установка

Компоненты изоляции HLO можно использовать через стандартный механизм сборки OpenXLA/TensorFlow (Bazel) или через бинарное распространение.

Настройка источника и Bazel

При сборке компилятора и набора инструментов из исходного кода необходимо включить или использовать следующие библиотеки:

  • Целевой API: //third_party/tensorflow/compiler/xla/tools/hlo_isolation:hlo_isolation_api
  • Цель командной строки: //third_party/tensorflow/compiler/xla/tools/hlo_isolation:hlo_isolation_test

Для сборки автономного инструмента командной строки:

bazel build -c opt //third_party/tensorflow/compiler/xla/tools/hlo_isolation:hlo_isolation_test

Интерфейс командной строки (CLI)

Интерфейс командной строки hlo_isolation_test позволяет изолировать и запускать проверки на несоответствие числовых значений и стабильность скомпилированных модулей HLO непосредственно из терминала. Он сравнивает результаты выполнения в различных средах (например, TPU против дефузированного TPU / CPU / интерпретатора).

Справочник флагов

Интерфейс командной строки поддерживает следующие флаги:

  • --hlo_file : Путь к входному файлу .hlo или .pbtxt для загрузки. Может быть в текстовом или протокольном формате (обязательно).
  • --test_platform : Целевая платформа для запуска основного теста (например, cpu , gpu , tpu ). По умолчанию используется cpu .
  • --reference_platform : Эталонная платформа для сравнения с базовым уровнем (например, interpreter ). Если поле пустое, сравнение с эталонной платформой отключено.
  • --filter_by_name : Регулярное выражение для поиска имени модуля. Будут запущены только соответствующие модули. По умолчанию используется .* .
  • --skip_by_name : Регулярное выражение для поиска имени модуля. Соответствующие модули будут пропущены.
  • --filter_by_opcode : Регулярное выражение для сопоставления кодов операций инструкций. Будут запущены только модули, содержащие хотя бы один совпадающий код операции. По умолчанию используется .* .
  • --skip_by_opcode : Регулярное выражение для сопоставления кодов операций инструкций. Модули, содержащие любой совпадающий код операции, будут пропущены.
  • --abs_error_bound : Абсолютная граница ошибки, используемая для сравнения. По умолчанию 0.01 .
  • --rel_error_bound : Относительная граница ошибки, используемая для сравнения. По умолчанию 0.1 .
  • --run_hlo_passes : Логический флаг, определяющий, следует ли выполнять стандартные проходы HLO для подмодулей. По умолчанию — false .
  • --shard_index : Конкретный индекс сегмента для запуска (с нулевой отсчетом). По умолчанию -1 (отключено).
  • --num_shards : Общее количество шардов. По умолчанию — 1 .

Базовый вызов 

./hlo_isolation_test \
  --hlo_file=/path/to/failing_fusion.hlo \
  --test_platform=gpu \
  --reference_platform=interpreter

Дампы результатов и артефактов

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

Содержимое свалки

При обнаружении несоответствия числовых значений инструмент записывает следующие отладочные артефакты:

  1. Текст неудачно завершившегося подмодуля HLO ( failed-module-<module_name>.txt ).
  2. Ожидаемый выходной литерал ( failed-<module_name>-expected.txt ).
  3. Фактический выходной текст ( failed-<module_name>-actual.txt ).
  4. Сводка по несоответствующим элементам ( failed-<module_name>-mismatches.txt ).

Место сброса

  • Среда тестирования: Если запуск осуществляется через bazel test или среду, определяющую переменную среды TEST_UNDECLARED_OUTPUTS_DIR , результаты помещаются непосредственно в указанный каталог с точными именами файлов, перечисленных выше (например, failed-<module_name>-expected.txt ).
  • Стандартный/ручной запуск: При ручном выполнении через командную строку артефакты записываются во временный каталог операционной системы (например, /tmp ), при этом сохраняются те же самые единые соглашения об именовании файлов (например, /tmp/failed-<module_name>-expected.txt ).

Интеграция с C++ и API

Для разработчиков, создающих собственные проходы компилятора, тестовые стенды или автоматизированные конвейеры, API C++ предоставляет прямой способ интеграции тестов на изоляцию.

Использование API напрямую

Основной API предоставляет функциональные интерфейсы для запуска модулей и получения структурированных отчетов:

#include "third_party/tensorflow/compiler/xla/tools/hlo_isolation/hlo_isolation_api.h"

xla::hlo_isolation::PipelineIsolationOptions options;
options.module_options.abs_error_bound = 0.01;
options.module_options.rel_error_bound = 0.1;
// Filter specific opcodes programmatically
options.filter_by_opcode = "exponential";

absl::StatusOr<std::vector<xla::HloIsolationTestResult>> results =
    xla::hlo_isolation::RunIsolationPipeline(
        input_hlo_module,
        &my_test_runner,
        &my_reference_runner,
        options);

Использование тестового миксина

При написании наборов тестов GoogleTest на C++ можно наследовать от HloIsolationTestMixin для встроенной обработки утверждений. Базовый класс должен предоставлять как средство запуска тестов, так и средство запуска ссылок (например, через HloPjRtInterpreterReferenceMixin ):

#include "third_party/tensorflow/compiler/xla/tests/hlo_pjrt_interpreter_reference_mixin.h"
#include "third_party/tensorflow/compiler/xla/tools/hlo_isolation/hlo_isolation_test_base.h"

class MyCustomPassIsolationTest : public xla::hlo_isolation::HloIsolationTestMixin<
    xla::HloPjRtInterpreterReferenceMixin<xla::HloPjRtTestBase>> {};

TEST_F(MyCustomPassIsolationTest, ChecksMyFusionSanity) {
  RunAndVerifyIsolationTest(my_failing_module);
}

Сегментированное выполнение (K8s/Slurm)

Для больших модулей или ресурсоемких тестовых матриц можно разделить выполнение на кластеры с несколькими устройствами (например, Google Kubernetes Engine или Slurm) с помощью флагов шардинга. Каждый shard_index детерминированно запускает изолированное подмножество декомпозированных подмодулей. Это позволяет обеспечить воспроизводимую распределенную проверку и целенаправленное повторное выполнение неудачных разделов.

Пример: задание Kubernetes

Каждый тестовый сегмент выполняется как отдельный под Kubernetes с использованием completionMode: Indexed . Параметр JOB_COMPLETION_INDEX передается непосредственно в флаг --shard_index интерфейса командной строки.

apiVersion: batch/v1
kind: Job
metadata:
  name: hlo-isolation-job
spec:
  completions: 50
  parallelism: 50
  completionMode: Indexed
  template:
    spec:
      containers:
      - name: test-runner
        image: gcr.io/my-project/hlo-isolation-tools:latest
        command: ["/bin/sh", "-c"]
        args:
        - |
          ./hlo_isolation_test \
            --num_shards=50 \
            --shard_index=$JOB_COMPLETION_INDEX \
            --hlo_file=/data/path/to/hlo.hlo
        volumeMounts:
        - name: hlo-data-volume
          mountPath: /data
      volumes:
      - name: hlo-data-volume
        csi:
          driver: gcsfuse.csi.storage.gke.io
          volumeAttributes:
            bucketName: my-xla-debug-bucket

Ключевые возможности

  • Портативность: Отделяет внутренние тестовые обертки от автономного API, что упрощает локальную отладку несоответствий HLO.
  • Детализация: Детальная фильтрация кодов операций и имен улучшает цикл отладки при работе с большими дампами HLO.
  • Расширяемость: Пользовательские функции обратного вызова для выполнения запуска и внедрения данных ( make_fake_arguments_fn ) позволяют полностью настраивать сложные рабочие процессы проверки.