Skip to content

Add profiler API for async capture#5969

Merged
jonb377 merged 2 commits intomasterfrom
jonbolin/trace
Dec 4, 2023
Merged

Add profiler API for async capture#5969
jonb377 merged 2 commits intomasterfrom
jonbolin/trace

Conversation

@jonb377
Copy link
Copy Markdown
Collaborator

@jonb377 jonb377 commented Dec 1, 2023

To more easily support programmatic captures, this changes add a new API trace_detached to torch_xla.debug.profiler to capture a profile from a background thread.

Tested by using the new API in Llama2 training.

@jonb377 jonb377 requested a review from JackCaoG December 1, 2023 20:19
@jonb377 jonb377 self-assigned this Dec 1, 2023
Copy link
Copy Markdown
Collaborator

@JackCaoG JackCaoG left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks a lot! Can we add a test to https://github.com/pytorch/xla/blob/master/test/test_profiler.py ? We already have a test for trace so I imagine you can mostly just copy paste the test and make small modifications.

@jonb377
Copy link
Copy Markdown
Collaborator Author

jonb377 commented Dec 1, 2023

Thanks a lot! Can we add a test to https://github.com/pytorch/xla/blob/master/test/test_profiler.py ? We already have a test for trace so I imagine you can mostly just copy paste the test and make small modifications.

Thanks @JackCaoG - I generalized the test_profile_mp_mnist script to support programmatic captures and added a test case

Copy link
Copy Markdown
Collaborator

@JackCaoG JackCaoG left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

@jonb377 jonb377 merged commit 6d73ca8 into master Dec 4, 2023
@jonb377 jonb377 deleted the jonbolin/trace branch December 4, 2023 19:38
chunnienc pushed a commit to chunnienc/xla that referenced this pull request Dec 14, 2023
* Add profiler API for async capture

* Add unit test
golechwierowicz pushed a commit that referenced this pull request Jan 12, 2024
* Add profiler API for async capture

* Add unit test
bhavya01 pushed a commit that referenced this pull request Apr 22, 2024
* Add profiler API for async capture

* Add unit test
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants