В посте, опубликованном на платформе Hugging Face, Amazon открывает серию материалов, посвященных анализу производительности вывода (инференса), то есть исполнения уже обученной модели. Эта первая часть объясняет, как определить, на что тратится время, с помощью инструментов профилирования, применяемых к чипу Trainium от AWS, созданному для обучения и инференса моделей глубокого обучения.
Руководство адресовано читателям без предварительного опыта в профилировании. Оно опирается на TorchNeuron — нативный open source бэкенд PyTorch для Trainium, который поддерживает стандартные рабочие процессы PyTorch, включая исполнение в eager-режиме и torch.compile, с минимальными изменениями кода. TorchNeuron также интегрируется с Neuron Kernel Interface (NKI), позволяющей создавать пользовательские операторы.
Настройка профилировщика
Профилирование основывается на двух объектах. NeuronConfig определяет потоки данных для захвата, NeuronCores для трассировки и выходную папку, а затем передается в torch.profiler.profile. NeuronProfiler предоставляет функцию экспорта, которая записывает трассу фреймворка в формате trace.json вместе с системными трассами Neuron. Параметр correlate_device_time=True позволяет отображать время, проведенное на устройстве, в таблице key_averages(), где его можно сравнить с временем CPU.
Затем трассы можно открыть в Neuron Explorer — инструменте визуализации, работающем в браузере. Он позволяет отследить операцию от распределения во фреймворке до её выполнения на оборудовании, включая планирование в рантайме. Для NeuronCore, то есть вычислительного ядра чипа, файл NEFF содержит скомпилированный исполняемый файл, в то время как соответствующий файл NTFF содержит трассу аппаратных инструкций.
Три версии одной операции
В посте исследуется функция RMSNorm — простая операция, состоящая из последовательности арифметических вычислений и редукций. Сравниваются три реализации: версия, написанная операция за операцией на PyTorch, вызов встроенной функции torch.rms_norm() и версия, написанная операция за операцией, но скомпилированная с помощью torch.compile.
В наивной версии таблица показывает, что три вызова aten::mul занимают наибольшую долю времени, проведенного на устройстве. Общее время Neuron достигает 2 355 µs для шести вызовов операторов. По словам автора, встроенная функция torch.rms_norm(), которая преобразуется в единое скомбинированное вычислительное ядро, занимает всего одну операцию, aten::_fused_rms_norm, с временем Neuron 403,3 µs. Версия, скомпилированная с torch.compile, отображается как единая скомпилированная область длительностью 375,5 µs.
В посте уточняется, что в этих тестах оба подхода занимают примерно одну шестую от времени вычислений наивной версии. Слияние устраняет перерывы между аппаратными исполнениями и позволяет избежать материализации промежуточных значений, таких как x_sq и mean_sq, между операциями, отправленными отдельно.
Автор указывает, что метод можно применять к более крупным моделям: определить узкое место в таблице, проверить его в трассе, а затем определить, может ли скомбинированная операция или torch.compile снизить затраты. Вторая часть серии будет посвящена профилированию на уровне устройства с анализом файла NEFF.
