[SPMD] Use sharding annotation from XLATensor data as ground truth.#4390
Merged
[SPMD] Use sharding annotation from XLATensor data as ground truth.#4390
Conversation
39b4aa0 to
4d9718d
Compare
f111f2f to
5f32719
Compare
5f32719 to
c025f99
Compare
JackCaoG
reviewed
Dec 27, 2022
JackCaoG
reviewed
Dec 27, 2022
354e1e9 to
b3b76e1
Compare
e0105e8 to
346a50c
Compare
JackCaoG
reviewed
Jan 4, 2023
JackCaoG
reviewed
Jan 4, 2023
JackCaoG
reviewed
Jan 4, 2023
ac0439c to
192fdfa
Compare
…Tensosr is created from a sharded IR node.
192fdfa to
ff267ab
Compare
89452a7 to
2c9f233
Compare
JackCaoG
reviewed
Jan 4, 2023
JackCaoG
reviewed
Jan 4, 2023
| ->SetSharding(sharding->sharding); | ||
| } | ||
| data()->ir_value = std::move(ir_value); | ||
| data()->generation += 1; |
Collaborator
There was a problem hiding this comment.
what if passed in ir_value has a sharding spec?
Contributor
Author
There was a problem hiding this comment.
The check will be done at the sharding_spec() call.
Collaborator
There was a problem hiding this comment.
I am afraid this will make debugging a bit difficult, since we can only see that there is a mismatch but it is hard to figure out why(error message is lazy). The code is cleaner this way I guess, so I am also OK with this approach.
Contributor
Author
There was a problem hiding this comment.
Yea, it's cleaner this way that we have a limited (consolidated) places, where we do the checks and the syncing. I see your point, though. Let's add some debugging probes TF_VLOG for that instead.
Collaborator
There was a problem hiding this comment.
let's do
TF_VLOG(5) << "Assign IR value " << ir_value->ToString() + sharding ? (" with sharding " + sharding->DebugString()) : "";
This way when TF_VLOG is not enabled, we don't construct the debug string.
JackCaoG
reviewed
Jan 4, 2023
JackCaoG
approved these changes
Jan 4, 2023
7ebc24b to
cb5a9db
Compare
* [SPMD] Add a debug probe in AssignIrValue()
cb5a9db to
406dea6
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Use
XLATensordata to hold a copy of user specified sharding annotation. This way, the IR node truncation and reset should not accidentally remove the sharding annotation, which should be removed via explicitclear_sharding.XLATensor::sharding_specreturns sharding fromXLATensor::data(source of the truth)XLATensor::Dataconstructors now havesharding(default tonullptr) argumentCollectSyncTensorscheck for sharding equality if IR sharding exists, and syncs the shardings fromXLATensor::Datato IR nodes otherwise.XLATensorshould copy sharding annotationXLATensorshould persist sharding annotationAlso,
AssignIrValue.