This document is relevant for: Inf2, Trn1, Trn2

PyTorch NeuronX Profiling API#

The profiler provides a method to generate a context manager to capture trace events at the operator or runtime level.

torch_neuronx.experimental.profiler.profile(port=9012, ms_duration=60000, neuron_tensorboard_plugin_dir='logs/plugins/neuron', profile_type='operator', auto_start=True, delete_working=True)#

The torch_neuronx.experimental.profiler.profile() method returns a profile context manager object. This object doesn’t need to be used directly, as default options are set to auto capture events based on the profile_type.

The context manager will wrap around the entire model and training/inference loop. The context-manager is backwards-compatible with the torch_xla.debug.profiler``

Required Arguments

None

Optional Keyword Arguments

Keyword Arguments:

  • port (int) – Port to run the profiling GRPC server on. Default is 9012.

  • ms_duration (int) – This defines how long the profiler will capture the HLO artifacts from the model to view in the profiler. The unit is in milliseconds. The default value is 60000 ms, or 1 minute.

  • neuron_tensorboard_plugin_dir (str) – The directory the neuron tensorboard plugin will file write to. This will be logs/plugins/neuron by default/

  • profile_type (str) – There is “trace” and “operator”. “trace” is the Torch Runtime Trace Level, while “operator” is the Model Operator Trace Level. Default is “operator”

  • auto_start (bool) – If set to true, the profiler will start profiling immediately. If set to false, the profiler can be set to start at a later condition. Refer to profile.start() for more details. Default is True.

  • delete_working (bool) – If set to False turns off the deletion of temporary files. Default True.

  • traced_only (str) – This should be set to True if profiling a model that has been traced with torch_neuronx.trace(). Default is False.

Returns:

The traced profile

Return type:

~profile

torch_neuronx.experimental.profiler.profile.start()#

The torch_neuronx.experimental.profiler.profile.start() method starts the profiler if not started (i.e when auto_start=False). This function does not take in any parameters, nor return anything.

Required Arguments

None

Optional Keyword Arguments

None

Returns:

None

This document is relevant for: Inf2, Trn1, Trn2