Add Chakra callback

[TaekyungHeo]

Important

The Update branch button must only be pressed in very rare occasions.
An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.

What does this PR do?

This PR introduces the ChakraCallback, a new PyTorch Lightning callback that collects Chakra host and device traces during training.

Collection: nemo.lightning.pytorch.callbacks

Changelog

• Added ChakraCallback to nemo.lightning.pytorch.callbacks.
• Implemented profiling support for capturing PyTorch execution traces (host) and Kineto traces (device).
• Updated __init__.py to include ChakraCallback.

Usage

Example of how to use the ChakraCallback:

from nemo.lightning.pytorch.callbacks.chakra import ChakraCallback

chakra_callback = ChakraCallback(
    start_step=5,
    end_step=10,
    warmup_steps=1,
    active_steps=1,
    trace_dir="/workspace/"
)

trainer = Trainer(callbacks=[chakra_callback])

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related PR: Profiling - support Chakra & Kineto trace dumping #11115

[titu1994]

Looks reasonable

[titu1994]

Rename to PytorchProfilerCallback rather than some code name.

[titu1994]

Inner function cannot be serialized, is it necessary? Maybe move it out.

[maanug-nv]

If trace_dir is required, should it have no default?

[maanug-nv]

Adding a short test would be good as well, to ensure all files that are expected to be created do exist. Fine with a separate PR for that

[github-actions]

This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.