Computing dependencies#

Consider the following piece of code being run on a single machine

importtorcha=torch.rand((3,3),requires_grad=True)b=torch.rand((3,3),requires_grad=True)c=torch.rand((3,3),requires_grad=True)d=a+be=b*cd.sum.().backward()

This is what the autograd graph for the code above would look like:

The first step the autograd engine performs as part of the backward pass iscomputing the number of dependencies for each node in the autograd graph. Thishelps the autograd engine know when a node in the graph is ready for execution.The numbers in brackets foradd(1) andmul(0) denote the number ofdependencies. As you can see, this means during the backward pass theaddnode needs 1 input and themul node doesn’t need any inputs (in otherwords doesn’t need to be executed). The local autograd engine computes thesedependencies by traversing the graph from the root nodes (d in this case).

The fact that certain nodes in the autograd graph might not be executed in thebackward pass poses a challenge for distributed autograd. Consider this pieceof code which uses RPC.

importtorchimporttorch.distributed.rpcasrpca=torch.rand((3,3),requires_grad=True)b=torch.rand((3,3),requires_grad=True)c=torch.rand((3,3),requires_grad=True)d=rpc.rpc_sync("worker1",torch.add,args=(a,b))e=rpc.rpc_sync("worker1",torch.mul,args=(b,c))loss=d.sum()

The associated autograd graph for the code above would be:

Computing dependencies of this distributed autograd graph is much morechallenging and requires some overhead (either in terms of computation ornetwork communication).

For performance sensitive applications we can avoid alot of overhead by assuming everysend andrecv function are valid aspart of the backward pass (most applications don’t perform RPCs that aren’tused). This simplifies the distributed autograd algorithm and is much moreefficient, but at the cost that the application needs to be aware of thelimitations. This algorithm is called theFAST mode algorithm and isdescribed in detail below.

In the general case it might not be necessary that everysend andrecvfunction is valid as part of the backward pass. To address this, we haveproposed aSMART mode algorithm which is described in a later section.Please note that currently, only theFAST mode algorithm is implemented.

FAST mode algorithm#

The key assumption of this algorithm is that eachsend function has adependency of 1 when we run a backward pass. In other words, we assume we’llreceive a gradient over RPC from another node.

The algorithm is as follows:

1. We start from the worker which has the roots for the backward pass(all roots must be local).

2. Lookup all thesend functions for the currentDistributed Autograd Context.

3. Compute dependencies locally starting from the provided roots and all thesend functions we retrieved.

4. After computing dependencies, kick off the local autograd engine with theprovided roots.

5. When the autograd engine executes therecv function, therecvfunction sends the input gradients via RPC to the appropriate worker.Eachrecv function knows the destination worker id since it is recordedas part of the forward pass. Therecv function also sends over theautograd_context_id andautograd_message_id to the remote host.

6. When this request is received on the remote host, we use theautograd_context_id andautograd_message_id to look up theappropriatesend function.

7. If this is the first time a worker has received a request for the givenautograd_context_id, it will compute dependencies locally as describedin points 1-3 above.

8. Thesend function retrieved in 6. is then enqueued for execution on thelocal autograd engine for that worker.

9. Finally, instead of accumulating the gradients on the.grad field of theTensor, we accumulate the gradients separately perDistributed Autograd Context. The gradients are stored in aDict[Tensor,Tensor], which is basically a map from Tensor to itsassociated gradient and this map can be retrieved using theget_gradients() API.

As an example the complete code with distributed autograd would be as follows:

importtorchimporttorch.distributed.autogradasdist_autogradimporttorch.distributed.rpcasrpcdefmy_add(t1,t2):returntorch.add(t1,t2)# On worker 0:# Setup the autograd context. Computations that take# part in the distributed backward pass must be within# the distributed autograd context manager.withdist_autograd.context()ascontext_id:t1=torch.rand((3,3),requires_grad=True)t2=torch.rand((3,3),requires_grad=True)# Perform some computation remotely.t3=rpc.rpc_sync("worker1",my_add,args=(t1,t2))# Perform some computation locally based on remote result.t4=torch.rand((3,3),requires_grad=True)t5=torch.mul(t3,t4)# Compute some loss.loss=t5.sum()# Run the backward pass.dist_autograd.backward(context_id,[loss])# Retrieve the gradients from the context.dist_autograd.get_gradients(context_id)

The distributed autograd graph with dependencies would be as follows (t5.sum() excluded for simplicity):

TheFAST mode algorithm applied to the above example would be as follows:

1. OnWorker0 we start from the rootsloss andsend1 to computedependencies. As a resultsend1 is marked with a dependency of 1 andmulonWorker0 is marked with a dependency of 1.

2. Now, we kickoff the local autograd engine onWorker0. We first executethemul function, accumulate its output in the autograd context as thegradient fort4. Then, we executerecv2 which sends the gradients toWorker1.

3. Since this is the first timeWorker1 has heard about this backward pass,it starts dependency computation and marks the dependencies forsend2,add andrecv1 appropriately.

4. Next, we enqueuesend2 on the local autograd engine ofWorker1, whichin turn executesadd andrecv1.

5. Whenrecv1 is executed it sends the gradients over toWorker0.

6. SinceWorker0 has already computed dependencies for this backward pass,it just enqueues and executessend1 locally.

7. Finally, gradients fort1,t2 andt4 are accumulated in theDistributed Autograd Context.

SMART mode algorithm#

Full details of this algorithm are still in the works, but for the general ideayou can refer toDistributed Autograd Algorithm Smart mode section in theRFC.