-
Notifications
You must be signed in to change notification settings - Fork 461
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
[PJRT] Only alias inputs and outputs when force_ltc_sync == True
#4575
Conversation
51a3639
to
00fd8db
Compare
I ran our imagenet benchmark on TPU v4 with the PJRT C API. Performance is actually a little better than the last time I checked, but that could be related to another change. |
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.
Generally speaking, LGTM. Let's see how CI respond.
@@ -1102,7 +1102,8 @@ XLAGraphExecutor::CompilationResult XLAGraphExecutor::Compile( | |||
// TODO(yeounoh) aliasing is disabled for partitioned computation, | |||
// since the current aliasing compares the unpartitioned input and output | |||
// shapes which can lead to an incorrect aliasing pairs if sharded. | |||
if (enable_aliasing && coll.config.sync_ltc_data && !is_sharded) { | |||
if (enable_aliasing && coll.config.sync_ltc_data && | |||
coll.config.force_ltc_data && !is_sharded) { |
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.
as a follow up, we should add a test to https://github.com/pytorch/xla/blob/master/test/test_input_output_aliases.py and make sure only mark_step
trigger the aliasing.
Updating logic to match comment.
According to the documentation for PJRT, when an input buffer is aliased to an output of an executable, the original input buffer is no longer available. When we reference the
DataPtr
that holds the buffer again, we will inevitably see this error:RuntimeError: ./tensorflow/compiler/xla/xla_client/pjrt_computation_client.h:140 : Check failed: HasValue()
Fixes the following failing C++ tests on CPU: