Hugging Face has published Part 1 of a TorchNeuron profiling series, a tutorial for tracing how RMSNorm inference runs on AWS Trainium. It explains how to collect PyTorch, Neuron runtime and device traces, then compare three implementations of the operation; the source material describes the method but does not report which implementation was fastest.

The walkthrough uses TorchNeuron, an open-source PyTorch backend for Trainium, and PyTorch’s profiler. The tutorial says TorchNeuron supports standard PyTorch workflows including eager execution and torch.compile, and connects with the Neuron Kernel Interface for custom operators. It explains that Neuron registers through PyTorch’s PrivateUse1 extension point, which allows an accelerator backend without a dedicated dispatch key such as CUDA’s.

For profiling, the tutorial configures NeuronConfig to select data streams, NeuronCores and an output directory, then passes it to torch.profiler.profile as experimental_config. It also uses NeuronProfiler to export the framework trace alongside Neuron system traces. The example enables CPU and PrivateUse1 activities, correlates device time with the profiler table, and synchronizes the device in the loop before advancing each profiler step.

The comparison covers three versions of RMSNorm: a naive sequence of PyTorch operations, the built-in torch.rms_norm() lowered to a fused Neuron kernel, and the naive version compiled with torch.compile. The source says the examples use the same inputs, profiler configuration and schedule. It points readers to the TorchNeuron examples directory for reproduction and describes Neuron Explorer as the browser-based tool for inspecting system traces, device timelines and per-engine activity.

Tracing Work Across the Neuron Stack

The tutorial addresses a practical distinction for machine-learning teams: recognizing that inference is slow is different from identifying where time is spent. A profiler table can show Neuron device time alongside Self CPU time, while a system trace can follow an operation from framework dispatch through runtime scheduling to hardware execution on a NeuronCore. Together, those views can help practitioners investigate whether delays appear in framework work, runtime handling or device execution.

That is useful when evaluating alternative implementations of the same operation. A built-in fused operation and a compiled sequence may change the work sent to the accelerator, but this source offers a method for comparing them, not evidence of a performance winner. The distinction matters for teams deciding whether to adopt a particular implementation or spend time optimizing a model.

Three RMSNorm Implementations

RMSNorm is used here as a small example with a clear sequence of arithmetic and reduction operations. The naive implementation squares its input, computes the mean of the squares, adds an epsilon, applies reciprocal square root, and scales the input by that result and a weight. The source uses a 4-by-512 bfloat16 input on a Neuron device and profiles repeated calls under a shared configuration.

The profiler schedule is set to wait for two steps, warm up for one, and record one active step across four calls. According to the tutorial, this leaves initial setup out of the recorded iteration. The instructions also say to call torch.neuron.synchronize() in the loop so the device finishes before export; without synchronization, the Neuron time column may be empty. Neuron Explorer can open a complete output directory for a system or combined profile, or matching NEFF and NTFF files for a device-only profile. The source describes NEFF as the compiled executable and NTFF as its hardware instruction trace.

“How to follow one operation through Neuron Explorer, from framework dispatch through runtime scheduling to hardware execution on a NeuronCore.”

— Hugging Face tutorial

Performance Results Are Not Included

The supplied source ends during the first profile walkthrough and does not include the resulting key_averages() table, Neuron Explorer findings, timing measurements or comparison among the three implementations. It does not establish which version uses less device time or CPU time, or whether the measured results generalize beyond the example inputs and setup. No publication date is supplied, and the source does not state when the series’ next part will appear.

Compare the Captured Traces

The next step within the described walkthrough is to inspect the exported profiler table and trace, then repeat the same process for the built-in and compiled RMSNorm versions. Readers can use the linked example scripts to reproduce the profiles and Neuron Explorer to inspect device and runtime activity. The source identifies the series as Part 1 but gives no schedule or details for a subsequent installment.

Where I land

I see this as a useful measurement tutorial because it connects profiler numbers to a trace of framework, runtime and hardware activity. That can give practitioners a clearer basis for investigating slow inference than timing a model call alone. The source, however, establishes a workflow rather than a performance conclusion.

The strongest counterargument is that a walkthrough centered on one small RMSNorm example may not tell teams much about their own full models, where shapes, batching and other operations affect results. I would revise my assessment if published measurements showed repeatable comparisons across representative workloads, with the inputs, software versions and hardware configuration documented. Until then, the value here is the profiling method, not a demonstrated speedup.

Key Questions

What does the TorchNeuron tutorial cover?

It shows how to profile RMSNorm inference on AWS Trainium with PyTorch Profiler and Neuron’s tools, including how to compare CPU and device timing.

Which RMSNorm versions are compared?

The tutorial describes a sequence of individual PyTorch operations, the built-in torch.rms_norm() lowered to a fused Neuron kernel, and the naive implementation compiled with torch.compile.

Does the source say which version is fastest?

No. The supplied material explains the profiling setup but includes no timing results or performance ranking.

Why does the example synchronize the Neuron device?

The tutorial says synchronization lets the device finish before the profiler exports its trace. Without it, the Neuron time column may be empty.

Source: Hugging Face