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

Many ML practitioners have some instinct about how you can make fashions run quick, however fewer know how you can analyze why a mannequin is gradual. In Half 1 of this sequence, you’ll learn to systematically analyze inference efficiency and establish the place the time is spent.

This put up begins from scratch and assumes no prior profiling expertise. The walkthrough makes use of AWS Trainium, a purpose-built AI chip for deep studying coaching and inference workloads, and TorchNeuron, its open supply native PyTorch backend. TorchNeuron helps commonplace PyTorch workflows on Trainium, together with keen execution, torch.compile, and commonplace PyTorch APIs, with minimal code adjustments. Moreover, TorchNeuron integrates with Neuron Kernel Interface (NKI) to allow customized operators and efficiency optimization for Trainium {hardware}. By the tip, you must know:

• Easy methods to configure torch.profiler with the Neuron backend and seize a clear hint.
• Easy methods to learn the key_averages() desk to establish which operations devour Neuron gadget time and examine their Self CPU time.
• Easy methods to observe one operation by way of Neuron Explorer, from framework dispatch by way of runtime scheduling to {hardware} execution on a NeuronCore.
• How these indicators change throughout three implementations of the identical operation.

Neuron as a PyTorch Backend

Earlier than you profile a workload, it helps to grasp how Neuron integrates with PyTorch.

The Neuron backend makes use of PyTorch’s PrivateUse1 extension level. This mechanism lets an accelerator combine with the framework with out including a devoted dispatch key reminiscent of cuda. The torch_neuronx bundle works with upstream PyTorch to register Neuron as a local backend for PyTorch code. For particulars about this extension mechanism, see Facilitating New Backend Integration by PrivateUse1.

Set Up the Profiler

You employ two objects to profile a Neuron workload:

• NeuronConfig controls which Neuron knowledge streams and NeuronCores to seize and the place to write down the output. You move it to torch.profiler.profile as experimental_config.
• NeuronProfiler takes a NeuronConfig and gives an export_trace callback. The callback writes the framework hint, hint.json, alongside the Neuron system traces.

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

The modes parameter selects the info streams to seize, reminiscent of gadget directions and runtime traces. The capture_enabled_for_nc parameter selects the NeuronCores to hint. Setting correlate_device_time=True propagates gadget time from the hint into the key_averages() desk, the place you’ll be able to examine CPU and Neuron gadget time. For the whole parameter reference and obtainable ProfileMode values, see the NeuronConfig Profile Modes documentation.

Tooling: Neuron Explorer

Neuron Explorer is a browser-based viewer for the system hint, gadget timeline, and per-engine breakdown. You’ll be able to entry it by way of an SSH tunnel or the Visible Studio Code extension. On this put up, you utilize Explorer to substantiate what key_averages() experiences numerically and examine what occurs between dispatches.

After profiling, the output listing accommodates the system hint and, once you seize gadget knowledge, matched .neff and .ntff recordsdata. The Neuron Executable File Format (NEFF) accommodates the compiled Neuron executable, whereas the matching NTFF accommodates its {hardware} instruction hint. The recordsdata share a numeric hash that identifies them as a pair. To open the profile in Neuron Explorer, add the whole output listing for a system or mixed profile. For a device-only profile, add the matching .neff and .ntff recordsdata individually. See Step 5: Add your profile for the UI and command-line directions.

Profile RMSNorm

For this primary evaluation, you profile RMSNorm, a small operation with a transparent sequence of arithmetic and discount operations.

You profile three implementations of the identical operation:

1. Naive RMSNorm with particular person PyTorch operations: multiply, imply, add, reciprocal sq. root, and multiply.
2. torch.rms_norm(), the PyTorch built-in lowered to a single fused Neuron kernel.
3. Naive RMSNorm with torch.compile, utilizing the identical operations as the primary implementation.

The scripts use the identical inputs, profiler configuration, and schedule so the outcomes are instantly comparable.

All scripts to breed this work can be found within the torch-neuronx examples listing.

Profile 1: Naive PyTorch Operations

Begin by writing RMSNorm with particular person PyTorch operations. Every line maps to a number of ATen operations and may create intermediate tensors.

gadget = torch.gadget('neuron')

x = torch.randn(4, 512, dtype=torch.bfloat16, gadget=gadget)
weight = torch.ones(512, dtype=torch.bfloat16, gadget=gadget)

def rms_norm_fn(t):
    x_sq = t * t                                    
    mean_sq = x_sq.imply(dim=-1, keepdim=True)       
    rms = torch.rsqrt(mean_sq + 1e-6)               
    return t * rms * weight                         

Name the profiler with two arguments required for the Neuron backend:

1. Embody ProfilerActivity.PrivateUse1 to gather exercise from the Neuron gadget.
2. Go NeuronConfig as experimental_config to allow the Neuron-specific Kineto options.

The schedule makes use of wait=2, warmup=1, and lively=1 throughout 4 prof.step() calls. This data one heat iteration with out together with preliminary setup overhead.

Name torch.neuron.synchronize() contained in the profiler loop. With out this synchronization, the profiler can export earlier than the gadget finishes and go away the Neuron time column empty.

with torch.profiler.profile(
    actions=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1],
    experimental_config=nc,
    record_shapes=False,
    on_trace_ready=exporter.export_trace,
    schedule=torch.profiler.schedule(wait=2, warmup=1, lively=1),
) as prof:
    for step in vary(4):
        with torch.no_grad():
            _ = rms_norm_fn(x)
        torch.neuron.synchronize()
        prof.step()

After profiling finishes, examine the key_averages() desk. Kind by self_privateuse1_time_total to position the operations with the very best gadget value first. The desk proven right here is filtered to the RMSNorm operations that report Neuron gadget time. The total key_averages() output additionally contains CPU-side occasions reminiscent of allocations and host-to-device transfers.

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

The important thing averages desk reveals that three aten::mul calls account for the biggest share of Neuron gadget time.

The desk reveals which ATen operations ran, how usually they ran, and the way a lot time they spent on the CPU and Neuron gadget. aten::mul has the very best complete gadget time as a result of the perform performs three multiplications.

Self NEURON complete: 2.355 ms

Every row maps on to the perform. aten::mul runs for x * x, x * rms, and the ultimate multiplication by weight. aten::imply reduces the final dimension. aten::add provides the epsilon worth and triggers a host-to-device copy as a result of the scalar begins on the CPU. aten::rsqrt computes the reciprocal sq. root.

Neuron Explorer follows an aten::mul operation from framework dispatch by way of runtime scheduling to execution on the Neuron {hardware} lanes.

Learn the Hint in Explorer

In Explorer, time flows from left to proper and the occasion lifecycle flows from high to backside. Every ATen operation follows the identical sample: creation within the framework lane, scheduling within the runtime lane as nrta_execute_schedule, and gadget execution within the {hardware} lane as nc_exec_running.

The system hint view organizes occasions into lanes. Within the Default grouping mode, labels observe the sample occasion/PID/subsystem/identifier:

In keen mode, every operation runs independently. Use the dependency chain viewer to observe aten::mul.Tensor from the framework layer by way of runtime scheduling to gadget execution.

The dependency chain viewer highlights the lifecycle of 1 operation throughout the framework, runtime, and {hardware} lanes.

Profile 2: torch.rms_norm

Exchange the person operations with PyTorch’s built-in rms_norm perform. Neuron lowers the operation to 1 fused kernel:

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

The fused RMSNorm seems as one gadget operation in the important thing averages desk.

Self NEURON complete: 403.3 μs

The fused RMSNorm follows one framework-to-hardware path as a substitute of dispatching every arithmetic operation individually.

Examine the Decreasing Path

Within the naive implementation, every operation crosses the dispatch boundary independently and follows its personal framework-to-runtime-to-hardware path. Within the fused implementation, one Python name covers the whole computation. The hint accommodates one dispatch and one {hardware} occasion, with no gaps between element operations.

Profile 3: Naive PyTorch with torch.compile

For the ultimate instance, compile the unique naive RMSNorm perform with torch.compile.

With @torch.compile(backend="neuron"), TorchDynamo captures the Python bytecode as an FX graph and AOTAutograd lowers it to ATen operations. The Neuron backend receives the whole graph as a substitute of dispatching every operation individually. It fuses the operations, optimizes reminiscence entry, and compiles the graph right into a single NEFF artifact. The artifact runs as one {hardware} dispatch.

@torch.compile(backend='neuron', fullgraph=True)
def naive_compiled(t):
    x_sq = t * t
    mean_sq = x_sq.imply(dim=-1, keepdim=True)
    rms = torch.rsqrt(mean_sq + 1e-6)
    return t * rms * weight

The compiled perform seems as one Torch-Compiled Area in the important thing averages desk.

Self NEURON complete: 375.5 μs

The compiled graph follows one path by way of runtime scheduling and {hardware} execution.

The desk reveals one gadget execution, as in Profile 2, however the code takes a unique path. As an alternative of calling a purpose-built fused operation, the compiler discovers the fusion from the unique Python perform. On this run, the Neuron time is corresponding to torch.rms_norm, at 375 μs in contrast with 403 μs.

Aspect-by-Aspect Abstract

Each torch.rms_norm and torch.compile fuse RMSNorm into one kernel dispatch. In these runs, both strategy takes about one-sixth of the naive implementation’s gadget time. Fusion removes the gaps between {hardware} executions and avoids materializing intermediate values reminiscent of x_sq and mean_sq between individually dispatched operations.

Conclusion

key_averages() reveals which operations devour gadget time, whereas its full output additionally reveals CPU-side setup and transfers. Explorer gives the corresponding visible hint. You’ll be able to observe one operation from framework dispatch by way of runtime scheduling to {hardware} execution and see the gaps that fusion removes.

You should use the identical workflow on bigger fashions: discover a hotspot within the desk, confirm it within the hint, after which decide whether or not a fused operation or torch.compile can scale back its value.

In Half 2, we are going to transfer into the device-level profile. You open the NEFF, examine the TensorEngine, VectorEngine, and direct reminiscence entry (DMA) breakdown, and place the kernel on the roofline to find out whether or not it’s compute-bound or memory-bound.

Particular due to Aritra Roy Gosthipaty for the inspiration to start this sequence, and to my friends Pinak Panigrahi, Samir Araujo, Akhil Azhikodan, Esha Lakhotia and the remaining for reviewing!