Модель затрат LHS

tldr;

На этой странице описывается внутренняя структура модели стоимости, используемой планировщиком скрытия задержек. Если вы заинтересованы в настройке модели, перейдите сразу к разделу «Настройка» .

Планировщик сокрытия задержек (LHS) — это проход компилятора, который планирует HLO DAG таким образом, чтобы минимизировать время ожидания.

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

Таблицы производительности – коллективы ICI

Таблица производительности состоит из двух основных компонентов: коллектора и интерполятора.

Коллекционер

Сборщик — это инструмент C++, отвечающий за генерацию таблиц производительности для коллективных операций. Он измеряет производительность отдельных операций HLO (например, all-gather или all-reduce ) в статически определённом пространстве параметров.

Как это работает

Инструмент выполняет сканирование по диапазону коллективных операций, размеров и схем передачи для заданного кластера. Он использует существующую инфраструктуру многохостового исполнителя HLO и данные ExecutionProfile для запуска сгенерированного HLO и сбора метрик производительности.

Параметры сбора данных

Таблицы задержек собираются для перекрестного произведения следующих параметров:

  • Тип коллектива :
    • all-reduce
    • all-gather
    • reduce-scatter
  • Размер передачи :
    • Логарифмическая шкала от 1024 Б до 2 ГБ (например, 1024 Б, 2048 Б, 4096 Б, ...)
  • Схема перевода :
    • rail-aligned
    • non-rail-aligned

Эта проверка выполняется для внутриузловых кластеров с 2, 4 и 8 устройствами .

Выход

Результатом сбора данных является таблица задержек в формате .pbtxt (примерно 116 КБ на платформу).

Интерполятор

Интерполятор — это компонент компилятора, который использует сгенерированные таблицы производительности для предоставления оценок времени выполнения во время компиляции.

Внутренняя структура данных

При инициализации интерполятор преобразует таблицу производительности в карту. Эта карта использует кортеж (collective_type, transfer_scheme) в качестве ключа .

Значение, связанное с каждым ключом, представляет собой двумерную евклидову плоскость. Эта плоскость индексирует пропускную способность сети (измеренную сборщиком) по двум осям:

  1. Размер передачи.
  2. Количество задействованных устройств.

Поиск и интерполяция

Когда компилятор встречает коллективную операцию, интерполятор выполняет следующие шаги:

  1. Он определяет правильную 2D-плоскость пропускной способности, используя операцию (collective_type, transfer_scheme) в качестве ключа карты.
  2. Затем он применяет средневзвешенный поиск (на основе евклидова расстояния) в пределах этой 2D-плоскости, используя операцию (transfer_size, num_devices) в качестве точки запроса.
  3. Результатом этого поиска является единственное уникальное значение пропускной способности сети .

Обоснование: производительность и экстраполяция

Система предназначена для хранения пропускной способности сети , а не чистой задержки. Такой подход значительно упрощает экстраполяцию производительности для объёмов передачи данных, явно не указанных в таблице.

Если таблицы задержек фиксируют насыщение пропускной способности сети при размере коллектива S , то пропускная способность T в этой точке считается максимальной. Для любого нового коллектива размером S' > S время выполнения можно оценить следующим образом:

\[\text{EstimatedTime}(S') = \frac{S'}{T_{\text{saturated} } }\]

Это позволяет модели оценивать производительность коллективов любого размера, даже тех, которые превышают максимальный размер в 2 ГиБ, измеренный Collector.

  • Недооценивать максимальную пропускную способность.
  • Следовательно, переоценивайте время выполнения для больших передач.

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

Таблицы производительности – GEMM

Подобно системе для коллективов, таблицы задержек GEMM поддерживаются двумя компонентами: сборщиком и интерполятором .

Коллекционер

Collector — это инструмент на C++, который вычисляет таблицы производительности для умножения матриц общего назначения (GEMM). Он измеряет производительность умножения матриц на уровне dot HLO.

Как это работает

Инструмент выполняет сканирование статического пространства измерений GEMM (пакетное, два несжимаемых и одно сжимаемое измерение) и типов данных.

  • Типы данных по умолчанию: LHS = bf16,f32 , RHS = bf16,f32 , OUT = bf16,f32 .
  • Инфраструктура: повторно использует профилировщик операций HLO.

Параметры сбора

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

  • партия: {1, 2, 4}
  • м (без подряда): {256, 512, ..., 4096}
  • n (недоговорной): {256, 512, ..., 4096}
  • k (контракт): {256, 512, ..., 4096}

Вывод и хранение

В результате полной развертки создается таблица задержек .pbtxt , готовая к использованию интерполятором.

Интерполятор

Интерполятор — это компонент компилятора, который использует сгенерированные таблицы для оценки производительности GEMM.

Обоснование: насыщение FLOPS

Собранные таблицы задержек позволяют интерполятору реконструировать FLOPS для каждой записи:

\[\text{FLOPS} = \frac{2 \times b \times m \times n \times k}{\text{runtime} }\]

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

Поиск и интерполяция

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

Аналитическая модель затрат — DCN

Модель коллективных затрат S-образной кривой

Модель S-образной кривой представляет собой полностью аналитическую сетевую модель линии крыши.

Обзор

Модель предназначена для оценки эффективности коллективных операций на основе набора фиксированных свойств сети.

Входные данные модели

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

  1. Фиксированные сетевые свойства (определяемые пользователем):

    • Коллективные накладные расходы на запуск
    • Скорость сетевой карты
    • RTT (время прохождения туда и обратно)

    По умолчанию XLA автоматически определяет платформу и использует значения для наиболее распространённых архитектур. Эти свойства настраиваются пользователем. Подробности см. в разделе «Настройка» .

  2. Коллективные вклады:

    • Коллективный тип (например, AllGather , ReduceScatter )
    • Размер передачи
    • Количество узлов, участвующих в коммуникации

Интеграция

Модель S-кривой интегрирована в XLA:GPU и используется в Hopper и Blackwell.

Аналитическая модель затрат – Fusions

Для других ядер мы используем модель стоимости производительности графического процессора для оценки корректного времени выполнения. Подробнее об этом можно узнать здесь .

Настройка

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

export NIC_SPEED_GBPS=... # NIC speed per GPU in Gigabytes
export GPUS_PER_NODE=... # Num of GPUs per cluster interconnected with fast network (e.g. NVLINK)
export XLA_FLAGS=--xla_gpu_analytical_latency_estimator_options="nic_speed_gbps=$NIC_SPEED_GBPS,gpus_per_node=$GPUS_PER_NODE"