Add Pytorch Profiler

[frobnitzem]

Description

This adds a training callback that stores pytorch profiling outputs following the simple config at https://pytorch.org/tutorials/intermediate/tensorboard_profiler_tutorial.html.

Motivation and Context

How Has This Been Tested?

This has been used to produce and view profile outputs.

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds or improves functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation improvement (updates to user guides, docstrings, or developer docs)

Checklist:

• My code follows the code style of this project and has been formatted using black.
• All new and existing tests passed, including on GPU (if relevant).
• I have added tests that cover my changes (if relevant).
• The option documentation (docs/options) has been updated with new or changed options.
• I have updated CHANGELOG.md.
• I have updated the documentation (if relevant).

[Linux-cpp-lisp]

Thanks @frobnitzem! Will be very happy to merge and have this feature.

One small thing: can you run the code through black: https://nequip.readthedocs.io/en/develop/dev/contributing.html#style-enforcement?

And one bigger thing:

The TensorBoard integration with the PyTorch profiler is now deprecated. Instead, use Perfetto or the Chrome trace to view trace.json files.
https://pytorch.org/tutorials/intermediate/tensorboard_profiler_tutorial.html

Would it make sense to add support in this callback for the newer export_chrome_trace API (https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html#using-tracing-functionality), with some option to flip between the older Tensorboard and newer Perfetto/chrome tracing in the on_trace_ready callback?

Finally, should any of the scheduler or profiler options be exposed optionally in the arguments to the callback?

Thanks again for the PR!!

[frobnitzem]

Still a work in progress. Trying this in October, I got much better profile info, so will probably scrap the tensorboard analysis in favor of the workflow you linked here.

[Linux-cpp-lisp]

Sounds great, thanks!