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

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

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

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

Neuron как бэкенд PyTorch

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

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

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

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

• NeuronConfig управляет тем, какие потоки данных и NeuronCore следует захватывать, а также местом записи результатов. Его передают в torch.profiler.profile как experimental_config.
• NeuronProfiler принимает NeuronConfig и предоставляет callback export_trace. Callback записывает трассировку фреймворка, trace.json, рядом с системными трассировками Neuron.

nc = NeuronConfig(
    modes=[ProfileMode.DEVICE, ProfileMode.RUNTIME],
    profile_output_dir=OUTPUT_DIR,
    capture_enabled_for_nc='0',
    correlate_device_time=True,
)
exporter = NeuronProfiler(nc)

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

Инструменты: Neuron Explorer

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

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

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

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

Профилируются три реализации одной и той же операции:

1. Наивная RMSNorm с отдельными операциями PyTorch: умножением, вычислением среднего, сложением, обратным квадратным корнем и умножением.
2. torch.rms_norm() — встроенная функция PyTorch, понижаемая до одного объединённого ядра Neuron.
3. Наивная RMSNorm с torch.compile, использующая те же операции, что и первая реализация.

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

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

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

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

device = torch.device('neuron')

x = torch.randn(4, 512, dtype=torch.bfloat16, device=device)
weight = torch.ones(512, dtype=torch.bfloat16, device=device)

def rms_norm_fn(t):
    x_sq = t * t                                    
    mean_sq = x_sq.mean(dim=-1, keepdim=True)       
    rms = torch.rsqrt(mean_sq + 1e-6)               
    return t * rms * weight                         

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

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

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

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

with torch.profiler.profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1],
    experimental_config=nc,
    record_shapes=False,
    on_trace_ready=exporter.export_trace,
    schedule=torch.profiler.schedule(wait=2, warmup=1, active=1),
) as prof:
    for step in range(4):
        with torch.no_grad():
            _ = rms_norm_fn(x)
        torch.neuron.synchronize()
        prof.step()

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

print(prof.key_averages().table(sort_by="self_privateuse1_time_total", row_limit=4))

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

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

Общее собственное время NEURON: 2.355 ms

Каждая строка напрямую соответствует функции. aten::mul выполняется для x * x, x * rms и финального умножения на weight. 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

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

def rms_norm_fn(t):
    return torch.rms_norm(t, [512], weight)

Объединённая RMSNorm отображается в таблице key averages как одна операция устройства.

Общее собственное время NEURON: 403.3 μs

Объединённая RMSNorm проходит по одному пути от фреймворка к оборудованию вместо отдельной диспетчеризации каждой арифметической операции.

Сравнение пути понижения

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

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

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

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

@torch.compile(backend='neuron', fullgraph=True)
def naive_compiled(t):
    x_sq = t * t
    mean_sq = x_sq.mean(dim=-1, keepdim=True)
    rms = torch.rsqrt(mean_sq + 1e-6)
    return t * rms * weight

Скомпилированная функция отображается в таблице key averages как один регион Torch-Compiled.

Общее собственное время NEURON: 375.5 μs

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

Таблица показывает одно исполнение на устройстве, как и в профиле 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 и остальным за рецензирование!

Наборы данных, упомянутые в этой статье 1

Другие материалы этого автора

Записывайте, обучайте и развёртывайте в одном месте с Strands Agents, LeRobot и корзинами хранения Hugging Face

• 

35

13 августа 2026 г.

От Hugging Face к Amazon SageMaker Studio одним щелчком

• 

21

7 июля 2026 г.