[kvshbg-aws]

Summary

On AWS Neuron (Trainium), build_torch_profiler recorded only CPU activity. This PR adds ProfilerActivity.PrivateUse1 when device_type == "neuron", in the same way the CUDA and XPU branches add their activities. No other device changes.

The Neuron branch also sets with_stack=True. Neuron Explorer uses the stacks to link framework events to source lines. The stacks add host overhead only on profiled steps.

Notes

• requires torch-neuronx>=2.33 version This PR adds no Neuron imports or configuration.
• TORCH_NEURONX_PROFILE_* environment variables set the Neuron profile options. The defaults capture rank 0.
• Logs, metrics, MFU, and memory snapshots already work on Neuron through device_module. This PR does not change them.

Test plan

pytest tests/unit_tests/cpu/observability/
pre-commit run --files torchtitan/observability/profiler.py tests/unit_tests/cpu/observability/test_profiler.py

• The unit tests give 135 passed. The new test checks the activities and with_stack on a mocked Neuron device.
• Ran the profiler on Neuron hw with Qwen3 8B (64 NeuronCores). Rank 0 trace:
  • 34K privateuse1_runtime and 22K privateuse1_driver events;
  • source-linked framework events;
  • a Neuron session that opens in Neuron Explorer.

[aws-arashj]

let's reword

requires a torch-neuronx change, that has been merged to the torch-neuronx repo

to requires totorch-neuronx ≥ 2.33 version

also

the trn2.48xlarge / 64-NeuronCore

to verified on Neuron hardware.

[kvshbg-aws]

closing this one since we have the upstream PR open - pytorch#5072