Профілювання TorchNeuron (частина 1): Профілювання на рівні системи на AWS Trainium.

Багато фахівців із машинного навчання мають певне уявлення про те, як змусити моделі працювати швидко, але значно менше людей знають, як проаналізувати, чому модель працює повільно. У частині 1 цієї серії ви дізнаєтеся, як систематично аналізувати продуктивність інференсу та визначати, де саме витрачається час.

Ця публікація починається з нуля й не передбачає попереднього досвіду профілювання. У покроковому керівництві використовується AWS Trainium — спеціалізований AI-чип для навчання моделей глибокого навчання та виконання інференсу — і TorchNeuron, його нативний бекенд PyTorch із відкритим кодом. TorchNeuron підтримує стандартні робочі процеси PyTorch на Trainium, зокрема нетерпляче виконання, torch.compile і стандартні API PyTorch, із мінімальними змінами коду. Крім того, TorchNeuron інтегрується з Neuron Kernel Interface (NKI), щоб уможливити створення власних операторів і оптимізацію продуктивності для апаратного забезпечення Trainium. Наприкінці ви знатимете:

• Як налаштувати torch.profiler із бекендом Neuron і отримати чистий трейс.
• Як читати таблицю key_averages(), щоб визначити, які операції споживають час пристрою Neuron, і порівняти їхній Self CPU time.
• Як простежити одну операцію в Neuron Explorer — від диспетчеризації фреймворком через планування середовищем виконання до виконання на апаратному NeuronCore.
• Як ці сигнали змінюються в трьох реалізаціях однієї й тієї самої операції.

Neuron як бекенд PyTorch

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

Бекенд Neuron використовує точку розширення PrivateUse1 у PyTorch. Цей механізм дає змогу прискорювачу інтегруватися з фреймворком без додавання спеціального ключа диспетчеризації, такого як cuda. Пакет torch_neuronx працює з основною версією PyTorch, реєструючи Neuron як нативний бекенд для коду PyTorch. Докладніше про цей механізм розширення див. у матеріалі Сприяння інтеграції нових бекендів за допомогою PrivateUse1.

Налаштування профайлера

Для профілювання робочого навантаження Neuron використовуються два об’єкти:

• NeuronConfig керує тим, які потоки даних Neuron і 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 див. у документації щодо режимів профілювання 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 має найбільший загальний час на пристрої, оскільки функція виконує три множення.

Загальний Self NEURON: 2.355 мс

Кожен рядок безпосередньо відповідає функції. 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:

У нетерплячому режимі кожна операція виконується незалежно. Використовуйте переглядач ланцюжка залежностей, щоб простежити 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.

Загальний Self 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

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

Загальний Self 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 Storage Buckets

• 

35

13 серпня 2026 року

Від Hugging Face до Amazon SageMaker Studio одним кліком

• 

21

7 липня 2026 року