Skip to content

Conversation

@jonb377
Copy link
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
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
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
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

3 participants