You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
…ad. (#53929) (#54358)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/53929
The local autograd engine performs appropriate stream synchronization
between autograd nodes in the graph to ensure a consumer's stream is
synchronized with the producer's stream before executing the consumer.
However in case of distributed autograd, the SendRpcBackward function receives
gradients over the wire and TensorPipe uses its own pool of streams for this
purpose. As a result, the tensors are received on TensorPipe's stream pool but
SendRpcBackward runs on a different stream during the backward pass and there
is no logic to synchronize these streams.
To fix this, I've enhanced DistEngine to synchronize these streams
appropriately when it receives grads over the wire.
ghstack-source-id: 124055277
(Note: this ignores all push blocking failures!)
Test Plan:
1) Added unit test which reproduced the issue.
2) waitforbuildbot.
Reviewed By: walterddr, wanchaol
Differential Revision: D27025307
fbshipit-source-id: 2944854e688e001cb3989d2741727b30d9278414
Co-authored-by: Pritam Damania <[email protected]>