[NVIDIA GPU] Fix mem p2p init in collective permute thunk #20086
Conversation
This file contains 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
9d3a8a4 to
10d3501
Compare
memcpy to make sure data consistency across ranks
10d3501 to
050bc59
Compare
copybara-service bot
pushed a commit
that referenced
this pull request
Dec 12, 2024
Imported from GitHub PR #20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad04 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59 by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59 PiperOrigin-RevId: 705647424
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 12, 2024
Imported from GitHub PR openxla/xla#20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad0445f27d7249b4bcebb4ac573188cf50cb0 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59c02732da728fe43bd6c4c12702d070c2c by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59c02732da728fe43bd6c4c12702d070c2c PiperOrigin-RevId: 705647424
copybara-service bot
pushed a commit
that referenced
this pull request
Dec 13, 2024
Imported from GitHub PR #20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad04 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59 by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59 PiperOrigin-RevId: 705647424
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 13, 2024
Imported from GitHub PR openxla/xla#20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad0445f27d7249b4bcebb4ac573188cf50cb0 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59c02732da728fe43bd6c4c12702d070c2c by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59c02732da728fe43bd6c4c12702d070c2c PiperOrigin-RevId: 705647424
copybara-service bot
pushed a commit
that referenced
this pull request
Dec 13, 2024
Imported from GitHub PR #20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad04 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59 by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59 PiperOrigin-RevId: 705647424
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 13, 2024
Imported from GitHub PR openxla/xla#20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad0445f27d7249b4bcebb4ac573188cf50cb0 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59c02732da728fe43bd6c4c12702d070c2c by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59c02732da728fe43bd6c4c12702d070c2c PiperOrigin-RevId: 705647424
copybara-service bot
pushed a commit
that referenced
this pull request
Dec 13, 2024
Imported from GitHub PR #20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad04 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59 by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59 PiperOrigin-RevId: 705647424
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 13, 2024
Imported from GitHub PR openxla/xla#20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad0445f27d7249b4bcebb4ac573188cf50cb0 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59c02732da728fe43bd6c4c12702d070c2c by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59c02732da728fe43bd6c4c12702d070c2c PiperOrigin-RevId: 705647424
thomasjoerg
requested changes
Dec 16, 2024
| if (!params.executor->HostMemoryUnregister(&barrier_flags_[current_id])) { | ||
| LOG(ERROR) << "Unregistering barrier flag failed."; | ||
| } | ||
| } |
There was a problem hiding this comment.
Several TensorFlow tests are failing with:
error: non-void function does not return a value in all control paths [-Werror,-Wreturn-type]
There was a problem hiding this comment.
added a returned status for the cleanup function
copybara-service bot
pushed a commit
that referenced
this pull request
Dec 16, 2024
Imported from GitHub PR #20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad04 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59 by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59 PiperOrigin-RevId: 705647424
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 16, 2024
Imported from GitHub PR openxla/xla#20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad0445f27d7249b4bcebb4ac573188cf50cb0 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59c02732da728fe43bd6c4c12702d070c2c by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59c02732da728fe43bd6c4c12702d070c2c PiperOrigin-RevId: 705647424
copybara-service bot
pushed a commit
that referenced
this pull request
Dec 16, 2024
Imported from GitHub PR #20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad04 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59 by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59 PiperOrigin-RevId: 705647424
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 16, 2024
Imported from GitHub PR openxla/xla#20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad0445f27d7249b4bcebb4ac573188cf50cb0 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59c02732da728fe43bd6c4c12702d070c2c by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59c02732da728fe43bd6c4c12702d070c2c PiperOrigin-RevId: 705647424
copybara-service bot
pushed a commit
that referenced
this pull request
Dec 16, 2024
Imported from GitHub PR #20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad04 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59 by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59 PiperOrigin-RevId: 705647424
copybara-service bot
pushed a commit
that referenced
this pull request
Dec 16, 2024
Imported from GitHub PR #20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad04 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59 by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59 PiperOrigin-RevId: 705647424
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 16, 2024
Imported from GitHub PR openxla/xla#20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad0445f27d7249b4bcebb4ac573188cf50cb0 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59c02732da728fe43bd6c4c12702d070c2c by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59c02732da728fe43bd6c4c12702d070c2c PiperOrigin-RevId: 705647424
thomasjoerg
approved these changes
Dec 17, 2024
copybara-service bot
pushed a commit
that referenced
this pull request
Dec 17, 2024
Imported from GitHub PR #20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad04 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59 by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop -- 1f75328 by TJ Xu <[email protected]>: Added return status for cleanup functions Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=#20086 from Tixxx:tixxx/memcpy_p2p_fix 1f75328 PiperOrigin-RevId: 707074350
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 17, 2024
Imported from GitHub PR openxla/xla#20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad0445f27d7249b4bcebb4ac573188cf50cb0 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59c02732da728fe43bd6c4c12702d070c2c by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop -- 1f7532815dfdbb6d047339d7189c1287dc72e6a3 by TJ Xu <[email protected]>: Added return status for cleanup functions Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#20086 from Tixxx:tixxx/memcpy_p2p_fix 1f7532815dfdbb6d047339d7189c1287dc72e6a3 PiperOrigin-RevId: 707074350
ddunl
approved these changes
Dec 17, 2024
copybara-service bot
pushed a commit
that referenced
this pull request
Dec 17, 2024
Imported from GitHub PR #20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad04 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59 by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59 PiperOrigin-RevId: 705647424
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 17, 2024
Imported from GitHub PR openxla/xla#20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad0445f27d7249b4bcebb4ac573188cf50cb0 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59c02732da728fe43bd6c4c12702d070c2c by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59c02732da728fe43bd6c4c12702d070c2c PiperOrigin-RevId: 705647424
copybara-service bot
pushed a commit
that referenced
this pull request
Dec 17, 2024
Imported from GitHub PR #20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad04 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59 by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop -- 1f75328 by TJ Xu <[email protected]>: Added return status for cleanup functions Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=#20086 from Tixxx:tixxx/memcpy_p2p_fix 1f75328 PiperOrigin-RevId: 707074350
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 17, 2024
Imported from GitHub PR openxla/xla#20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad0445f27d7249b4bcebb4ac573188cf50cb0 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59c02732da728fe43bd6c4c12702d070c2c by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop -- 1f7532815dfdbb6d047339d7189c1287dc72e6a3 by TJ Xu <[email protected]>: Added return status for cleanup functions Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#20086 from Tixxx:tixxx/memcpy_p2p_fix 1f7532815dfdbb6d047339d7189c1287dc72e6a3 PiperOrigin-RevId: 707074350
copybara-service bot
pushed a commit
that referenced
this pull request
Dec 17, 2024
Imported from GitHub PR #20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad04 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59 by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59 PiperOrigin-RevId: 705647424
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 17, 2024
Imported from GitHub PR openxla/xla#20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad0445f27d7249b4bcebb4ac573188cf50cb0 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59c02732da728fe43bd6c4c12702d070c2c by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59c02732da728fe43bd6c4c12702d070c2c PiperOrigin-RevId: 705647424
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 17, 2024
Imported from GitHub PR openxla/xla#20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad0445f27d7249b4bcebb4ac573188cf50cb0 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59c02732da728fe43bd6c4c12702d070c2c by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop -- 1f7532815dfdbb6d047339d7189c1287dc72e6a3 by TJ Xu <[email protected]>: Added return status for cleanup functions Merging this change closes #20086 PiperOrigin-RevId: 707145351
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 17, 2024
…alls. FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#20086 from Tixxx:tixxx/memcpy_p2p_fix 1f7532815dfdbb6d047339d7189c1287dc72e6a3 PiperOrigin-RevId: 690686233
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 17, 2024
Imported from GitHub PR openxla/xla#20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad0445f27d7249b4bcebb4ac573188cf50cb0 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59c02732da728fe43bd6c4c12702d070c2c by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59c02732da728fe43bd6c4c12702d070c2c PiperOrigin-RevId: 705647424
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 18, 2024
Imported from GitHub PR openxla/xla#20086 Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait. Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy. Copybara import of the project: -- ba4ad0445f27d7249b4bcebb4ac573188cf50cb0 by TJ Xu <[email protected]>: Moved pointer init to thunk init stage and add a sync point before doing memcpy to make sure data consistency across ranks -- 050bc59c02732da728fe43bd6c4c12702d070c2c by TJ Xu <[email protected]>: Added e2e test for mem cpy p2p in a loop Merging this change closes #20086 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#20086 from Tixxx:tixxx/memcpy_p2p_fix 050bc59c02732da728fe43bd6c4c12702d070c2c PiperOrigin-RevId: 705647424
|
I'm seeing this regress the Maxtext LLama 7B model with 4-way FSDP and 2-way TP when collective matmul is enabled, using the script you gave me @Tixxx. I see |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Move pointer initialization to the thunk init stage instead of runtime to get rid of the runtime blocking wait.
Add a device sync point using nccl allreduce before doing memcpy to make sure all gpus arrive at the same stage. Otherwise it's possible to have data corruptions when the receiving rank hasn't arrived at the memcpy.