Skip to content

Enable eager mode for PyTorch/XLA#7234

Merged
JackCaoG merged 10 commits intomasterfrom
JackCaoG/eager_mode
Jun 11, 2024
Merged

Enable eager mode for PyTorch/XLA#7234
JackCaoG merged 10 commits intomasterfrom
JackCaoG/eager_mode

Conversation

@JackCaoG
Copy link
Copy Markdown
Collaborator

@JackCaoG JackCaoG commented Jun 10, 2024

Add a eager mode for PyTorch/XLA. We already have an eager mode for debug which is UseEagerDebugMode that's controlled by a env var. This was contributed by the AWS folks(cc @aws-rhsoln @amithrm @jeffhataws). I decided to make a offical top level api for it.

In this pr, I

  1. add api use_eager_mode along with all the pybinds
  2. make sure that PT_XLA_DEBUG and PT_XLA_DEBUG_LEVEL does not print analysis for eager execution
  3. add new metrics for eager mode compilation and eager mode executions

For future pr

  1. increase the maximum number of queued task (for eager mode)
  2. add api/decorator to compile a specified function in the eager mode.
  3. add docs
  4. test eager + SPMD

For a 2 layer decoder only model, eager mode can acheive ~40% of throuput compared to the full compiled mode.

@JackCaoG JackCaoG added the usability Bugs/features related to improving the usability of PyTorch/XLA label Jun 10, 2024
@lsy323
Copy link
Copy Markdown
Collaborator

lsy323 commented Jun 10, 2024

Let's add a test, otherwise LGTM

@JackCaoG JackCaoG added the tpuci label Jun 10, 2024
@JackCaoG
Copy link
Copy Markdown
Collaborator Author

The correctness check is handled by run_eager_debug "$CDIR/test_operations.py" "$@" --verbosity=$VERBOSITY. I will expand this in the future.

@JackCaoG
Copy link
Copy Markdown
Collaborator Author

Ah.. I need to move this api under experimental.. Let me update the pr...

@JackCaoG
Copy link
Copy Markdown
Collaborator Author

I am going to skip the gpu tests and it takes forever to get resouce. From last few runs of the CI in this pr, there is nothing specified about the GPU.

@JackCaoG JackCaoG merged commit 4d0b94f into master Jun 11, 2024
@JackCaoG JackCaoG deleted the JackCaoG/eager_mode branch June 11, 2024 20:12
@JackCaoG JackCaoG added the eager PyTorch/XLA eager-mode label Jun 20, 2024
JackCaoG added a commit that referenced this pull request Jul 2, 2024
JackCaoG added a commit that referenced this pull request Jul 8, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

eager PyTorch/XLA eager-mode usability Bugs/features related to improving the usability of PyTorch/XLA

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants