-
Notifications
You must be signed in to change notification settings - Fork 293
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add XLA support to moco
benchmark.
#2292
Conversation
The CI failures don't seem related to this PR. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sounds good to me.
FYI @aaronenyeshi and @JackCaoG
) | ||
except RuntimeError: | ||
pass # already initialized? | ||
elif device.startswith("xla"): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit for consistency device == "xla" ?
4c46871
to
eb46de2
Compare
@xuzhao9 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
This PR tweaks
moco
benchmark, so that it will also run on XLA devices. Previously,moco
hardcoded the CUDA device in two ways:ProcessGroup
withnccl
backend, onlycuda
, explicitlyIn order to add XLA support, this PR:
xla*
devices and, if detected, initializes theProcessGroup
withxla
backendcc @lezcano