Системный профилировщик TorchNeuron на платформе AWS Trainium

Анализ производительности инференса моделей машинного обучения позволяет выявить узкие места с помощью системного профилирования на базе процессоров AWS Trainium и бэкенда TorchNeuron. В данном материале подробно рассматривается системный профировщик TorchNeuron для оценки времени выполнения операций.

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

Интеграция Neuron как бэкенда PyTorch

Бэкенд Neuron взаимодействует с экосистемой PyTorch посредством механизма расширения PrivateUse1. Такой подход дает возможность интеграции аппаратного ускорителя без выделенного ключа диспетчеризации вроде cuda. Специальный пакет torch_neuronx регистрирует Neuron в качестве нативного бэкенда для выполнения исходного кода.

Настройка инструментов профилирования

Для сбора данных о выполнении рабочей нагрузки на чипах Trainium применяются два основных компонента. Объект NeuronConfig задает параметры захвата потоков данных и выбор конкретных ядер NeuronCore, а NeuronProfiler управляет экспортом собранной информации.

Параметры конфигурации NeuronConfig

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

Инструмент визуализации Neuron Explorer

Для детального изучения полученных данных используется графический обозреватель Neuron Explorer, доступный через SSH-туннель или расширение для среды разработки. Он помогает визуально подтвердить числовые показатели, полученные из табличных отчетов, и проанализировать паузы между операциями.

Форматы файлов трассировки

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

Методика анализа операции RMSNorm

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

Первый вариант: реализация базовыми операциями PyTorch

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

Запуск сбора метрик и синхронизация

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

Анализ таблицы средних значений

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

Второй вариант: использование встроенной функции torch.rms_norm

Альтернативный подход заменяет набор разрозненных команд на встроенную функцию PyTorch. В этом случае компилятор Neuron оптимизирует вычисления, объединяя их в одно цельное ядро обработки.

Сравнение путей выполнения

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

Третий вариант: компиляция базовых операций через torch.compile

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

Особенности скомпилированного графа

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

Сводные результаты тестирования

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

Заключение

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