AWS Trainium Profiling: A Practical PyTorch Guide

TorchNeuron's place in the stack

TorchNeuron connects Trainium to PyTorch using PyTorch's PrivateUse1 extension point. That's how an accelerator backend can participate in PyTorch without requiring a dedicated dispatch key like cuda. In practice, you can keep using standard PyTorch APIs and common execution patterns (eager mode and torch.compile included). The Neuron backend handles the execution on the accelerator.

That compatibility is convenient, but it can also hide details that matter for performance. Two bits of PyTorch code that "do the same math" may still end up producing different backend operations and different hardware dispatches. So the Python side alone isn't a trustworthy performance model.

What system-level profiling can answer

A good first profile should help you answer questions like:

• Which operations take the most Neuron device time?
• How many times does each operation run?
• Is a big device-time total caused by one expensive kernel, or by several smaller executions?
• How much CPU time does the framework spend around those operations?
• Does the timeline show repeated dispatches or idle gaps that fusion could potentially reduce?

One thing to keep straight: a system trace shows where time is spent and how work moves through the stack, but it won't automatically explain every hardware-level mechanism. If you need that, you move on to a device-level investigation (engine activity, memory behavior, and so on).

A practical flow is: use the table to pick a candidate, use the trace to confirm how it runs, and only then test an optimization. This avoids a common mistake - optimizing something that looks complicated instead of something that measurably costs time.

Include the Neuron activity and export callback

To get Neuron device activity, the profiler needs CPU activity plus ProfilerActivity.PrivateUse1. Pass the Neuron configuration as experimental_config, and register the exporter so the framework trace gets written alongside the Neuron system traces.

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()

The schedule above skips two steps, warms up for one, then records one active step. That extra loop iteration is what actually advances the schedule far enough to capture the active window. For your own model, pick a schedule that matches the phase you care about.

Synchronization matters. Accelerator work can run asynchronously: the host may enqueue operations and keep going before the device is done. Calling torch.neuron.synchronize() inside the profiling loop ensures the work finishes before the trace is finalized. If you skip it, the device-time numbers can look oddly incomplete or host-heavy, and you'll chase the wrong explanation.

Keep the capture focused

A profile is only useful if it answers a specific question. Start with the smallest representative workload that exercises the behavior you want to understand. When you compare implementations, keep inputs, shapes, data types, profiler settings, and the execution schedule consistent.

Also, don't confuse profiling overhead with normal execution. Think of the trace as an inspection tool, not a substitute for an end-to-end benchmark. Locate the cause with the trace, then measure the real workload with a repeatable benchmark setup. If your model compiles, keep compile time separate from steady-state inference time - mixing them makes it easy to misread runtime performance.

Follow one operation through Neuron Explorer

Neuron Explorer is a browser-based viewer for system traces, device timelines, and per-engine information. The system trace lays out events across lanes. Time goes left to right, while the lifecycle of each event spans framework, runtime, and hardware lanes.

A typical eager operation usually maps to three stages:

1. Framework dispatch: PyTorch dispatches the operation on a framework thread.
2. Runtime scheduling: the Neuron runtime schedules execution (shown in the walkthrough as nrta_execute_schedule).
3. Hardware execution: the operation runs on a NeuronCore (often shown as a hardware event like nc_exec_running).

Lane labels tell you which subsystem and grouping you're looking at. Framework lanes can be grouped by thread or stream; runtime and hardware lanes are organized by NeuronCore grouping. The exact label names can vary across tooling and configuration, so don't memorize them - use the lane context and event relationships instead.

The dependency-chain view is particularly helpful when you're not sure which hardware event corresponds to a framework operation. Follow one operation from dispatch to scheduling to execution. Then check the timeline for long gaps, repeated small executions, or patterns that line up with the call counts you saw in the table.

Know what files you are opening

A captured profile can include a system trace and, when device data is collected, matching .neff and .ntff files. NEFF is the compiled Neuron executable format; NTFF contains the corresponding hardware instruction trace. A shared numeric identifier links the files.

For a system or combined profile, upload the full output directory to Explorer. For a device-only profile, provide the matching NEFF and NTFF files. If you upload only part of the output, correlation gets harder and you may lose context needed to follow the execution path.

Variant 1: separate eager operations

The naive implementation uses individual operations that create intermediate tensors:

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

Because those operations dispatch separately in eager mode, the device runs multiple executions. In the example run, the profile shows 2,355 μs of total Neuron time and six operator calls in the summary. Intermediates like x_sq and mean_sq are part of the execution path, and the profiler makes the cost of that path easy to see.

Variant 2: the built-in operation

If you replace the sequence with the PyTorch built-in, the backend can lower it as a fused kernel:

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

Here the profile shows one device operation, aten::_fused_rms_norm, with 403.3 μs of Neuron time. That's about one-sixth of the naive version's device time in this run. The change isn't just "less Python" - it's that the computation runs as a fused operation instead of dispatching each arithmetic step independently.

Variant 3: compile the naive function

You can also feed the same sequence into torch.compile using the Neuron backend:

@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

In that flow, TorchDynamo turns the Python function into an FX graph, and then the compilation path sends it to the Neuron backend. From there, the backend can fuse the operations and compile the graph into a single NEFF artifact. The profile shows one compiled region and 375.5 μs of Neuron time.

All three approaches reduce the number of device executions to one in this example. The built-in version expresses a known operation directly; compilation finds the fusion opportunity from the original sequence. The measured device times are close, with the compiled RMSNorm slightly faster in this run. That difference isn't enough to declare a general winner.

What you choose depends on the real workload. Built-ins can be clean and convenient when they match the semantics and shapes you need. Compilation can be great when multiple ops can fuse into a larger region. But compiled graphs come with tradeoffs: graph breaks, unsupported behavior, compile time, and sensitivity to shape changes. Test the actual model instead of extrapolating from a micro-example.

A practical decision process

Use this workflow when you profile a larger inference workload:

1. Find the hotspot. Sort the averages table by device time and note both call count and total contribution.
2. Confirm the execution path. Use Explorer to connect the framework event to runtime scheduling and hardware execution.
3. Form a specific hypothesis. For example: repeated dispatches or materialized intermediates might make fusion worth testing.
4. Change one thing. Try a supported fused operator or compile the relevant region.
5. Measure the full workload again. Re-check device time, CPU overhead, end-to-end latency, and correctness/constraint impact.

Fusion isn't automatically beneficial. It can reduce dispatch overhead and intermediate materialization, but the fused kernel might behave differently in terms of resources. A system-level trace tells you where to look; once you know the candidate, the device-level profile is the next step for engine utilization, memory traffic, and whether the kernel is compute-bound or memory-bound.