diff --git a/_static/img/thumbnails/cropped/Combining-Distributed-DataParallel-with-Distributed-RPC-Framework.png b/_static/img/thumbnails/cropped/Combining-Distributed-DataParallel-with-Distributed-RPC-Framework.png deleted file mode 100644 index 426a14d98f5..00000000000 Binary files a/_static/img/thumbnails/cropped/Combining-Distributed-DataParallel-with-Distributed-RPC-Framework.png and /dev/null differ diff --git a/_static/img/thumbnails/cropped/Distributed-Pipeline-Parallelism-Using-RPC.png b/_static/img/thumbnails/cropped/Distributed-Pipeline-Parallelism-Using-RPC.png deleted file mode 100644 index 426a14d98f5..00000000000 Binary files a/_static/img/thumbnails/cropped/Distributed-Pipeline-Parallelism-Using-RPC.png and /dev/null differ diff --git a/_static/img/thumbnails/cropped/Getting Started with Distributed-RPC-Framework.png b/_static/img/thumbnails/cropped/Getting Started with Distributed-RPC-Framework.png deleted file mode 100644 index 426a14d98f5..00000000000 Binary files a/_static/img/thumbnails/cropped/Getting Started with Distributed-RPC-Framework.png and /dev/null differ diff --git a/_static/img/thumbnails/cropped/Getting-Started-with Distributed RPC Framework.png b/_static/img/thumbnails/cropped/Getting-Started-with Distributed RPC Framework.png deleted file mode 100644 index 426a14d98f5..00000000000 Binary files a/_static/img/thumbnails/cropped/Getting-Started-with Distributed RPC Framework.png and /dev/null differ diff --git a/_static/img/thumbnails/cropped/Getting-Started-with-Distributed-RPC-Framework.png b/_static/img/thumbnails/cropped/Getting-Started-with-Distributed-RPC-Framework.png deleted file mode 100644 index 426a14d98f5..00000000000 Binary files a/_static/img/thumbnails/cropped/Getting-Started-with-Distributed-RPC-Framework.png and /dev/null differ diff --git a/_static/img/thumbnails/cropped/Implementing-Batch-RPC-Processing-Using-Asynchronous-Executions.png b/_static/img/thumbnails/cropped/Implementing-Batch-RPC-Processing-Using-Asynchronous-Executions.png deleted file mode 100644 index 426a14d98f5..00000000000 Binary files a/_static/img/thumbnails/cropped/Implementing-Batch-RPC-Processing-Using-Asynchronous-Executions.png and /dev/null differ diff --git a/_static/img/thumbnails/cropped/Implementing-a-Parameter-Server-Using-Distributed-RPC-Framework.png b/_static/img/thumbnails/cropped/Implementing-a-Parameter-Server-Using-Distributed-RPC-Framework.png deleted file mode 100644 index 426a14d98f5..00000000000 Binary files a/_static/img/thumbnails/cropped/Implementing-a-Parameter-Server-Using-Distributed-RPC-Framework.png and /dev/null differ diff --git a/advanced_source/rpc_ddp_tutorial.rst b/advanced_source/rpc_ddp_tutorial.rst deleted file mode 100644 index 5c7aeffb2f9..00000000000 --- a/advanced_source/rpc_ddp_tutorial.rst +++ /dev/null @@ -1,162 +0,0 @@ -Combining Distributed DataParallel with Distributed RPC Framework -================================================================= -**Authors**: `Pritam Damania `_ and `Yi Wang `_ - -.. note:: - |edit| View and edit this tutorial in `github `__. - -This tutorial uses a simple example to demonstrate how you can combine -`DistributedDataParallel `__ (DDP) -with the `Distributed RPC framework `__ -to combine distributed data parallelism with distributed model parallelism to -train a simple model. Source code of the example can be found `here `__. - -Previous tutorials, -`Getting Started With Distributed Data Parallel `__ -and `Getting Started with Distributed RPC Framework `__, -described how to perform distributed data parallel and distributed model -parallel training respectively. Although, there are several training paradigms -where you might want to combine these two techniques. For example: - -1) If we have a model with a sparse part (large embedding table) and a dense - part (FC layers), we might want to put the embedding table on a parameter - server and replicate the FC layer across multiple trainers using `DistributedDataParallel `__. - The `Distributed RPC framework `__ - can be used to perform embedding lookups on the parameter server. -2) Enable hybrid parallelism as described in the `PipeDream `__ paper. - We can use the `Distributed RPC framework `__ - to pipeline stages of the model across multiple workers and replicate each - stage (if needed) using `DistributedDataParallel `__. - -| -In this tutorial we will cover case 1 mentioned above. We have a total of 4 -workers in our setup as follows: - - -1) 1 Master, which is responsible for creating an embedding table - (nn.EmbeddingBag) on the parameter server. The master also drives the - training loop on the two trainers. -2) 1 Parameter Server, which basically holds the embedding table in memory and - responds to RPCs from the Master and Trainers. -3) 2 Trainers, which store an FC layer (nn.Linear) which is replicated amongst - themselves using `DistributedDataParallel `__. - The trainers are also responsible for executing the forward pass, backward - pass and optimizer step. - -| -The entire training process is executed as follows: - -1) The master creates a `RemoteModule `__ - that holds an embedding table on the Parameter Server. -2) The master, then kicks off the training loop on the trainers and passes the - remote module to the trainers. -3) The trainers create a ``HybridModel`` which first performs an embedding lookup - using the remote module provided by the master and then executes the - FC layer which is wrapped inside DDP. -4) The trainer executes the forward pass of the model and uses the loss to - execute the backward pass using `Distributed Autograd `__. -5) As part of the backward pass, the gradients for the FC layer are computed - first and synced to all trainers via allreduce in DDP. -6) Next, Distributed Autograd propagates the gradients to the parameter server, - where the gradients for the embedding table are updated. -7) Finally, the `Distributed Optimizer `__ is used to update all the parameters. - - -.. attention:: - - You should always use `Distributed Autograd `__ - for the backward pass if you're combining DDP and RPC. - - -Now, let's go through each part in detail. Firstly, we need to setup all of our -workers before we can perform any training. We create 4 processes such that -ranks 0 and 1 are our trainers, rank 2 is the master and rank 3 is the -parameter server. - -We initialize the RPC framework on all 4 workers using the TCP init_method. -Once RPC initialization is done, the master creates a remote module that holds an `EmbeddingBag `__ -layer on the Parameter Server using `RemoteModule `__. -The master then loops through each trainer and kicks off the training loop by -calling ``_run_trainer`` on each trainer using `rpc_async `__. -Finally, the master waits for all training to finish before exiting. - -The trainers first initialize a ``ProcessGroup`` for DDP with world_size=2 -(for two trainers) using `init_process_group `__. -Next, they initialize the RPC framework using the TCP init_method. Note that -the ports are different in RPC initialization and ProcessGroup initialization. -This is to avoid port conflicts between initialization of both frameworks. -Once the initialization is done, the trainers just wait for the ``_run_trainer`` -RPC from the master. - -The parameter server just initializes the RPC framework and waits for RPCs from -the trainers and master. - - -.. literalinclude:: ../advanced_source/rpc_ddp_tutorial/main.py - :language: py - :start-after: BEGIN run_worker - :end-before: END run_worker - -Before we discuss details of the Trainer, let's introduce the ``HybridModel`` that -the trainer uses. As described below, the ``HybridModel`` is initialized using a -remote module that holds an embedding table (``remote_emb_module``) on the parameter server and the ``device`` -to use for DDP. The initialization of the model wraps an -`nn.Linear `__ -layer inside DDP to replicate and synchronize this layer across all trainers. - -The forward method of the model is pretty straightforward. It performs an -embedding lookup on the parameter server using RemoteModule's ``forward`` -and passes its output onto the FC layer. - - -.. literalinclude:: ../advanced_source/rpc_ddp_tutorial/main.py - :language: py - :start-after: BEGIN hybrid_model - :end-before: END hybrid_model - -Next, let's look at the setup on the Trainer. The trainer first creates the -``HybridModel`` described above using a remote module that holds the embedding table on the -parameter server and its own rank. - -Now, we need to retrieve a list of RRefs to all the parameters that we would -like to optimize with `DistributedOptimizer `__. -To retrieve the parameters for the embedding table from the parameter server, -we can call RemoteModule's `remote_parameters `__, -which basically walks through all the parameters for the embedding table and returns -a list of RRefs. The trainer calls this method on the parameter server via RPC -to receive a list of RRefs to the desired parameters. Since the -DistributedOptimizer always takes a list of RRefs to parameters that need to -be optimized, we need to create RRefs even for the local parameters for our -FC layers. This is done by walking ``model.fc.parameters()``, creating an RRef for -each parameter and appending it to the list returned from ``remote_parameters()``. -Note that we cannnot use ``model.parameters()``, -because it will recursively call ``model.remote_emb_module.parameters()``, -which is not supported by ``RemoteModule``. - -Finally, we create our DistributedOptimizer using all the RRefs and define a -CrossEntropyLoss function. - -.. literalinclude:: ../advanced_source/rpc_ddp_tutorial/main.py - :language: py - :start-after: BEGIN setup_trainer - :end-before: END setup_trainer - -Now we're ready to introduce the main training loop that is run on each trainer. -``get_next_batch`` is just a helper function to generate random inputs and -targets for training. We run the training loop for multiple epochs and for each -batch: - -1) Setup a `Distributed Autograd Context `__ - for Distributed Autograd. -2) Run the forward pass of the model and retrieve its output. -3) Compute the loss based on our outputs and targets using the loss function. -4) Use Distributed Autograd to execute a distributed backward pass using the loss. -5) Finally, run a Distributed Optimizer step to optimize all the parameters. - -.. literalinclude:: ../advanced_source/rpc_ddp_tutorial/main.py - :language: py - :start-after: BEGIN run_trainer - :end-before: END run_trainer -.. code:: python - -Source code for the entire example can be found `here `__. diff --git a/advanced_source/rpc_ddp_tutorial/main.py b/advanced_source/rpc_ddp_tutorial/main.py deleted file mode 100644 index 459741ff0f4..00000000000 --- a/advanced_source/rpc_ddp_tutorial/main.py +++ /dev/null @@ -1,185 +0,0 @@ -import random - -import torch -import torch.distributed as dist -import torch.distributed.autograd as dist_autograd -import torch.distributed.rpc as rpc -import torch.multiprocessing as mp -import torch.optim as optim -from torch.distributed.nn import RemoteModule -from torch.distributed.optim import DistributedOptimizer -from torch.distributed.rpc import RRef -from torch.distributed.rpc import TensorPipeRpcBackendOptions -from torch.nn.parallel import DistributedDataParallel as DDP - -NUM_EMBEDDINGS = 100 -EMBEDDING_DIM = 16 - -# BEGIN hybrid_model -class HybridModel(torch.nn.Module): - r""" - The model consists of a sparse part and a dense part. - 1) The dense part is an nn.Linear module that is replicated across all trainers using DistributedDataParallel. - 2) The sparse part is a Remote Module that holds an nn.EmbeddingBag on the parameter server. - This remote model can get a Remote Reference to the embedding table on the parameter server. - """ - - def __init__(self, remote_emb_module, device): - super(HybridModel, self).__init__() - self.remote_emb_module = remote_emb_module - self.fc = DDP(torch.nn.Linear(16, 8).cuda(device), device_ids=[device]) - self.device = device - - def forward(self, indices, offsets): - emb_lookup = self.remote_emb_module.forward(indices, offsets) - return self.fc(emb_lookup.cuda(self.device)) -# END hybrid_model - -# BEGIN setup_trainer -def _run_trainer(remote_emb_module, rank): - r""" - Each trainer runs a forward pass which involves an embedding lookup on the - parameter server and running nn.Linear locally. During the backward pass, - DDP is responsible for aggregating the gradients for the dense part - (nn.Linear) and distributed autograd ensures gradients updates are - propagated to the parameter server. - """ - - # Setup the model. - model = HybridModel(remote_emb_module, rank) - - # Retrieve all model parameters as rrefs for DistributedOptimizer. - - # Retrieve parameters for embedding table. - model_parameter_rrefs = model.remote_emb_module.remote_parameters() - - # model.fc.parameters() only includes local parameters. - # NOTE: Cannot call model.parameters() here, - # because this will call remote_emb_module.parameters(), - # which supports remote_parameters() but not parameters(). - for param in model.fc.parameters(): - model_parameter_rrefs.append(RRef(param)) - - # Setup distributed optimizer - opt = DistributedOptimizer( - optim.SGD, - model_parameter_rrefs, - lr=0.05, - ) - - criterion = torch.nn.CrossEntropyLoss() - # END setup_trainer - - # BEGIN run_trainer - def get_next_batch(rank): - for _ in range(10): - num_indices = random.randint(20, 50) - indices = torch.LongTensor(num_indices).random_(0, NUM_EMBEDDINGS) - - # Generate offsets. - offsets = [] - start = 0 - batch_size = 0 - while start < num_indices: - offsets.append(start) - start += random.randint(1, 10) - batch_size += 1 - - offsets_tensor = torch.LongTensor(offsets) - target = torch.LongTensor(batch_size).random_(8).cuda(rank) - yield indices, offsets_tensor, target - - # Train for 100 epochs - for epoch in range(100): - # create distributed autograd context - for indices, offsets, target in get_next_batch(rank): - with dist_autograd.context() as context_id: - output = model(indices, offsets) - loss = criterion(output, target) - - # Run distributed backward pass - dist_autograd.backward(context_id, [loss]) - - # Tun distributed optimizer - opt.step(context_id) - - # Not necessary to zero grads as each iteration creates a different - # distributed autograd context which hosts different grads - print("Training done for epoch {}".format(epoch)) - # END run_trainer - -# BEGIN run_worker -def run_worker(rank, world_size): - r""" - A wrapper function that initializes RPC, calls the function, and shuts down - RPC. - """ - - # We need to use different port numbers in TCP init_method for init_rpc and - # init_process_group to avoid port conflicts. - rpc_backend_options = TensorPipeRpcBackendOptions() - rpc_backend_options.init_method = "tcp://localhost:29501" - - # Rank 2 is master, 3 is ps and 0 and 1 are trainers. - if rank == 2: - rpc.init_rpc( - "master", - rank=rank, - world_size=world_size, - rpc_backend_options=rpc_backend_options, - ) - - remote_emb_module = RemoteModule( - "ps", - torch.nn.EmbeddingBag, - args=(NUM_EMBEDDINGS, EMBEDDING_DIM), - kwargs={"mode": "sum"}, - ) - - # Run the training loop on trainers. - futs = [] - for trainer_rank in [0, 1]: - trainer_name = "trainer{}".format(trainer_rank) - fut = rpc.rpc_async( - trainer_name, _run_trainer, args=(remote_emb_module, trainer_rank) - ) - futs.append(fut) - - # Wait for all training to finish. - for fut in futs: - fut.wait() - elif rank <= 1: - # Initialize process group for Distributed DataParallel on trainers. - dist.init_process_group( - backend="gloo", rank=rank, world_size=2, init_method="tcp://localhost:29500" - ) - - # Initialize RPC. - trainer_name = "trainer{}".format(rank) - rpc.init_rpc( - trainer_name, - rank=rank, - world_size=world_size, - rpc_backend_options=rpc_backend_options, - ) - - # Trainer just waits for RPCs from master. - else: - rpc.init_rpc( - "ps", - rank=rank, - world_size=world_size, - rpc_backend_options=rpc_backend_options, - ) - # parameter server do nothing - pass - - # block until all rpcs finish - rpc.shutdown() - - -if __name__ == "__main__": - # 2 trainers, 1 parameter server, 1 master. - world_size = 4 - mp.spawn(run_worker, args=(world_size,), nprocs=world_size, join=True) -# END run_worker diff --git a/beginner_source/profiler.py b/beginner_source/profiler.py index 6a3894cbbd3..547740304f5 100644 --- a/beginner_source/profiler.py +++ b/beginner_source/profiler.py @@ -317,6 +317,5 @@ def forward(self, input, mask): # Read more about Profiler here: # # - `Profiler Usage Recipe `__ -# - `Profiling RPC-Based Workloads `__ # - `Profiler API Docs `__ # - `Hugging Face PyTorch Profiler Blog Post `__ diff --git a/distributed.rst b/distributed.rst index 6ef33db8889..c77cfb01e0d 100644 --- a/distributed.rst +++ b/distributed.rst @@ -15,7 +15,6 @@ PyTorch with each method having their advantages in certain use cases: * `Fully Sharded Data Parallel (FSDP2) <#learn-fsdp>`__ * `Tensor Parallel (TP) <#learn-tp>`__ * `Device Mesh <#device-mesh>`__ -* `Remote Procedure Call (RPC) distributed training <#learn-rpc>`__ * `Monarch Framework <#learn-monarch>`__ * `Custom Extensions <#custom-extensions>`__ @@ -111,55 +110,6 @@ Learn DeviceMesh +++ :octicon:`code;1em` Code -.. _learn-rpc: - -Learn RPC ---------- - -.. grid:: 3 - - .. grid-item-card:: :octicon:`file-code;1em` - Getting Started with Distributed RPC Framework - :link: https://pytorch.org/tutorials/intermediate/rpc_tutorial.html?utm_source=distr_landing&utm_medium=rpc_getting_started - :link-type: url - - This tutorial demonstrates how to get started with RPC-based distributed - training. - +++ - :octicon:`code;1em` Code - - .. grid-item-card:: :octicon:`file-code;1em` - Implementing a Parameter Server Using Distributed RPC Framework - :link: https://pytorch.org/tutorials/intermediate/rpc_param_server_tutorial.html?utm_source=distr_landing&utm_medium=rpc_param_server_tutorial - :link-type: url - - This tutorial walks you through a simple example of implementing a - parameter server using PyTorch’s Distributed RPC framework. - +++ - :octicon:`code;1em` Code - - .. grid-item-card:: :octicon:`file-code;1em` - Implementing Batch RPC Processing Using Asynchronous Executions - :link: https://pytorch.org/tutorials/intermediate/rpc_async_execution.html?utm_source=distr_landing&utm_medium=rpc_async_execution - :link-type: url - - In this tutorial you will build batch-processing RPC applications - with the @rpc.functions.async_execution decorator. - +++ - :octicon:`code;1em` Code - -.. grid:: 3 - - .. grid-item-card:: :octicon:`file-code;1em` - Combining Distributed DataParallel with Distributed RPC Framework - :link: https://pytorch.org/tutorials/advanced/rpc_ddp_tutorial.html?utm_source=distr_landing&utm_medium=rpc_plus_ddp - :link-type: url - - In this tutorial you will learn how to combine distributed data - parallelism with distributed model parallelism. - +++ - :octicon:`code;1em` Code - .. _learn-monarch: Learn Monarch @@ -207,11 +157,7 @@ Custom Extensions intermediate/TP_tutorial intermediate/pipelining_tutorial intermediate/process_group_cpp_extension_tutorial - intermediate/rpc_tutorial - intermediate/rpc_param_server_tutorial - intermediate/rpc_async_execution intermediate/monarch_distributed_tutorial intermediate/debug_hangs_with_flight_recorder - advanced/rpc_ddp_tutorial advanced/generic_join beginner/distributed_training_with_ray_tutorial diff --git a/en-wordlist.txt b/en-wordlist.txt index baf75d75ac0..3e9e58c71b6 100644 --- a/en-wordlist.txt +++ b/en-wordlist.txt @@ -215,7 +215,6 @@ RNNs ROCm RPC RPN -RRef RTX Radford ReLU diff --git a/index.rst b/index.rst index 8dbc3e79eca..475f9a67157 100644 --- a/index.rst +++ b/index.rst @@ -667,20 +667,6 @@ Welcome to PyTorch Tutorials :link: intermediate/process_group_cpp_extension_tutorial.html :tags: Parallel-and-Distributed-Training -.. customcarditem:: - :header: Getting Started with Distributed RPC Framework - :card_description: Learn how to build distributed training using the torch.distributed.rpc package. - :image: _static/img/thumbnails/cropped/Getting Started with Distributed-RPC-Framework.png - :link: intermediate/rpc_tutorial.html - :tags: Parallel-and-Distributed-Training - -.. customcarditem:: - :header: Implementing a Parameter Server Using Distributed RPC Framework - :card_description: Walk through a through a simple example of implementing a parameter server using PyTorch’s Distributed RPC framework. - :image: _static/img/thumbnails/cropped/Implementing-a-Parameter-Server-Using-Distributed-RPC-Framework.png - :link: intermediate/rpc_param_server_tutorial.html - :tags: Parallel-and-Distributed-Training - .. customcarditem:: :header: Introduction to Distributed Pipeline Parallelism :card_description: Demonstrate how to implement pipeline parallelism using torch.distributed.pipelining @@ -688,20 +674,6 @@ Welcome to PyTorch Tutorials :link: intermediate/pipelining_tutorial.html :tags: Parallel-and-Distributed-Training -.. customcarditem:: - :header: Implementing Batch RPC Processing Using Asynchronous Executions - :card_description: Learn how to use rpc.functions.async_execution to implement batch RPC - :image: _static/img/thumbnails/cropped/Implementing-Batch-RPC-Processing-Using-Asynchronous-Executions.png - :link: intermediate/rpc_async_execution.html - :tags: Parallel-and-Distributed-Training - -.. customcarditem:: - :header: Combining Distributed DataParallel with Distributed RPC Framework - :card_description: Walk through a through a simple example of how to combine distributed data parallelism with distributed model parallelism. - :image: _static/img/thumbnails/cropped/Combining-Distributed-DataParallel-with-Distributed-RPC-Framework.png - :link: advanced/rpc_ddp_tutorial.html - :tags: Parallel-and-Distributed-Training - .. customcarditem:: :header: Getting Started with Fully Sharded Data Parallel (FSDP2) :card_description: Learn how to train models with Fully Sharded Data Parallel (fully_shard) package. diff --git a/intermediate_source/rpc_async_execution.rst b/intermediate_source/rpc_async_execution.rst deleted file mode 100644 index 4c7739104cc..00000000000 --- a/intermediate_source/rpc_async_execution.rst +++ /dev/null @@ -1,523 +0,0 @@ -Implementing Batch RPC Processing Using Asynchronous Executions -=============================================================== -**Author**: `Shen Li `_ - -.. note:: - |edit| View and edit this tutorial in `github `__. - -Prerequisites: - -- `PyTorch Distributed Overview <../beginner/dist_overview.html>`__ -- `Getting started with Distributed RPC Framework `__ -- `Implementing a Parameter Server using Distributed RPC Framework `__ -- `RPC Asynchronous Execution Decorator `__ - -This tutorial demonstrates how to build batch-processing RPC applications with -the `@rpc.functions.async_execution `__ -decorator, which helps to speed up training by reducing the number of blocked -RPC threads and consolidating CUDA operations on the callee. - -.. note:: This tutorial requires PyTorch v1.6.0 or above. - -Basics ------- - -Previous tutorials have shown the steps to build distributed training -applications using `torch.distributed.rpc `__, -but they didn't elaborate on what happens on the callee side when processing an -RPC request. As of PyTorch v1.5, each RPC request will block one thread on the -callee to execute the function in that request until that function returns. -This works for many use cases, but there is one caveat. If the user function -blocks on IO, e.g., with nested RPC invocation, or signaling, e.g., waiting for -a different RPC request to unblock, the RPC thread on the callee will have to -idle waiting until the IO finishes or the signaling event occurs. As a result, -RPC callees are likely to use more threads than necessary. The cause of this -problem is that RPC treats user functions as black boxes, and knows very little -about what happens in the function. To allow user functions to yield and free -RPC threads, more hints need to be provided to the RPC system. - -Since v1.6.0, PyTorch addresses this problem by introducing two new concepts: - -* A `torch.futures.Future `__ type - that encapsulates an asynchronous execution, which also supports installing - callback functions. -* An `@rpc.functions.async_execution `__ - decorator that allows applications to tell the callee that the target function - will return a future and can pause and yield multiple times during execution. - -With these two tools, the application code can break a user function into -multiple smaller functions, chain them together as callbacks on ``Future`` -objects, and return the ``Future`` that contains the final result. On the callee -side, when getting the ``Future`` object, it installs subsequent RPC response -preparation and communication as callbacks as well, which will be triggered -when the final result is ready. In this way, the callee no longer needs to block -one thread and wait until the final return value is ready. Please refer to the -API doc of -`@rpc.functions.async_execution `__ -for simple examples. - -Besides reducing the number of idle threads on the callee, these tools also help -to make batch RPC processing easier and faster. The following two sections of -this tutorial demonstrate how to build distributed batch-updating parameter -server and batch-processing reinforcement learning applications using the -`@rpc.functions.async_execution `__ -decorator. - -Batch-Updating Parameter Server -------------------------------- - -Consider a synchronized parameter server training application with one parameter -server (PS) and multiple trainers. In this application, the PS holds the -parameters and waits for all trainers to report gradients. In every iteration, -it waits until receiving gradients from all trainers and then updates all -parameters in one shot. The code below shows the implementation of the PS class. -The ``update_and_fetch_model`` method is decorated using -``@rpc.functions.async_execution`` and will be called by trainers. Each -invocation returns a ``Future`` object that will be populated with the updated -model. Invocations launched by most trainers just accumulate gradients to the -``.grad`` field, return immediately, and yield the RPC thread on the PS. The -last arriving trainer will trigger the optimizer step and consume all previously -reported gradients. Then it sets the ``future_model`` with the updated model, -which in turn notifies all previous requests from other trainers through the -``Future`` object and sends out the updated model to all trainers. - -.. code:: python - - import threading - import torchvision - import torch - import torch.distributed.rpc as rpc - from torch import optim - - num_classes, batch_update_size = 30, 5 - - class BatchUpdateParameterServer(object): - def __init__(self, batch_update_size=batch_update_size): - self.model = torchvision.models.resnet50(num_classes=num_classes) - self.lock = threading.Lock() - self.future_model = torch.futures.Future() - self.batch_update_size = batch_update_size - self.curr_update_size = 0 - self.optimizer = optim.SGD(self.model.parameters(), lr=0.001, momentum=0.9) - for p in self.model.parameters(): - p.grad = torch.zeros_like(p) - - def get_model(self): - return self.model - - @staticmethod - @rpc.functions.async_execution - def update_and_fetch_model(ps_rref, grads): - # Using the RRef to retrieve the local PS instance - self = ps_rref.local_value() - with self.lock: - self.curr_update_size += 1 - # accumulate gradients into .grad field - for p, g in zip(self.model.parameters(), grads): - p.grad += g - - # Save the current future_model and return it to make sure the - # returned Future object holds the correct model even if another - # thread modifies future_model before this thread returns. - fut = self.future_model - - if self.curr_update_size >= self.batch_update_size: - # update the model - for p in self.model.parameters(): - p.grad /= self.batch_update_size - self.curr_update_size = 0 - self.optimizer.step() - self.optimizer.zero_grad() - # by settiing the result on the Future object, all previous - # requests expecting this updated model will be notified and - # the their responses will be sent accordingly. - fut.set_result(self.model) - self.future_model = torch.futures.Future() - - return fut - -For the trainers, they are all initialized using the same set of -parameters from the PS. In every iteration, each trainer first runs the forward -and the backward passes to generate gradients locally. Then, each trainer -reports its gradients to the PS using RPC, and fetches back the updated -parameters through the return value of the same RPC request. In the trainer's -implementation, whether the target function is marked with -``@rpc.functions.async_execution`` or not makes no difference. The -trainer simply calls ``update_and_fetch_model`` using ``rpc_sync`` which will -block on the trainer until the updated model is returned. - -.. code:: python - - batch_size, image_w, image_h = 20, 64, 64 - - class Trainer(object): - def __init__(self, ps_rref): - self.ps_rref, self.loss_fn = ps_rref, torch.nn.MSELoss() - self.one_hot_indices = torch.LongTensor(batch_size) \ - .random_(0, num_classes) \ - .view(batch_size, 1) - - def get_next_batch(self): - for _ in range(6): - inputs = torch.randn(batch_size, 3, image_w, image_h) - labels = torch.zeros(batch_size, num_classes) \ - .scatter_(1, self.one_hot_indices, 1) - yield inputs.cuda(), labels.cuda() - - def train(self): - name = rpc.get_worker_info().name - # get initial model parameters - m = self.ps_rref.rpc_sync().get_model().cuda() - # start training - for inputs, labels in self.get_next_batch(): - self.loss_fn(m(inputs), labels).backward() - m = rpc.rpc_sync( - self.ps_rref.owner(), - BatchUpdateParameterServer.update_and_fetch_model, - args=(self.ps_rref, [p.grad for p in m.cpu().parameters()]), - ).cuda() - -We skip the code that launches multiple processes in this tutorial and please -refer to the `examples `__ -repo for the full implementation. Note that, it is possible to implement batch -processing without the -`@rpc.functions.async_execution `__ -decorator. However, that would require either blocking more RPC threads on -the PS or use another round of RPC to fetch updated models, where the latter -would add both more code complexity and more communication overhead. - -This section uses a simple parameter sever training example to show how to -implement batch RPC applications using the -`@rpc.functions.async_execution `__ -decorator. In the next section, we re-implement the reinforcement learning -example in the previous -`Getting started with Distributed RPC Framework `__ -tutorial using batch processing, and demonstrate its impact on the training -speed. - -Batch-Processing CartPole Solver --------------------------------- - -This section uses CartPole-v1 from OpenAI Gym as -an example to show the performance impact of batch processing RPC. Please note -that since the goal is to demonstrate the usage of -`@rpc.functions.async_execution `__ -instead of building the best CartPole solver or solving most different RL -problems, we use very simple policies and reward calculation strategies and -focus on the multi-observer single-agent batch RPC implementation. We use a -similar ``Policy`` model as the previous tutorial which is shown below. Compared -to the previous tutorial, the difference is that its constructor takes an -additional ``batch`` argument which controls the ``dim`` parameter for -``F.softmax`` because with batching, the ``x`` argument in the ``forward`` -function contains states from multiple observers and hence the dimension needs -to change properly. Everything else stays intact. - -.. code:: python - - import argparse - import torch.nn as nn - import torch.nn.functional as F - - parser = argparse.ArgumentParser(description='PyTorch RPC Batch RL example') - parser.add_argument('--gamma', type=float, default=1.0, metavar='G', - help='discount factor (default: 1.0)') - parser.add_argument('--seed', type=int, default=543, metavar='N', - help='random seed (default: 543)') - parser.add_argument('--num-episode', type=int, default=10, metavar='E', - help='number of episodes (default: 10)') - args = parser.parse_args() - - torch.manual_seed(args.seed) - - class Policy(nn.Module): - def __init__(self, batch=True): - super(Policy, self).__init__() - self.affine1 = nn.Linear(4, 128) - self.dropout = nn.Dropout(p=0.6) - self.affine2 = nn.Linear(128, 2) - self.dim = 2 if batch else 1 - - def forward(self, x): - x = self.affine1(x) - x = self.dropout(x) - x = F.relu(x) - action_scores = self.affine2(x) - return F.softmax(action_scores, dim=self.dim) - - -The constructor of the ``Observer`` adjusts accordingly as well. It also takes a -``batch`` argument, which governs which ``Agent`` function it uses to select -actions. In batch mode, it calls ``select_action_batch`` function on ``Agent`` -which will be presented shortly, and this function will be decorated with -`@rpc.functions.async_execution `__. - - -.. code:: python - - import gym - import torch.distributed.rpc as rpc - - class Observer: - def __init__(self, batch=True): - self.id = rpc.get_worker_info().id - 1 - self.env = gym.make('CartPole-v1') - self.env.seed(args.seed) - self.select_action = Agent.select_action_batch if batch else Agent.select_action - -Compared to the previous tutorial -`Getting started with Distributed RPC Framework `__, -observers behave a little differently. Instead of exiting when the environment -is stopped, it always runs ``n_steps`` iterations in every episode. When the -environment returns, the observer simply resets the environment and start over -again. With this design, the agent will receive a fixed number of states from -every observer and hence can pack them into a fixed-size tensor. In every -step, the ``Observer`` uses RPC to send its state to the ``Agent`` and fetches -the action through the return value. At the end of every episode, it returns the -rewards of all steps to ``Agent``. Note that this ``run_episode`` function will -be called by the ``Agent`` using RPC. So the ``rpc_sync`` call in this function -will be a nested RPC invocation. We could mark this function as ``@rpc.functions.async_execution`` -too to avoid blocking one thread on the ``Observer``. However, as the bottleneck -is the ``Agent`` instead of the ``Observer``, it should be OK to block one -thread on the ``Observer`` process. - - -.. code:: python - - import torch - - class Observer: - ... - - def run_episode(self, agent_rref, n_steps): - state, ep_reward = self.env.reset(), NUM_STEPS - rewards = torch.zeros(n_steps) - start_step = 0 - for step in range(n_steps): - state = torch.from_numpy(state).float().unsqueeze(0) - # send the state to the agent to get an action - action = rpc.rpc_sync( - agent_rref.owner(), - self.select_action, - args=(agent_rref, self.id, state) - ) - - # apply the action to the environment, and get the reward - state, reward, done, _ = self.env.step(action) - rewards[step] = reward - - if done or step + 1 >= n_steps: - curr_rewards = rewards[start_step:(step + 1)] - R = 0 - for i in range(curr_rewards.numel() -1, -1, -1): - R = curr_rewards[i] + args.gamma * R - curr_rewards[i] = R - state = self.env.reset() - if start_step == 0: - ep_reward = min(ep_reward, step - start_step + 1) - start_step = step + 1 - - return [rewards, ep_reward] - -The constructor of the ``Agent`` also takes a ``batch`` argument, which controls -how action probs are batched. In batch mode, the ``saved_log_probs`` contains a -list of tensors, where each tensor contains action robs from all observers in -one step. Without batching, the ``saved_log_probs`` is a dictionary where the -key is the observer id and the value is a list of action probs for that -observer. - -.. code:: python - - import threading - from torch.distributed.rpc import RRef - - class Agent: - def __init__(self, world_size, batch=True): - self.ob_rrefs = [] - self.agent_rref = RRef(self) - self.rewards = {} - self.policy = Policy(batch).cuda() - self.optimizer = optim.Adam(self.policy.parameters(), lr=1e-2) - self.running_reward = 0 - - for ob_rank in range(1, world_size): - ob_info = rpc.get_worker_info(OBSERVER_NAME.format(ob_rank)) - self.ob_rrefs.append(rpc.remote(ob_info, Observer, args=(batch,))) - self.rewards[ob_info.id] = [] - - self.states = torch.zeros(len(self.ob_rrefs), 1, 4) - self.batch = batch - self.saved_log_probs = [] if batch else {k:[] for k in range(len(self.ob_rrefs))} - self.future_actions = torch.futures.Future() - self.lock = threading.Lock() - self.pending_states = len(self.ob_rrefs) - -The non-batching ``select_acion`` simply runs the state throw the policy, saves -the action prob, and returns the action to the observer right away. - -.. code:: python - - from torch.distributions import Categorical - - class Agent: - ... - - @staticmethod - def select_action(agent_rref, ob_id, state): - self = agent_rref.local_value() - probs = self.policy(state.cuda()) - m = Categorical(probs) - action = m.sample() - self.saved_log_probs[ob_id].append(m.log_prob(action)) - return action.item() - -With batching, the state is stored in a 2D tensor ``self.states``, using the -observer id as the row id. Then, it chains a ``Future`` by installing a callback -function to the batch-generated ``self.future_actions`` ``Future`` object, which -will be populated with the specific row indexed using the id of that observer. -The last arriving observer runs all batched states through the policy in one -shot and set ``self.future_actions`` accordingly. When this occurs, all the -callback functions installed on ``self.future_actions`` will be triggered and -their return values will be used to populate the chained ``Future`` object, -which in turn notifies the ``Agent`` to prepare and communicate responses for -all previous RPC requests from other observers. - -.. code:: python - - class Agent: - ... - - @staticmethod - @rpc.functions.async_execution - def select_action_batch(agent_rref, ob_id, state): - self = agent_rref.local_value() - self.states[ob_id].copy_(state) - future_action = self.future_actions.then( - lambda future_actions: future_actions.wait()[ob_id].item() - ) - - with self.lock: - self.pending_states -= 1 - if self.pending_states == 0: - self.pending_states = len(self.ob_rrefs) - probs = self.policy(self.states.cuda()) - m = Categorical(probs) - actions = m.sample() - self.saved_log_probs.append(m.log_prob(actions).t()[0]) - future_actions = self.future_actions - self.future_actions = torch.futures.Future() - future_actions.set_result(actions.cpu()) - return future_action - -Now let's define how different RPC functions are stitched together. The ``Agent`` -controls the execution of every episode. It first uses ``rpc_async`` to kick off -the episode on all observers and block on the returned futures which will be -populated with observer rewards. Note that the code below uses the RRef helper -``ob_rref.rpc_async()`` to launch the ``run_episode`` function on the owner -of the ``ob_rref`` RRef with the provided arguments. -It then converts the saved action probs and returned observer rewards into -expected data format, and launch the training step. Finally, it resets all -states and returns the reward of the current episode. This function is the entry -point to run one episode. - -.. code:: python - - class Agent: - ... - - def run_episode(self, n_steps=0): - futs = [] - for ob_rref in self.ob_rrefs: - # make async RPC to kick off an episode on all observers - futs.append(ob_rref.rpc_async().run_episode(self.agent_rref, n_steps)) - - # wait until all obervers have finished this episode - rets = torch.futures.wait_all(futs) - rewards = torch.stack([ret[0] for ret in rets]).cuda().t() - ep_rewards = sum([ret[1] for ret in rets]) / len(rets) - - # stack saved probs into one tensor - if self.batch: - probs = torch.stack(self.saved_log_probs) - else: - probs = [torch.stack(self.saved_log_probs[i]) for i in range(len(rets))] - probs = torch.stack(probs) - - policy_loss = -probs * rewards / len(rets) - policy_loss.sum().backward() - self.optimizer.step() - self.optimizer.zero_grad() - - # reset variables - self.saved_log_probs = [] if self.batch else {k:[] for k in range(len(self.ob_rrefs))} - self.states = torch.zeros(len(self.ob_rrefs), 1, 4) - - # calculate running rewards - self.running_reward = 0.5 * ep_rewards + 0.5 * self.running_reward - return ep_rewards, self.running_reward - -The rest of the code is normal processes launching and logging which are -similar to other RPC tutorials. In this tutorial, all observers passively -waiting for commands from the agent. Please refer to the -`examples `__ -repo for the full implementation. - -.. code:: python - - def run_worker(rank, world_size, n_episode, batch, print_log=True): - os.environ['MASTER_ADDR'] = 'localhost' - os.environ['MASTER_PORT'] = '29500' - if rank == 0: - # rank0 is the agent - rpc.init_rpc(AGENT_NAME, rank=rank, world_size=world_size) - - agent = Agent(world_size, batch) - for i_episode in range(n_episode): - last_reward, running_reward = agent.run_episode(n_steps=NUM_STEPS) - - if print_log: - print('Episode {}\tLast reward: {:.2f}\tAverage reward: {:.2f}'.format( - i_episode, last_reward, running_reward)) - else: - # other ranks are the observer - rpc.init_rpc(OBSERVER_NAME.format(rank), rank=rank, world_size=world_size) - # observers passively waiting for instructions from agents - rpc.shutdown() - - - def main(): - for world_size in range(2, 12): - delays = [] - for batch in [True, False]: - tik = time.time() - mp.spawn( - run_worker, - args=(world_size, args.num_episode, batch), - nprocs=world_size, - join=True - ) - tok = time.time() - delays.append(tok - tik) - - print(f"{world_size}, {delays[0]}, {delays[1]}") - - - if __name__ == '__main__': - main() - -Batch RPC helps to consolidate the action inference into less CUDA operations, -and hence reduces the amortized overhead. The above ``main`` function runs the -same code on both batch and no-batch modes using different numbers of observers, -ranging from 1 to 10. The figure below plots the execution time of different -world sizes using default argument values. The results confirmed our expectation -that batch processing helped to speed up training. - - -.. figure:: /_static/img/rpc-images/batch.png - :alt: - -Learn More ----------- - -- `Batch-Updating Parameter Server Source Code `__ -- `Batch-Processing CartPole Solver `__ -- `Distributed Autograd `__ diff --git a/intermediate_source/rpc_param_server_tutorial.rst b/intermediate_source/rpc_param_server_tutorial.rst deleted file mode 100644 index 324331646c5..00000000000 --- a/intermediate_source/rpc_param_server_tutorial.rst +++ /dev/null @@ -1,386 +0,0 @@ - -Implementing a Parameter Server Using Distributed RPC Framework -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -**Author**\ : `Rohan Varma `_ - -.. note:: - |edit| View and edit this tutorial in `github `__. - -Prerequisites: - -- `PyTorch Distributed Overview <../beginner/dist_overview.html>`__ -- `RPC API documents `__ - -This tutorial walks through a simple example of implementing a parameter server using PyTorch's `Distributed RPC framework `_. The parameter server framework is a paradigm in which a set of servers store parameters, such as large embedding tables, and several trainers query the parameter servers in order to retrieve the most up to date parameters. These trainers can run a training loop locally and occasionally synchronize with the parameter server to get the latest parameters. For more reading on the parameter server approach, check out `this paper `_. - -Using the Distributed RPC Framework, we'll build an example where multiple trainers use RPC to communicate with the same parameter server and use `RRef `_ to access states on the remote parameter server instance. Each trainer will launch its dedicated backward pass in a distributed fashion through stitching of the autograd graph across multiple nodes using distributed autograd. - -**Note**\ : This tutorial covers the use of the Distributed RPC Framework, which is useful for splitting a model onto multiple machines, or for implementing a parameter-server training strategy where network trainers fetch parameters hosted on a different machine. If instead you are looking for replicating your model across many GPUs, please see the `Distributed Data Parallel tutorial `_. There is also another `RPC tutorial `_ that covers reinforcement learning and RNN use cases. - -Let's start with the familiar: importing our required modules and defining a simple ConvNet that will train on the MNIST dataset. The below network is largely adopted from the network defined in the `pytorch/examples repo `_. - -.. code-block:: python - - import argparse - import os - import time - from threading import Lock - - import torch - import torch.distributed.autograd as dist_autograd - import torch.distributed.rpc as rpc - import torch.multiprocessing as mp - import torch.nn as nn - import torch.nn.functional as F - from torch import optim - from torch.distributed.optim import DistributedOptimizer - from torchvision import datasets, transforms - - # --------- MNIST Network to train, from pytorch/examples ----- - - class Net(nn.Module): - def __init__(self, num_gpus=0): - super(Net, self).__init__() - print(f"Using {num_gpus} GPUs to train") - self.num_gpus = num_gpus - device = torch.device( - "cuda:0" if torch.cuda.is_available() and self.num_gpus > 0 else "cpu") - print(f"Putting first 2 convs on {str(device)}") - # Put conv layers on the first cuda device, or CPU if no cuda device - self.conv1 = nn.Conv2d(1, 32, 3, 1).to(device) - self.conv2 = nn.Conv2d(32, 64, 3, 1).to(device) - # Put rest of the network on the 2nd cuda device, if there is one - if "cuda" in str(device) and num_gpus > 1: - device = torch.device("cuda:1") - - print(f"Putting rest of layers on {str(device)}") - self.dropout1 = nn.Dropout2d(0.25).to(device) - self.dropout2 = nn.Dropout2d(0.5).to(device) - self.fc1 = nn.Linear(9216, 128).to(device) - self.fc2 = nn.Linear(128, 10).to(device) - - def forward(self, x): - x = self.conv1(x) - x = F.relu(x) - x = self.conv2(x) - x = F.max_pool2d(x, 2) - - x = self.dropout1(x) - x = torch.flatten(x, 1) - # Move tensor to next device if necessary - next_device = next(self.fc1.parameters()).device - x = x.to(next_device) - - x = self.fc1(x) - x = F.relu(x) - x = self.dropout2(x) - x = self.fc2(x) - output = F.log_softmax(x, dim=1) - return output -Next, let's define some helper functions that will be useful for the rest of our script. The following uses `rpc_sync `_ and `RRef `_ in order to define a function that invokes a given method on an object living on a remote node. Below, our handle to the remote object is given by the ``rref`` argument, and we run it on its owning node: ``rref.owner()``. On the caller node, we run this command synchronously through the use of ``rpc_sync``\ , meaning that we will block until a response is received. - -.. code-block:: python - - # --------- Helper Methods -------------------- - - # On the local node, call a method with first arg as the value held by the - # RRef. Other args are passed in as arguments to the function called. - # Useful for calling instance methods. method could be any matching function, including - # class methods. - def call_method(method, rref, *args, **kwargs): - return method(rref.local_value(), *args, **kwargs) - - # Given an RRef, return the result of calling the passed in method on the value - # held by the RRef. This call is done on the remote node that owns - # the RRef and passes along the given argument. - # Example: If the value held by the RRef is of type Foo, then - # remote_method(Foo.bar, rref, arg1, arg2) is equivalent to calling - # .bar(arg1, arg2) on the remote node and getting the result - # back. - - def remote_method(method, rref, *args, **kwargs): - args = [method, rref] + list(args) - return rpc.rpc_sync(rref.owner(), call_method, args=args, kwargs=kwargs) -Now, we're ready to define our parameter server. We will subclass ``nn.Module`` and save a handle to our network defined above. We'll also save an input device which will be the device our input is transferred to before invoking the model. - -.. code-block:: python - - # --------- Parameter Server -------------------- - class ParameterServer(nn.Module): - def __init__(self, num_gpus=0): - super().__init__() - model = Net(num_gpus=num_gpus) - self.model = model - self.input_device = torch.device( - "cuda:0" if torch.cuda.is_available() and num_gpus > 0 else "cpu") -Next, we'll define our forward pass. Note that regardless of the device of the model output, we move the output to CPU, as the Distributed RPC Framework currently only supports sending CPU tensors over RPC. We have intentionally disabled sending CUDA tensors over RPC due to the potential for different devices (CPU/GPU) on on the caller/callee, but may support this in future releases. - -.. code-block:: python - - class ParameterServer(nn.Module): - ... - def forward(self, inp): - inp = inp.to(self.input_device) - out = self.model(inp) - # This output is forwarded over RPC, which as of 1.5.0 only accepts CPU tensors. - # Tensors must be moved in and out of GPU memory due to this. - out = out.to("cpu") - return out -Next, we'll define a few miscellaneous functions useful for training and verification purposes. The first, ``get_dist_gradients``\ , will take in a Distributed Autograd context ID and call into the ``dist_autograd.get_gradients`` API in order to retrieve gradients computed by distributed autograd. More information can be found in the `distributed autograd documentation `_. Note that we also iterate through the resulting dictionary and convert each tensor to a CPU tensor, as the framework currently only supports sending tensors over RPC. Next, ``get_param_rrefs`` will iterate through our model parameters and wrap them as a (local) `RRef `_. This method will be invoked over RPC by trainer nodes and will return a list of the parameters to be optimized. This is required as input to the `Distributed Optimizer `_\ , which requires all parameters it must optimize as a list of ``RRef``\ s. - -.. code-block:: python - - # Use dist autograd to retrieve gradients accumulated for this model. - # Primarily used for verification. - def get_dist_gradients(self, cid): - grads = dist_autograd.get_gradients(cid) - # This output is forwarded over RPC, which as of 1.5.0 only accepts CPU tensors. - # Tensors must be moved in and out of GPU memory due to this. - cpu_grads = {} - for k, v in grads.items(): - k_cpu, v_cpu = k.to("cpu"), v.to("cpu") - cpu_grads[k_cpu] = v_cpu - return cpu_grads - - # Wrap local parameters in a RRef. Needed for building the - # DistributedOptimizer which optimizes paramters remotely. - def get_param_rrefs(self): - param_rrefs = [rpc.RRef(param) for param in self.model.parameters()] - return param_rrefs -Finally, we'll create methods to initialize our parameter server. Note that there will only be one instance of a parameter server across all processes, and all trainers will talk to the same parameter server and update the same stored model. As seen in ``run_parameter_server``\ , the server itself does not take any independent actions; it waits for requests from trainers (which are yet to be defined) and responds to them by running the requested function. - -.. code-block:: python - - # The global parameter server instance. - param_server = None - # A lock to ensure we only have one parameter server. - global_lock = Lock() - - - def get_parameter_server(num_gpus=0): - """ - Returns a singleton parameter server to all trainer processes - """ - global param_server - # Ensure that we get only one handle to the ParameterServer. - with global_lock: - if not param_server: - # construct it once - param_server = ParameterServer(num_gpus=num_gpus) - return param_server - - def run_parameter_server(rank, world_size): - # The parameter server just acts as a host for the model and responds to - # requests from trainers. - # rpc.shutdown() will wait for all workers to complete by default, which - # in this case means that the parameter server will wait for all trainers - # to complete, and then exit. - print("PS master initializing RPC") - rpc.init_rpc(name="parameter_server", rank=rank, world_size=world_size) - print("RPC initialized! Running parameter server...") - rpc.shutdown() - print("RPC shutdown on parameter server.") -Note that above, ``rpc.shutdown()`` will not immediately shut down the Parameter Server. Instead, it will wait for all workers (trainers in this case) to also call into ``rpc.shutdown()``. This gives us the guarantee that the parameter server will not go offline before all trainers (yet to be define) have completed their training process. - -Next, we'll define our ``TrainerNet`` class. This will also be a subclass of ``nn.Module``\ , and our ``__init__`` method will use the ``rpc.remote`` API to obtain an RRef, or Remote Reference, to our parameter server. Note that here we are not copying the parameter server to our local process, instead, we can think of ``self.param_server_rref`` as a distributed shared pointer to the parameter server that lives on a separate process. - -.. code-block:: python - - # --------- Trainers -------------------- - - # nn.Module corresponding to the network trained by this trainer. The - # forward() method simply invokes the network on the given parameter - # server. - class TrainerNet(nn.Module): - def __init__(self, num_gpus=0): - super().__init__() - self.num_gpus = num_gpus - self.param_server_rref = rpc.remote( - "parameter_server", get_parameter_server, args=(num_gpus,)) -Next, we'll define a method called ``get_global_param_rrefs``. To motivate the need for this method, it is worth it to read through the documentation on `DistributedOptimizer `_, specifically the API signature. The optimizer must be passed a list of ``RRef``\ s corresponding to the remote parameters to be optimized, so here we obtain the necessary ``RRef``\ s. Since the only remote worker that a given ``TrainerNet`` interacts with is the ``ParameterServer``\ , we simply invoke a ``remote_method`` on the ``ParameterServer``. We use the ``get_param_rrefs`` method which we defined in the ``ParameterServer`` class. This method will return a list of ``RRef``\ s to the parameters that need to be optimized. Note that in this case our ``TrainerNet`` does not define its own paramaters; if it did, we would need to wrap each parameter in an ``RRef`` as well and include it into our input to ``DistributedOptimizer``. - -.. code-block:: python - - class TrainerNet(nn.Module): - ... - def get_global_param_rrefs(self): - remote_params = remote_method( - ParameterServer.get_param_rrefs, - self.param_server_rref) - return remote_params -Now, we're ready to define our ``forward`` method, which will invoke (synchronous) RPC to run the forward pass of the network defined on the ``ParameterServer``. Note that we pass in ``self.param_server_rref``\ , which is a remote handle to our ``ParameterServer``\ , to our RPC call. This call will send an RPC to the node on which our ``ParameterServer`` is running, invoke the ``forward`` pass, and return the ``Tensor`` corresponding to the model's output. - -.. code-block:: python - - class TrainerNet(nn.Module): - ... - def forward(self, x): - model_output = remote_method( - ParameterServer.forward, self.param_server_rref, x) - return model_output -With our trainer fully defined, it's now time to write our neural network training loop that will create our network and optimizer, run some inputs through the network and compute the loss. The training loop looks a lot like that of a local training program, with some modifications due to the nature of our network being distributed across machines. - -Below, we initialize our ``TrainerNet`` and build a ``DistributedOptimizer``. Note that as mentioned above, we must pass in all of the global (across all nodes participating in distributed training) parameters that we want to be optimized. In addition, we pass in the local optimizer to be used, in this case, SGD. Note that we can configure the underlying optimizer algorithm in the same way as creating a local optimizer - all arguments for ``optimizer.SGD`` will be forwarded properly. As an example, we pass in a custom learning rate that will be used as the learning rate for all local optimizers. - -.. code-block:: python - - def run_training_loop(rank, num_gpus, train_loader, test_loader): - # Runs the typical nueral network forward + backward + optimizer step, but - # in a distributed fashion. - net = TrainerNet(num_gpus=num_gpus) - # Build DistributedOptimizer. - param_rrefs = net.get_global_param_rrefs() - opt = DistributedOptimizer(optim.SGD, param_rrefs, lr=0.03) -Next, we define our main training loop. We loop through iterables given by PyTorch's `DataLoader `_. Before writing our typical forward/backward/optimizer loop, we first wrap the logic within a `Distributed Autograd context `_. Note that this is needed to record RPCs invoked in the model's forward pass, so that an appropriate graph can be constructed which includes all participating distributed workers in the backward pass. The distributed autograd context returns a ``context_id`` which serves as an identifier for accumulating and optimizing gradients corresponding to a particular iteration. - -As opposed to calling the typical ``loss.backward()`` which would kick off the backward pass on this local worker, we call ``dist_autograd.backward()`` and pass in our context_id as well as ``loss``\ , which is the root at which we want the backward pass to begin. In addition, we pass this ``context_id`` into our optimizer call, which is required to be able to look up the corresponding gradients computed by this particular backwards pass across all nodes. - -.. code-block:: python - - def run_training_loop(rank, num_gpus, train_loader, test_loader): - ... - for i, (data, target) in enumerate(train_loader): - with dist_autograd.context() as cid: - model_output = net(data) - target = target.to(model_output.device) - loss = F.nll_loss(model_output, target) - if i % 5 == 0: - print(f"Rank {rank} training batch {i} loss {loss.item()}") - dist_autograd.backward(cid, [loss]) - # Ensure that dist autograd ran successfully and gradients were - # returned. - assert remote_method( - ParameterServer.get_dist_gradients, - net.param_server_rref, - cid) != {} - opt.step(cid) - - print("Training complete!") - print("Getting accuracy....") - get_accuracy(test_loader, net) -The following simply computes the accuracy of our model after we're done training, much like a traditional local model. However, note that the ``net`` we pass into this function above is an instance of ``TrainerNet`` and therefore the forward pass invokes RPC in a transparent fashion. - -.. code-block:: python - - def get_accuracy(test_loader, model): - model.eval() - correct_sum = 0 - # Use GPU to evaluate if possible - device = torch.device("cuda:0" if model.num_gpus > 0 - and torch.cuda.is_available() else "cpu") - with torch.no_grad(): - for i, (data, target) in enumerate(test_loader): - out = model(data, -1) - pred = out.argmax(dim=1, keepdim=True) - pred, target = pred.to(device), target.to(device) - correct = pred.eq(target.view_as(pred)).sum().item() - correct_sum += correct - - print(f"Accuracy {correct_sum / len(test_loader.dataset)}") -Next, similar to how we defined ``run_parameter_server`` as the main loop for our ``ParameterServer`` that is responsible for initializing RPC, let's define a similar loop for our trainers. The difference will be that our trainers must run the training loop we defined above: - -.. code-block:: python - - # Main loop for trainers. - def run_worker(rank, world_size, num_gpus, train_loader, test_loader): - print(f"Worker rank {rank} initializing RPC") - rpc.init_rpc( - name=f"trainer_{rank}", - rank=rank, - world_size=world_size) - - print(f"Worker {rank} done initializing RPC") - - run_training_loop(rank, num_gpus, train_loader, test_loader) - rpc.shutdown() -Note that similar to ``run_parameter_server``\ , ``rpc.shutdown()`` will by default wait for all workers, both trainers and ParameterServers, to call into ``rpc.shutdown()`` before this node exits. This ensures that nodes are terminated gracefully and no node goes offline while another is expecting it to be online. - -We've now completed our trainer and parameter server specific code, and all that's left is to add code to launch trainers and parameter servers. First, we must take in various arguments that apply to our parameter server and trainers. ``world_size`` corresponds to the total number of nodes that will participate in training, and is the sum of all trainers and the parameter server. We also must pass in a unique ``rank`` for each individual process, from 0 (where we will run our single parameter server) to ``world_size - 1``. ``master_addr`` and ``master_port`` are arguments that can be used to identify where the rank 0 process is running, and will be used by individual nodes to discover each other. To test this example out locally, simply pass in ``localhost`` and the same ``master_port`` to all instances spawned. Note that for demonstration purposes, this example supports only between 0-2 GPUs, although the pattern can be extended to make use of additional GPUs. - -.. code-block:: python - - if __name__ == '__main__': - parser = argparse.ArgumentParser( - description="Parameter-Server RPC based training") - parser.add_argument( - "--world_size", - type=int, - default=4, - help="""Total number of participating processes. Should be the sum of - master node and all training nodes.""") - parser.add_argument( - "--rank", - type=int, - default=None, - help="Global rank of this process. Pass in 0 for master.") - parser.add_argument( - "--num_gpus", - type=int, - default=0, - help="""Number of GPUs to use for training, Currently supports between 0 - and 2 GPUs. Note that this argument will be passed to the parameter servers.""") - parser.add_argument( - "--master_addr", - type=str, - default="localhost", - help="""Address of master, will default to localhost if not provided. - Master must be able to accept network traffic on the address + port.""") - parser.add_argument( - "--master_port", - type=str, - default="29500", - help="""Port that master is listening on, will default to 29500 if not - provided. Master must be able to accept network traffic on the host and port.""") - - args = parser.parse_args() - assert args.rank is not None, "must provide rank argument." - assert args.num_gpus <= 3, f"Only 0-2 GPUs currently supported (got {args.num_gpus})." - os.environ['MASTER_ADDR'] = args.master_addr - os.environ["MASTER_PORT"] = args.master_port -Now, we'll create a process corresponding to either a parameter server or trainer depending on our command line arguments. We'll create a ``ParameterServer`` if our passed in rank is 0, and a ``TrainerNet`` otherwise. Note that we're using ``torch.multiprocessing`` to launch a subprocess corresponding to the function that we want to execute, and waiting on this process's completion from the main thread with ``p.join()``. In the case of initializing our trainers, we also use PyTorch's `dataloaders `_ in order to specify train and test data loaders on the MNIST dataset. - -.. code-block:: python - - processes = [] - world_size = args.world_size - if args.rank == 0: - p = mp.Process(target=run_parameter_server, args=(0, world_size)) - p.start() - processes.append(p) - else: - # Get data to train on - train_loader = torch.utils.data.DataLoader( - datasets.MNIST('../data', train=True, download=True, - transform=transforms.Compose([ - transforms.ToTensor(), - transforms.Normalize((0.1307,), (0.3081,)) - ])), - batch_size=32, shuffle=True,) - test_loader = torch.utils.data.DataLoader( - datasets.MNIST( - '../data', - train=False, - transform=transforms.Compose([ - transforms.ToTensor(), - transforms.Normalize((0.1307,), (0.3081,)) - ])), - batch_size=32, - shuffle=True, - ) - # start training worker on this node - p = mp.Process( - target=run_worker, - args=( - args.rank, - world_size, args.num_gpus, - train_loader, - test_loader)) - p.start() - processes.append(p) - - for p in processes: - p.join() -To run the example locally, run the following command worker for the server and each worker you wish to spawn, in separate terminal windows: ``python rpc_parameter_server.py --world_size=WORLD_SIZE --rank=RANK``. For example, for a master node with world size of 2, the command would be ``python rpc_parameter_server.py --world_size=2 --rank=0``. The trainer can then be launched with the command ``python rpc_parameter_server.py --world_size=2 --rank=1`` in a separate window, and this will begin training with one server and a single trainer. Note that this tutorial assumes that training occurs using between 0 and 2 GPUs, and this argument can be configured by passing ``--num_gpus=N`` into the training script. - -You can pass in the command line arguments ``--master_addr=ADDRESS`` and ``--master_port=PORT`` to indicate the address and port that the master worker is listening on, for example, to test functionality where trainers and master nodes run on different machines. diff --git a/intermediate_source/rpc_tutorial.rst b/intermediate_source/rpc_tutorial.rst deleted file mode 100644 index 791ecf86d35..00000000000 --- a/intermediate_source/rpc_tutorial.rst +++ /dev/null @@ -1,622 +0,0 @@ -Getting Started with Distributed RPC Framework -================================================= -**Author**: `Shen Li `_ - -.. note:: - |edit| View and edit this tutorial in `github `__. - -Prerequisites: - -- `PyTorch Distributed Overview <../beginner/dist_overview.html>`__ -- `RPC API documents `__ - -This tutorial uses two simple examples to demonstrate how to build distributed -training with the `torch.distributed.rpc `__ -package which was first introduced as an experimental feature in PyTorch v1.4. -Source code of the two examples can be found in -`PyTorch examples `__. - -Previous tutorials, -`Getting Started With Distributed Data Parallel `__ -and `Writing Distributed Applications With PyTorch `__, -described `DistributedDataParallel `__ -which supports a specific training paradigm where the model is replicated across -multiple processes and each process handles a split of the input data. -Sometimes, you might run into scenarios that require different training -paradigms. For example: - -1) In reinforcement learning, it might be relatively expensive to acquire - training data from environments while the model itself can be quite small. In - this case, it might be useful to spawn multiple observers running in parallel - and share a single agent. In this case, the agent takes care of the training - locally, but the application would still need libraries to send and receive - data between observers and the trainer. -2) Your model might be too large to fit in GPUs on a single machine, and hence - would need a library to help split the model onto multiple machines. Or you - might be implementing a `parameter server `__ - training framework, where model parameters and trainers live on different - machines. - - -The `torch.distributed.rpc `__ package -can help with the above scenarios. In case 1, `RPC `__ -and `RRef `__ allow sending data -from one worker to another while easily referencing remote data objects. In -case 2, `distributed autograd `__ -and `distributed optimizer `__ -make executing backward pass and optimizer step as if it is local training. In -the next two sections, we will demonstrate APIs of -`torch.distributed.rpc `__ using a -reinforcement learning example and a language model example. Please note, this -tutorial does not aim at building the most accurate or efficient models to -solve given problems, instead, the main goal here is to show how to use the -`torch.distributed.rpc `__ package to -build distributed training applications. - - - -Distributed Reinforcement Learning using RPC and RRef ------------------------------------------------------ - -This section describes steps to build a toy distributed reinforcement learning -model using RPC to solve CartPole-v1 from `OpenAI Gym `__. -The policy code is mostly borrowed from the existing single-thread -`example `__ -as shown below. We will skip details of the ``Policy`` design, and focus on RPC -usages. - -.. code:: python - - import torch.nn as nn - import torch.nn.functional as F - - class Policy(nn.Module): - - def __init__(self): - super(Policy, self).__init__() - self.affine1 = nn.Linear(4, 128) - self.dropout = nn.Dropout(p=0.6) - self.affine2 = nn.Linear(128, 2) - - def forward(self, x): - x = self.affine1(x) - x = self.dropout(x) - x = F.relu(x) - action_scores = self.affine2(x) - return F.softmax(action_scores, dim=1) - - -We are ready to present the observer. In this example, each observer creates its -own environment, and waits for the agent's command to run an episode. In each -episode, one observer loops at most ``n_steps`` iterations, and in each -iteration, it uses RPC to pass its environment state to the agent and gets an -action back. Then it applies that action to its environment, and gets the reward -and the next state from the environment. After that, the observer uses another -RPC to report the reward to the agent. Again, please note that, this is -obviously not the most efficient observer implementation. For example, one -simple optimization could be packing current state and last reward in one RPC to -reduce the communication overhead. However, the goal is to demonstrate RPC API -instead of building the best solver for CartPole. So, let's keep the logic -simple and the two steps explicit in this example. - -.. code:: python - - import argparse - import gym - import torch.distributed.rpc as rpc - - parser = argparse.ArgumentParser( - description="RPC Reinforcement Learning Example", - formatter_class=argparse.ArgumentDefaultsHelpFormatter, - ) - - parser.add_argument('--world_size', default=2, type=int, metavar='W', - help='number of workers') - parser.add_argument('--log_interval', type=int, default=10, metavar='N', - help='interval between training status logs') - parser.add_argument('--gamma', type=float, default=0.99, metavar='G', - help='how much to value future rewards') - parser.add_argument('--seed', type=int, default=1, metavar='S', - help='random seed for reproducibility') - args = parser.parse_args() - - class Observer: - - def __init__(self): - self.id = rpc.get_worker_info().id - self.env = gym.make('CartPole-v1') - self.env.seed(args.seed) - - def run_episode(self, agent_rref): - state, ep_reward = self.env.reset(), 0 - for _ in range(10000): - # send the state to the agent to get an action - action = agent_rref.rpc_sync().select_action(self.id, state) - - # apply the action to the environment, and get the reward - state, reward, done, _ = self.env.step(action) - - # report the reward to the agent for training purpose - agent_rref.rpc_sync().report_reward(self.id, reward) - - # finishes after the number of self.env._max_episode_steps - if done: - break - - -The code for agent is a little more complex, and we will break it into multiple -pieces. In this example, the agent serves as both the trainer and the master, -such that it sends command to multiple distributed observers to run episodes, -and it also records all actions and rewards locally which will be used during -the training phase after each episode. The code below shows ``Agent`` -constructor where most lines are initializing various components. The loop at -the end initializes observers remotely on other workers, and holds ``RRefs`` to -those observers locally. The agent will use those observer ``RRefs`` later to -send commands. Applications don't need to worry about the lifetime of ``RRefs``. -The owner of each ``RRef`` maintains a reference counting map to track its -lifetime, and guarantees the remote data object will not be deleted as long as -there is any live user of that ``RRef``. Please refer to the ``RRef`` -`design doc `__ for details. - - -.. code:: python - - import gym - import numpy as np - - import torch - import torch.distributed.rpc as rpc - import torch.optim as optim - from torch.distributed.rpc import RRef, rpc_async, remote - from torch.distributions import Categorical - - class Agent: - def __init__(self, world_size): - self.ob_rrefs = [] - self.agent_rref = RRef(self) - self.rewards = {} - self.saved_log_probs = {} - self.policy = Policy() - self.optimizer = optim.Adam(self.policy.parameters(), lr=1e-2) - self.eps = np.finfo(np.float32).eps.item() - self.running_reward = 0 - self.reward_threshold = gym.make('CartPole-v1').spec.reward_threshold - for ob_rank in range(1, world_size): - ob_info = rpc.get_worker_info(OBSERVER_NAME.format(ob_rank)) - self.ob_rrefs.append(remote(ob_info, Observer)) - self.rewards[ob_info.id] = [] - self.saved_log_probs[ob_info.id] = [] - - -Next, the agent exposes two APIs to observers for selecting actions and -reporting rewards. Those functions only run locally on the agent, but will -be triggered by observers through RPC. - - -.. code:: python - - class Agent: - ... - def select_action(self, ob_id, state): - state = torch.from_numpy(state).float().unsqueeze(0) - probs = self.policy(state) - m = Categorical(probs) - action = m.sample() - self.saved_log_probs[ob_id].append(m.log_prob(action)) - return action.item() - - def report_reward(self, ob_id, reward): - self.rewards[ob_id].append(reward) - - -Let's add a ``run_episode`` function on agent which tells all observers -to execute an episode. In this function, it first creates a list to collect -futures from asynchronous RPCs, and then loop over all observer ``RRefs`` to -make asynchronous RPCs. In these RPCs, the agent also passes an ``RRef`` of -itself to the observer, so that the observer can call functions on the agent as -well. As shown above, each observer will make RPCs back to the agent, which are -nested RPCs. After each episode, the ``saved_log_probs`` and ``rewards`` will -contain the recorded action probs and rewards. - - -.. code:: python - - class Agent: - ... - def run_episode(self): - futs = [] - for ob_rref in self.ob_rrefs: - # make async RPC to kick off an episode on all observers - futs.append( - rpc_async( - ob_rref.owner(), - ob_rref.rpc_sync().run_episode, - args=(self.agent_rref,) - ) - ) - - # wait until all obervers have finished this episode - for fut in futs: - fut.wait() - - -Finally, after one episode, the agent needs to train the model, which -is implemented in the ``finish_episode`` function below. There is no RPCs in -this function and it is mostly borrowed from the single-thread -`example `__. -Hence, we skip describing its contents. - - - -.. code:: python - - class Agent: - ... - def finish_episode(self): - # joins probs and rewards from different observers into lists - R, probs, rewards = 0, [], [] - for ob_id in self.rewards: - probs.extend(self.saved_log_probs[ob_id]) - rewards.extend(self.rewards[ob_id]) - - # use the minimum observer reward to calculate the running reward - min_reward = min([sum(self.rewards[ob_id]) for ob_id in self.rewards]) - self.running_reward = 0.05 * min_reward + (1 - 0.05) * self.running_reward - - # clear saved probs and rewards - for ob_id in self.rewards: - self.rewards[ob_id] = [] - self.saved_log_probs[ob_id] = [] - - policy_loss, returns = [], [] - for r in rewards[::-1]: - R = r + args.gamma * R - returns.insert(0, R) - returns = torch.tensor(returns) - returns = (returns - returns.mean()) / (returns.std() + self.eps) - for log_prob, R in zip(probs, returns): - policy_loss.append(-log_prob * R) - self.optimizer.zero_grad() - policy_loss = torch.cat(policy_loss).sum() - policy_loss.backward() - self.optimizer.step() - return min_reward - - -With ``Policy``, ``Observer``, and ``Agent`` classes, we are ready to launch -multiple processes to perform the distributed training. In this example, all -processes run the same ``run_worker`` function, and they use the rank to -distinguish their role. Rank 0 is always the agent, and all other ranks are -observers. The agent serves as master by repeatedly calling ``run_episode`` and -``finish_episode`` until the running reward surpasses the reward threshold -specified by the environment. All observers passively waiting for commands -from the agent. The code is wrapped by -`rpc.init_rpc `__ and -`rpc.shutdown `__, -which initializes and terminates RPC instances respectively. More details are -available in the `API page `__. - - -.. code:: python - - import os - from itertools import count - - import torch.multiprocessing as mp - - AGENT_NAME = "agent" - OBSERVER_NAME="obs{}" - - def run_worker(rank, world_size): - os.environ['MASTER_ADDR'] = 'localhost' - os.environ['MASTER_PORT'] = '29500' - if rank == 0: - # rank0 is the agent - rpc.init_rpc(AGENT_NAME, rank=rank, world_size=world_size) - - agent = Agent(world_size) - print(f"This will run until reward threshold of {agent.reward_threshold}" - " is reached. Ctrl+C to exit.") - for i_episode in count(1): - agent.run_episode() - last_reward = agent.finish_episode() - - if i_episode % args.log_interval == 0: - print(f"Episode {i_episode}\tLast reward: {last_reward:.2f}\tAverage reward: " - f"{agent.running_reward:.2f}") - if agent.running_reward > agent.reward_threshold: - print(f"Solved! Running reward is now {agent.running_reward}!") - break - else: - # other ranks are the observer - rpc.init_rpc(OBSERVER_NAME.format(rank), rank=rank, world_size=world_size) - # observers passively waiting for instructions from the agent - - # block until all rpcs finish, and shutdown the RPC instance - rpc.shutdown() - - - mp.spawn( - run_worker, - args=(args.world_size, ), - nprocs=args.world_size, - join=True - ) - -Below are some sample outputs when training with `world_size=2`. - -:: - - This will run until reward threshold of 475.0 is reached. Ctrl+C to exit. - Episode 10 Last reward: 26.00 Average reward: 10.01 - Episode 20 Last reward: 16.00 Average reward: 11.27 - Episode 30 Last reward: 49.00 Average reward: 18.62 - Episode 40 Last reward: 45.00 Average reward: 26.09 - Episode 50 Last reward: 44.00 Average reward: 30.03 - Episode 60 Last reward: 111.00 Average reward: 42.23 - Episode 70 Last reward: 131.00 Average reward: 70.11 - Episode 80 Last reward: 87.00 Average reward: 76.51 - Episode 90 Last reward: 86.00 Average reward: 95.93 - Episode 100 Last reward: 13.00 Average reward: 123.93 - Episode 110 Last reward: 33.00 Average reward: 91.39 - Episode 120 Last reward: 73.00 Average reward: 76.38 - Episode 130 Last reward: 137.00 Average reward: 88.08 - Episode 140 Last reward: 89.00 Average reward: 104.96 - Episode 150 Last reward: 97.00 Average reward: 98.74 - Episode 160 Last reward: 150.00 Average reward: 100.87 - Episode 170 Last reward: 126.00 Average reward: 104.38 - Episode 180 Last reward: 500.00 Average reward: 213.74 - Episode 190 Last reward: 322.00 Average reward: 300.22 - Episode 200 Last reward: 165.00 Average reward: 272.71 - Episode 210 Last reward: 168.00 Average reward: 233.11 - Episode 220 Last reward: 184.00 Average reward: 195.02 - Episode 230 Last reward: 284.00 Average reward: 208.32 - Episode 240 Last reward: 395.00 Average reward: 247.37 - Episode 250 Last reward: 500.00 Average reward: 335.42 - Episode 260 Last reward: 500.00 Average reward: 386.30 - Episode 270 Last reward: 500.00 Average reward: 405.29 - Episode 280 Last reward: 500.00 Average reward: 443.29 - Episode 290 Last reward: 500.00 Average reward: 464.65 - Solved! Running reward is now 475.3163778435275! - - -In this example, we show how to use RPC as the communication vehicle to pass -data across workers, and how to use RRef to reference remote objects. It is true -that you could build the entire structure directly on top of ``ProcessGroup`` -``send`` and ``recv`` APIs or use other communication/RPC libraries. However, -by using `torch.distributed.rpc`, you can get the native support and -continuously optimized performance under the hood. - -Next, we will show how to combine RPC and RRef with distributed autograd and -distributed optimizer to perform distributed model parallel training. - - - -Distributed RNN using Distributed Autograd and Distributed Optimizer --------------------------------------------------------------------- - -In this section, we use an RNN model to show how to build distributed model -parallel training with the RPC API. The example RNN model is very small and -can easily fit into a single GPU, but we still divide its layers onto two -different workers to demonstrate the idea. Developer can apply the similar -techniques to distribute much larger models across multiple devices and -machines. - -The RNN model design is borrowed from the word language model in PyTorch -`example `__ -repository, which contains three main components, an embedding table, an -``LSTM`` layer, and a decoder. The code below wraps the embedding table and the -decoder into sub-modules, so that their constructors can be passed to the RPC -API. In the ``EmbeddingTable`` sub-module, we intentionally put the -``Embedding`` layer on GPU to cover the use case. In v1.4, RPC always creates -CPU tensor arguments or return values on the destination worker. If the function -takes a GPU tensor, you need to move it to the proper device explicitly. - - -.. code:: python - - class EmbeddingTable(nn.Module): - r""" - Encoding layers of the RNNModel - """ - def __init__(self, ntoken, ninp, dropout): - super(EmbeddingTable, self).__init__() - self.drop = nn.Dropout(dropout) - self.encoder = nn.Embedding(ntoken, ninp).cuda() - self.encoder.weight.data.uniform_(-0.1, 0.1) - - def forward(self, input): - return self.drop(self.encoder(input.cuda()).cpu() - - - class Decoder(nn.Module): - def __init__(self, ntoken, nhid, dropout): - super(Decoder, self).__init__() - self.drop = nn.Dropout(dropout) - self.decoder = nn.Linear(nhid, ntoken) - self.decoder.bias.data.zero_() - self.decoder.weight.data.uniform_(-0.1, 0.1) - - def forward(self, output): - return self.decoder(self.drop(output)) - - -With the above sub-modules, we can now piece them together using RPC to -create an RNN model. In the code below ``ps`` represents a parameter server, -which hosts parameters of the embedding table and the decoder. The constructor -uses the `remote `__ -API to create an ``EmbeddingTable`` object and a ``Decoder`` object on the -parameter server, and locally creates the ``LSTM`` sub-module. During the -forward pass, the trainer uses the ``EmbeddingTable`` ``RRef`` to find the -remote sub-module and passes the input data to the ``EmbeddingTable`` using RPC -and fetches the lookup results. Then, it runs the embedding through the local -``LSTM`` layer, and finally uses another RPC to send the output to the -``Decoder`` sub-module. In general, to implement distributed model parallel -training, developers can divide the model into sub-modules, invoke RPC to create -sub-module instances remotely, and use on ``RRef`` to find them when necessary. -As you can see in the code below, it looks very similar to single-machine model -parallel training. The main difference is replacing ``Tensor.to(device)`` with -RPC functions. - - -.. code:: python - - class RNNModel(nn.Module): - def __init__(self, ps, ntoken, ninp, nhid, nlayers, dropout=0.5): - super(RNNModel, self).__init__() - - # setup embedding table remotely - self.emb_table_rref = rpc.remote(ps, EmbeddingTable, args=(ntoken, ninp, dropout)) - # setup LSTM locally - self.rnn = nn.LSTM(ninp, nhid, nlayers, dropout=dropout) - # setup decoder remotely - self.decoder_rref = rpc.remote(ps, Decoder, args=(ntoken, nhid, dropout)) - - def forward(self, input, hidden): - # pass input to the remote embedding table and fetch emb tensor back - emb = _remote_method(EmbeddingTable.forward, self.emb_table_rref, input) - output, hidden = self.rnn(emb, hidden) - # pass output to the rremote decoder and get the decoded output back - decoded = _remote_method(Decoder.forward, self.decoder_rref, output) - return decoded, hidden - -Before introducing the distributed optimizer, let's add a helper function to -generate a list of RRefs of model parameters, which will be consumed by the -distributed optimizer. In local training, applications could call -``Module.parameters()`` to grab references to all parameter tensors, and pass it -to the local optimizer for subsequent updates. However, the same API does not -work in distributed training scenarios as some parameters live on remote -machines. Therefore, instead of taking a list of parameter ``Tensors``, the -distributed optimizer takes a list of ``RRefs``, one ``RRef`` per model -parameter for both local and remote model parameters. The helper function is -pretty simple, just call ``Module.parameters()`` and creates a local ``RRef`` on -each of the parameters. - - -.. code:: python - - def _parameter_rrefs(module): - param_rrefs = [] - for param in module.parameters(): - param_rrefs.append(RRef(param)) - return param_rrefs - - -Then, as the ``RNNModel`` contains three sub-modules, we need to call -``_parameter_rrefs`` three times, and wrap that into another helper function. - - -.. code:: python - - class RNNModel(nn.Module): - ... - def parameter_rrefs(self): - remote_params = [] - # get RRefs of embedding table - remote_params.extend(_remote_method(_parameter_rrefs, self.emb_table_rref)) - # create RRefs for local parameters - remote_params.extend(_parameter_rrefs(self.rnn)) - # get RRefs of decoder - remote_params.extend(_remote_method(_parameter_rrefs, self.decoder_rref)) - return remote_params - - -Now, we are ready to implement the training loop. After initializing model -arguments, we create the ``RNNModel`` and the ``DistributedOptimizer``. The -distributed optimizer will take a list of parameter ``RRefs``, find all distinct -owner workers, and create the given local optimizer (i.e., ``SGD`` in this case, -you can use other local optimizers as well) on each of the owner worker using -the given arguments (i.e., ``lr=0.05``). - -In the training loop, it first creates a distributed autograd context, which -will help the distributed autograd engine to find gradients and involved RPC -send/recv functions. The design details of the distributed autograd engine can -be found in its `design note `__. -Then, it kicks off the forward pass as if it is a local -model, and run the distributed backward pass. For the distributed backward, you -only need to specify a list of roots, in this case, it is the loss ``Tensor``. -The distributed autograd engine will traverse the distributed graph -automatically and write gradients properly. Next, it runs the ``step`` -function on the distributed optimizer, which will reach out to all involved -local optimizers to update model parameters. Compared to local training, one -minor difference is that you don't need to run ``zero_grad()`` because each -autograd context has dedicated space to store gradients, and as we create a -context per iteration, those gradients from different iterations will not -accumulate to the same set of ``Tensors``. - - -.. code:: python - - def run_trainer(): - batch = 5 - ntoken = 10 - ninp = 2 - - nhid = 3 - nindices = 3 - nlayers = 4 - hidden = ( - torch.randn(nlayers, nindices, nhid), - torch.randn(nlayers, nindices, nhid) - ) - - model = rnn.RNNModel('ps', ntoken, ninp, nhid, nlayers) - - # setup distributed optimizer - opt = DistributedOptimizer( - optim.SGD, - model.parameter_rrefs(), - lr=0.05, - ) - - criterion = torch.nn.CrossEntropyLoss() - - def get_next_batch(): - for _ in range(5): - data = torch.LongTensor(batch, nindices) % ntoken - target = torch.LongTensor(batch, ntoken) % nindices - yield data, target - - # train for 10 iterations - for epoch in range(10): - for data, target in get_next_batch(): - # create distributed autograd context - with dist_autograd.context() as context_id: - hidden[0].detach_() - hidden[1].detach_() - output, hidden = model(data, hidden) - loss = criterion(output, target) - # run distributed backward pass - dist_autograd.backward(context_id, [loss]) - # run distributed optimizer - opt.step(context_id) - # not necessary to zero grads since they are - # accumulated into the distributed autograd context - # which is reset every iteration. - print("Training epoch {}".format(epoch)) - - -Finally, let's add some glue code to launch the parameter server and the trainer -processes. - - -.. code:: python - - def run_worker(rank, world_size): - os.environ['MASTER_ADDR'] = 'localhost' - os.environ['MASTER_PORT'] = '29500' - if rank == 1: - rpc.init_rpc("trainer", rank=rank, world_size=world_size) - _run_trainer() - else: - rpc.init_rpc("ps", rank=rank, world_size=world_size) - # parameter server do nothing - pass - - # block until all rpcs finish - rpc.shutdown() - - - if __name__=="__main__": - world_size = 2 - mp.spawn(run_worker, args=(world_size, ), nprocs=world_size, join=True) diff --git a/recipes_index.rst b/recipes_index.rst index a7116b0478d..e0efc670cf3 100644 --- a/recipes_index.rst +++ b/recipes_index.rst @@ -281,13 +281,6 @@ from our full-length tutorials. :link: recipes/zero_redundancy_optimizer.html :tags: Distributed-Training -.. customcarditem:: - :header: Direct Device-to-Device Communication with TensorPipe RPC - :card_description: How to use RPC with direct GPU-to-GPU communication. - :image: _static/img/thumbnails/cropped/profiler.png - :link: recipes/cuda_rpc.html - :tags: Distributed-Training - .. customcarditem:: :header: Getting Started with Distributed Checkpoint (DCP) :card_description: Learn how to checkpoint distributed models with Distributed Checkpoint package. diff --git a/redirects.py b/redirects.py index b69d1a6c627..f68c23e25b6 100644 --- a/redirects.py +++ b/redirects.py @@ -2,6 +2,7 @@ "advanced/cpp_extension.html": "https://docs.pytorch.org/tutorials/advanced/custom_ops_landing_page.html", "advanced/cpp_cuda_graphs.html": "../index.html", "advanced/dynamic_quantization_tutorial.html": "../index.html", + "advanced/rpc_ddp_tutorial.html": "https://docs.pytorch.org/docs/main/distributed/rpc.html", "advanced/static_quantization_tutorial.html": "../index.html", "advanced/super_resolution_with_onnxruntime.html": "../index.html", "advanced/torch_script_custom_classes": "../index.html", @@ -17,6 +18,9 @@ "intermediate/dynamic_quantization_bert_tutorial.html": "../index.html", "intermediate/flask_rest_api_tutorial": "../index.html", "intermediate/quantized_transfer_learning_tutorial.html": "../index.html", + "intermediate/rpc_async_execution.html": "https://docs.pytorch.org/docs/main/distributed/rpc.html", + "intermediate/rpc_param_server_tutorial.html": "https://docs.pytorch.org/docs/main/distributed/rpc.html", + "intermediate/rpc_tutorial.html": "https://docs.pytorch.org/docs/main/distributed/rpc.html", "intermediate/torchserve_with_ipex.html": "../index.html", "intermediate/torchserve_with_ipex_2.html": "../index.html", "intermediate/tiatoolbox_tutorial.html": "../index.html", @@ -50,6 +54,7 @@ "recipes/torchserve_vertexai_tutorial.html": "../index.html", "recipes/amx.html": "../index.html", "unstable_source/vulkan_workflow.rst": "../index.html", + "unstable/distributed_rpc_profiling.html": "https://docs.pytorch.org/docs/main/distributed/rpc.html", "unstable/semi_structured_sparse.html": "https://docs.pytorch.org/tutorials/advanced/semi_structured_sparse.html", "unstable/skip_param_init.html": "https://docs.pytorch.org/tutorials/recipes/recipes/module_load_state_dict_tips.html", "unstable_source/backend_config_tutorial.rst": "../index.html", diff --git a/unstable_source/README.txt b/unstable_source/README.txt index 55a94b43626..088593b2406 100644 --- a/unstable_source/README.txt +++ b/unstable_source/README.txt @@ -1,13 +1,9 @@ Prototype Tutorials ------------------ -1. distributed_rpc_profiling.rst - Profiling PyTorch RPC-Based Workloads - https://github.com/pytorch/tutorials/blob/main/unstable_source/distributed_rpc_profiling.rst - -2. flight_recorder_tutorial.rst +1. flight_recorder_tutorial.rst Flight Recorder User Guide https://pytorch.org/tutorials/prototype/flight_recorder_tutorial.html -3. python_extension_autoload.rst +2. python_extension_autoload.rst Autoloading Out-of-Tree Extension https://github.com/pytorch/tutorials/blob/main/unstable_source/python_extension_autoload.rst diff --git a/unstable_source/distributed_rpc_profiling.rst b/unstable_source/distributed_rpc_profiling.rst deleted file mode 100644 index ab2ecbb2d9b..00000000000 --- a/unstable_source/distributed_rpc_profiling.rst +++ /dev/null @@ -1,10 +0,0 @@ -Profiling PyTorch RPC-Based Workloads -====================================== - -This tutorial has been deprecated. - -Redirecting to homepage... - -.. raw:: html - -