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

Много специалисти по машинно обучение имат известна интуиция как да накарат моделите да работят бързо, но по-малко хора знаят как да анализират защо даден модел е бавен. В част 1 от тази поредица ще научите как систематично да анализирате производителността при инференция и да определяте къде се изразходва времето.

Тази публикация започва от нулата и не предполага предишен опит с профилиране. Разходката използва AWS Trainium — специално създаден AI чип за натоварвания за обучение и инференция на модели за дълбоко обучение — и TorchNeuron, неговия нативен бекенд за PyTorch с отворен код. TorchNeuron поддържа стандартни работни процеси на PyTorch върху Trainium, включително eager изпълнение, 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 използва разширителната точка PrivateUse1 на PyTorch. Този механизъм позволява на ускорител да се интегрира с фреймуърка, без да добавя специален ключ за изпращане като cuda. Пакетът torch_neuronx работи с upstream 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 вижте документацията за режимите на 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))

Таблицата с ключовите средни стойности показва, че три извиквания на aten::mul представляват най-големия дял от времето на устройството Neuron.

Таблицата показва кои ATen операции са се изпълнили, колко често са се изпълнили и колко време са прекарали върху CPU и устройството Neuron. aten::mul има най-високо общо време на устройството, тъй като функцията извършва три умножения.

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

Всеки ред съответства директно на функцията. aten::mul се изпълнява за x * x, x * rms и финалното умножение по weight. aten::mean редуцира последното измерение. aten::add добавя стойността epsilon и задейства копиране от хост към устройство, тъй като скаларът започва върху 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

Заменете отделните операции с вградената функция rms_norm на PyTorch. Neuron свежда операцията до едно слято ядро:

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

Слятият RMSNorm се появява като една операция на устройството в таблицата с ключовите средни стойности.

Общо собствено 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 в таблицата с ключовите средни стойности.

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

Компилираният граф следва един път през планирането от средата за изпълнение и хардуерното изпълнение.

Таблицата показва едно изпълнение на устройството, както в профил 2, но кодът следва различен път. Вместо да извиква специално създадена слята операция, компилаторът открива сливането от оригиналната функция на Python. При това изпълнение времето на Neuron е сравнимо с torch.rms_norm — 375 μs спрямо 403 μs.

Обобщение едно до друго

И 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 и Storage Buckets на Hugging Face

• 

35

13 август 2026 г.

От Hugging Face до Amazon SageMaker Studio с едно кликване

• 

21

7 юли 2026 г.