Профилирование TorchNeuron (Часть 1): Профилирование на системном уровне в AWS Trainium.

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

Эта статья написана с нуля и не предполагает предварительного опыта профилирования. В пошаговом руководстве используется AWS Trainium, специализированный чип искусственного интеллекта для рабочих нагрузок глубокого обучения (обучения и инференса), а также — его собственный бэкенд PyTorch с открытым исходным кодом. TorchNeuron поддерживает стандартные рабочие процессы PyTorch на Trainium, включая нетерпеливое выполнение (eager execution), torch.compile и стандартные API PyTorch с минимальными изменениями кода. Кроме того, TorchNeuron интегрируется с Neuron Kernel Interface (NKI) для поддержки пользовательских операторов и оптимизации производительности оборудования Trainium. К концу статьи вы будете знать:

• Как настроить torch.profiler с бэкендом Neuron и записать чистую трассировку.
• Как читать таблицу key_averages(), чтобы определить, какие операции потребляют время устройства Neuron, и сравнить их время Self CPU.
• Как проследить за одной операцией через Neuron Explorer: от распределения в фреймворке до планирования во время выполнения и аппаратного исполнения на NeuronCore.
• Как эти сигналы изменяются в трех реализациях одной и той же операции.

Neuron как бэкенд PyTorch

Прежде чем профилировать рабочую нагрузку, полезно понять, как интегрируется с PyTorch.

Бэкенд Neuron использует точку расширения PrivateUse1 в PyTorch. Этот механизм позволяет ускорителю интегрироваться в фреймворк без добавления специального ключа распределения, такого как cuda. Пакет torch_neuronx работает совместно с upstream PyTorch для регистрации Neuron в качестве нативного бэкенда для кода PyTorch. Подробную информацию об этом механизме расширения см. в Facilitating New Backend Integration by PrivateUse1.

Настройка профилировщика

Для профилирования рабочей нагрузки Neuron используются два объекта:

• NeuronConfig управляет тем, какие потоки данных Neuron и какие NeuronCores записывать, а также куда выводить результат. Вы передаете его в torch.profiler.profile в качестве аргумента experimental_config.
• NeuronProfiler принимает NeuronConfig и предоставляет обратный вызов export_trace. Этот обратный вызов записывает трассировку фреймворка, trace.json, вместе с системными трассировками Neuron.

Параметр modes выбирает потоки данных для записи, такие как инструкции устройства и трассировки среды выполнения. Параметр capture_enabled_for_nc выбирает ядра NeuronCore для трассировки. Установка correlate_device_time=True передает время устройства из трассировки в таблицу key_averages(), где вы можете сравнить время центрального процессора (CPU) и устройства Neuron. Полный справочник параметров и доступные значения ProfileMode см. в документации по режимам профилирования NeuronConfig.

Инструментарий: Neuron Explorer

Neuron Explorer — это средство просмотра на базе браузера для системной трассировки, временной шкалы устройства и детализации по отдельным движкам. Вы можете получить к нему доступ через SSH-туннель или расширение Visual Studio Code. В этой статье вы используете Explorer, чтобы подтвердить числовые данные из key_averages() и изучить, что происходит между моментами диспетчеризации.

После профилирования выходной каталог содержит системную трассировку и, при захвате данных устройства, соответствующие файлы .neff и .ntff. Формат исполняемых файлов Neuron (NEFF) содержит скомпилированный исполняемый файл Neuron, а соответствующий файл NTFF содержит его трассировку аппаратных инструкций. Файлы имеют общий числовой хэш, который определяет их как пару. Чтобы открыть профиль в Neuron Explorer, загрузите полный выходной каталог для системного или комбинированного профиля. Для профиля, содержащего только данные устройства, загрузите соответствующие файлы .neff и .ntff по отдельности. Инструкции по интерфейсу пользователя и командной строке см. в Step 5: Upload your profile.

Профилирование RMSNorm

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

Вы профилируете три реализации одной и той же операции:

• Наивная RMSNorm с использованием отдельных операций PyTorch: умножение, среднее значение, сложение, обратный квадратный корень и умножение.
• torch.rms_norm(), встроенная функция PyTorch, низкоуровнево преобразованная в одно слитное ядро Neuron.
• Наивная RMSNorm с torch.compile, использующая те же операции, что и первая реализация.

В скриптах используются одинаковые входные данные, конфигурация профилировщика и расписание, поэтому результаты можно напрямую сравнивать.

Все скрипты для воспроизведения этой работы доступны в каталоге примеров torch-neuronx.

Профиль 1: Наивные операции PyTorch

Начните с написания RMSNorm с использованием отдельных операций PyTorch. Каждая строка соответствует одной или нескольким операциям ATen и может создавать промежуточные тензоры.

Вызовите профилировщик с двумя аргументами, необходимыми для бэкенда Neuron:

• Включите ProfilerActivity.PrivateUse1 для сбора активности с устройства Neuron.
• Передайте NeuronConfig в качестве experimental_config для включения функций Kineto, специфичных для Neuron.

В расписании используются wait=2, warmup=1 и active=1 в четырех вызовах prof.step(). Это позволяет записать одну итерацию прогрева без учета накладных расходов на начальную настройку.

Вызовите torch.neuron.synchronize() внутри цикла профилировщика. Без этой синхронизации профилировщик может выполнить экспорт до завершения работы устройства, из-за чего столбец времени Neuron останется пустым.

После завершения профилирования изучите таблицу key_averages(). Отсортируйте по self_privateuse1_time_total, чтобы операции с наибольшими затратами ресурсов устройства оказались на первом месте. Показанная здесь таблица отфильтрована по операциям RMSNorm, для которых регистрируется время устройства Neuron. Полный вывод key_averages() также включает события на стороне CPU, такие как выделение памяти и передачи данных с хоста на устройство.

Таблица средних значений ключевых показателей показывает, что три вызова aten::mul занимают наибольшую долю времени устройства Neuron.

Таблица показывает, какие операции ATen выполнялись, как часто они запускались и сколько времени они провели на CPU и на устройстве Neuron. aten::mul имеет наибольшее общее время работы устройства, поскольку функция выполняет три умножения.

Общее время Self NEURON: 2.355 мс

Каждая строка напрямую соотносится с функцией. aten::mul выполняется для x * x, x * rms и окончательного умножения на вес. aten::mean выполняет редукцию по последнему измерению. aten::add добавляет значение эпсилон и инициирует копирование с хоста на устройство, так как скаляр находится на CPU. aten::rsqrt вычисляет обратный квадратный корень.

Neuron Explorer отслеживает операцию aten::mul от диспетчеризации в фреймворке через планирование во время выполнения до выполнения на аппаратных линиях Neuron.

Чтение трассировки в Explorer

В Explorer время течет слева направо, а жизненный цикл события — сверху вниз. Каждая операция ATen следует одному и тому же паттерну: создание на линии фреймворка, планирование на линии среды выполнения как nrta_execute_schedule и выполнение устройства на аппаратной линии как nc_exec_running.

Представление системной трассировки организует события по линиям. В режиме группировки по умолчанию метки имеют следующий формат: instance/PID/subsystem/identifier:

В режиме eager каждая операция выполняется независимо. Используйте средство просмотра цепочки зависимостей, чтобы проследить за aten::mul.Tensor от уровня фреймворка через планирование среды выполнения до выполнения на устройстве.

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

Профиль 2: torch.rms_norm

Замените отдельные операции встроенной функцией rms_norm в PyTorch. Neuron сводит операцию к одному слитному ядру:

Слитный RMSNorm отображается как одна операция устройства в таблице средних значений ключевых показателей.

Всего Self NEURON: 403.3 мкс

Слитный RMSNorm следует по единому пути от фреймворка к железу вместо отправки каждой арифметической операции по отдельности.

Сравнение пути трансляции

В наивной реализации каждая операция независимо пересекает границу отправки и следует по собственному пути от фреймворка к среде выполнения и железу. В слитной реализации один вызов Python охватывает все вычисления целиком. Трейс содержит одну отправку и одно аппаратное событие без каких-либо зазоров между составными операциями.

Профиль 3: Наивный PyTorch с torch.compile

В качестве финального примера скомпилируйте исходную наивную функцию RMSNorm с помощью torch.compile.

С помощью @torch.compile(backend="neuron") TorchDynamo захватывает байт-код Python в виде графа FX, а AOTAutograd транслирует его в операции ATen. Бэкенд Neuron получает весь граф целиком вместо отправки каждой операции по отдельности. Он объединяет операции, оптимизирует доступ к памяти и компилирует граф в единый артефакт NEFF. Артефакт выполняется как одна аппаратная операция.

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

Всего Self NEURON: 375.5 мкс

Скомпилированный граф следует по единому пути через планирование среды выполнения и аппаратное выполнение.

Таблица показывает одно выполнение на устройстве, как и в Профиле 2, но код идет по другому пути. Вместо вызова специально созданной слитной операции компилятор обнаруживает слияние из исходной функции Python. В этом запуске время Neuron сопоставимо с torch.rms_norm и составляет 375 мкс по сравнению с 403 мкс.

Сводка сравнения

И torch.rms_norm, и torch.compile объединяют RMSNorm в отправку одного ядра. В этих запусках любой из подходов занимает примерно одну шестую от времени выполнения наивной реализации на устройстве. Слияние устраняет зазоры между аппаратными выполнениями и позволяет избежать материализации промежуточных значений, таких как x_sq и mean_sq, между раздельно отправляемыми операциями.

Заключение

key_averages() показывает, какие операции потребляют время устройства, в то время как полный вывод также раскрывает настройку на стороне CPU и передачи данных. Explorer предоставляет соответствующий визуальный трейс. Вы можете проследить за одной операцией от отправки из фреймворка через планирование среды выполнения до аппаратного выполнения и увидеть зазоры, которые устраняет слияние.

Вы можете использовать тот же рабочий процесс для больших моделей: найти «узкое место» в таблице, проверить его в трейсе и затем определить, может ли слитная операция или torch.compile снизить его стоимость.

В Части 2 мы перейдем к профилированию на уровне устройства. Вы откроете NEFF, изучите детализацию TensorEngine, VectorEngine и прямого доступа к памяти (DMA), а также разместите ядро на модели roofline, чтобы определить, ограничено ли оно производительностью вычислений или пропускной способностью памяти.

Особая благодарность Аритре Рою Гоштипати (Aritra Roy Gosthipaty) за вдохновение начать эту серию, а моим коллегам Пинаку Паниграхи (Pinak Panigrahi), Самиру Араужо (Samir Araujo), Ахилу Ажикодану (Akhil Azhikodan), Эше Лахотиа (Esha Lakhotia) и остальным — за рецензирование!