Первая часть разбора Profiling TorchNeuron посвящена системному профилированию PyTorch-нагрузок на AWS Trainium. Эксперимент с RMSNorm показывает, что структура графа вычислений может влиять на время на устройстве сильнее, чем отдельные математические операции: объединение нескольких вызовов в один kernel dispatch устраняет промежуточные данные и паузы между запусками. Для разработчиков это практический способ проверять эффект от torch.compile и встроенных операций на Trainium, а не оценивать оптимизацию только по коду модели.
Профиль связывает уровни программного стека.
TorchNeuron позиционируется как нативный backend PyTorch для Trainium. Как следует из анонса PyTorch, устройство регистрируется через механизм PrivateUse1: стандартные API, включая torch.nn, torch.distributed и torch.profiler, должны работать с устройством neuron. Это меняет характер диагностики: разработчик видит не отдельный лог компилятора или метрики чипа, а последовательность от Python-вызова и диспетчеризации PyTorch до событий Neuron Runtime и исполнения на NeuronCore.
Для такого анализа стандартный torch.profiler дополняют объектом NeuronConfig. В документации Neuron Explorer описаны четыре режима сбора: DEVICE записывает низкоуровневые инструкции NeuronCore, RUNTIME — системную трассу Neuron Runtime, CPU_UTIL — загрузку процессора хоста, HOST_MEMORY — использование его памяти. Совместный профиль позволяет отличить задержку на стороне Python или CPU от простоя ускорителя, проблем с памятью либо неудачного разбиения вычислений на kernels.
Корректность процедуры здесь критична. Перед измерением документация рекомендует выполнить не менее трёх прогревочных итераций. В профиле должны быть включены CPU и PrivateUse1, а перед экспортом трассы требуется torch.neuron.synchronize(). Иначе асинхронное устройство может не завершить работу к моменту выгрузки, и время Neuron окажется неполным или не попадёт в отчёт. В разборе также предлагается передавать correlate_device_time в NeuronConfig, чтобы время устройства отображалось в таблице агрегированных операций.
RMSNorm показывает цену разрозненных операций.
В тесте RMSNorm реализована тремя способами: как набор обычных операций PyTorch, через встроенную torch.rms_norm и как исходная наивная версия, обёрнутая в torch.compile с backend='neuron'. В одном измерении вариант из отдельных операций занял 2 355 мкс времени Neuron и потребовал шести вызовов операторов. Встроенная реализация заняла 403 мкс, а скомпилированная наивная версия — 376 мкс: в обоих случаях вычисление было сведено к одному запуску kernel.
Разница объясняется не изменением формулы RMSNorm. Последовательная реализация отдельно выполняет умножение, усреднение, сложение, обратный квадратный корень и финальное умножение. Между вызовами создаются промежуточные тензоры, включая квадрат входа и среднее квадратов, а в расписании устройства появляются промежутки. Fusion позволяет не материализовать эти данные и передать весь фрагмент одним исполняемым блоком. При этом приведённые микросекунды относятся к конкретному запуску RMSNorm, а не являются сравнением производительности Trainium с другими ускорителями или оценкой ускорения всей модели.
От трассы к изменению кода.
Практический цикл оптимизации начинается с поиска операций с заметным device time и проверки, сколько запусков они создают. Затем разработчик сопоставляет событие с исходной строкой, вызовами рантайма и зависимостями в Neuron Explorer. После замены операции на встроенную, применения torch.compile или переработки участка графа профиль снимают повторно в тех же условиях. Такой подход полезен и для слоёв внимания, и для блоков MoE: узкое место может находиться в kernel, передаче данных, CPU fallback или в лишней фрагментации вычислений.
Этот сценарий уже используется в прикладной работе, хотя его результаты нельзя переносить на все модели. В описании переноса Nova-3 компанией Deepgram и AWS сначала запускали модель в eager mode и проверяли операции, уходящие на CPU, а затем применяли torch.compile. Авторы сообщили, что в скомпилированном пути не обнаружили CPU fallback для операций модели. Последующую настройку они строили на профилях Neuron Explorer, готовых kernels из NKILib и, там где это было необходимо, на собственных kernels.
Удобство backend не отменяет ограничений.
Amazon в своём сообщении о TorchNeuron демонстрирует перенос скрипта с изменением настройки устройства и заявляет о поддержке Trainium в инструментах Hugging Face. Однако это не означает, что любой PyTorch-проект сразу получит оптимальный граф. Опубликованный разбор ограничений TorchNeuron указывает на требование статических форм и отмечает, что backend neuron компилирует каждый FX-граф в отдельный исполняемый файл Neuron. Динамические формы, неподдержанные операции и границы между графами остаются предметом проверки в профиле.
Для команд, использующих открытые модели, главный вывод состоит в том, что перенос на новый ускоритель нужно рассматривать как измеряемую инженерную задачу. Сначала следует подтвердить отсутствие fallback и собрать трассу после прогрева, затем найти фрагменты с избыточным числом dispatch и лишь после этого выбирать между встроенными операциями, torch.compile и собственными kernels. Системный профиль не гарантирует ускорения, но делает видимыми те места, где оно теряется.
