Profiling TorchNeuron (Part 1): System-level profiling on AWS Trainium.

What it claims

The torch.profiler workflow on Trainium needs a few Neuron-specific settings. Without them, the device-time column comes back empty or traces are hard to open. The test case is a small RMSNorm: a bfloat16 input of shape (4, 512) and a 512-wide weight, profiled on NeuronCore 0. On that input, both fused paths (torch.rms_norm at 403.3 μs and torch.compile at 375.5 μs) use about one-sixth of the naive eager device time of 2.355 ms. The post credits fusion with removing the gaps between hardware executions and not materializing intermediates such as x_sq and mean_sq.

How it works

TorchNeuron registers Neuron through PyTorch's PrivateUse1 extension point, not through a dedicated dispatch key. Profiling therefore uses ProfilerActivity.PrivateUse1, and device time appears under self_privateuse1_time_total.

In eager mode, each ATen op goes through three lanes in Neuron Explorer:

1. Creation in the framework lane.
2. nrta_execute_schedule in the runtime lane.
3. nc_exec_running in the hardware lane.

The naive version dispatches three aten::mul, plus aten::mean, aten::add and aten::rsqrt. The aten::add also triggers a host-to-device copy, because the epsilon scalar starts on the CPU.

With torch.compile(backend='neuron'), TorchDynamo captures an FX graph, which AOTAutograd lowers to ATen. The Neuron backend then compiles the whole graph into one NEFF, which runs as a single hardware dispatch and appears as one "Torch-Compiled Region".

How to use it

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

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

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

Practical points from the page:

• correlate_device_time=True is what puts Neuron device time into the key_averages() table next to CPU time.
• torch.neuron.synchronize() must sit inside the loop. Without it, export can happen before the device finishes, and the Neuron time column is empty.
• The schedule records one warm iteration, leaving out setup overhead.
• The output directory holds trace.json, the system traces and matched .neff/.ntff pairs that share a numeric hash.
• Neuron Explorer runs in the browser through an SSH tunnel or the VS Code extension. Upload the whole directory for a system or combined profile, or the matching .neff/.ntff pair for a device-only profile.
• Compile with @torch.compile(backend='neuron', fullgraph=True).
• The scripts are in the torch-neuronx examples directory.

Limits

Each figure comes from a single warm iteration of one small tensor on one NeuronCore. The page reports no repeated runs or variance, and it calls the 375 μs versus 403 μs gap "comparable" rather than a win for torch.compile. No larger model is profiled. NeuronConfig is passed through experimental_config, the field PyTorch uses for experimental profiler options. Device-level analysis (TensorEngine, VectorEngine and DMA breakdown, and roofline placement) is left to Part 2.