TensorPipeRpcAgent "set_device_map" should implicitly use the same order for the backward pass.

[pritamdamania87]

🚀 Feature

TensorPipeRpcAgent should implicitly use the reverse device mapping specified in "set_device_map" for the backward pass.

Motivation

In the current implementation, user's need to specify the forward and backward pass mapping themselves as follows to run the backward pass appropriately:

        # Forward pass.
        options.set_device_map(dst, {self.rank: (self.rank + 1) % self.world_size})
        # Backward pass.
        reverse_rank = (self.rank - 1 + self.world_size) % self.world_size
        options.set_device_map(worker_name(reverse_rank), {self.rank: reverse_rank})

This isn't a good user experience since it would be hard for users to predict which all nodes are involved in a backward pass and how to set up the backward pass device mapping. Instead, it would be ideal if the distributed autograd engine remembers the forward mapping (as part of the send-recv nodes) and uses the inverse of that during the backward pass.

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @rohan-varma @xush6528 @jjlilley @osalpekar @jiayisuse @lw @beauby

[lw]

Interesting issue... Just to clarify: if the forward pass goes from A to B, then the backward pass consists in an RPC request initiated by B towards A, right?

We (actually, @mrshenli) did implement the logic that makes responses use the reverse map of the request, but we didn't think that there may be new requests that we want to behave like responses.

One quick fix we could imagine for 1.7 is the following: if the user specified a device mapping from A to B and another one from B to A, we check that one is the inverse of the other and fail if it's not the case. If the user specified only one of them, then we set the other one to its inverse. If the user didn't specify any, then we use the identity mapping for both directions.

This is a rather strict requirement but I think it's fine if the first iteration is restrictive as we can always relax it later on without breaking backwards compatibility.