Skip to content

Commit 06ef4a7

Browse files
rohan-varmafacebook-github-bot
rohan-varma
authored and
facebook-github-bot
committed
Add docs for RPC, dist autograd, and RRef modules (pytorch#29276)
Summary: Closes pytorch#28983. Documentation for `torch.distributed.rpc` and `torch.distributed.autograd` modules. Also fixes/tidies up some of the docstrings in rpc/autograd, and moves some functions to be private so they don't show up in the documentation. Note: Much of the text to describe/explain the RPC/RRef layers are taken from the following RFCs: pytorch#23110, pytorch#26759 Pull Request resolved: pytorch#29276 Differential Revision: D18478754 Pulled By: rohan-varma fbshipit-source-id: e9a7089baf5275304e5408d319eb9bf98e53fff8
1 parent ce70583 commit 06ef4a7

7 files changed

Lines changed: 88 additions & 15 deletions

File tree

‎docs/source/index.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ PyTorch is an optimized tensor library for deep learning using GPUs and CPUs.
4444
onnx
4545
optim
4646
quantization
47+
rpc
4748
torch.random <random>
4849
sparse
4950
storage

‎docs/source/rpc.rst‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
.. role:: hidden
2+
:class: hidden-section
3+
4+
Distributed RPC Framework
5+
=========================
6+
7+
The distributed RPC framework provides mechanisms for multi-machine model
8+
training through a set of primitives to allow for remote communication, and a
9+
higher-level API to automatically differentiate models split across several
10+
machines.
11+
12+
RPC and RRef Framework
13+
----------------------
14+
15+
Before using RPC and distributed autograd primitives, initialization must take
16+
place. First, a backend over which RPCs can be sent over must be initialized.
17+
The default (and currently, only available) implementation is the `ProcessGroup`
18+
backend, and must be initialized with `torch.distributed.init_process_group
19+
<https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group>`_
20+
before using other functions. See the `documentation for
21+
torch.distributed <https://pytorch.org/docs/stable/distributed.html>`_ for
22+
additional details. Next, to initialize the RPC framework we need to use
23+
`init_model_parallel` which would initialize the RPC framework, RRef framework
24+
and distributed autograd.
25+
26+
.. automodule:: torch.distributed.rpc
27+
.. autofunction:: init_model_parallel
28+
29+
RRef
30+
----
31+
32+
An `RRef` (Remote REFerence) is a reference to a value of some type `T`
33+
(e.g. `Tensor`) on a remote worker. This handle keeps the referenced remote
34+
value alive on the owner, but there is no implication that the value will be
35+
transferred to the local worker in the future. RRefs can be used in
36+
multi-machine training by holding references to `nn.Modules
37+
<https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`_ that exist on
38+
other workers, and calling the appropriate functions to retrieve or modify their
39+
parameters during training.
40+
41+
.. autoclass:: RRef
42+
:members:
43+
44+
RPC and RRef primitives
45+
-----------------------
46+
47+
This library provides primitives allowing users to create and modify references
48+
(RRefs) to remote data as well as remotely execute functions.
49+
50+
.. automodule:: torch.distributed.rpc.api
51+
:members:
52+
53+
Distributed Autograd Framework
54+
------------------------------
55+
56+
This module provides an RPC-based distributed autograd framework that can be
57+
used for applications such as model parallel training. In short, applications
58+
may send and receive gradient recording tensors over RPC. In the forward pass,
59+
we record when gradient recording tensors are sent over RPC and during the
60+
backward pass we use this information to perform a distributed backward pass
61+
using RPC. For more details see the design doc
62+
`here <https://github.com/pytorch/pytorch/pull/29175>`_.
63+
64+
.. automodule:: torch.distributed.autograd
65+
:members:

‎torch/csrc/distributed/autograd/init.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,7 @@ using the ``get_gradients`` API.
133133
computation. All the tensors should be scalars.
134134
135135
Example::
136+
136137
>> import torch.distributed.autograd as dist_autograd
137138
>> with dist_autograd.context() as context_id:
138139
>> pred = model.forward()
@@ -158,6 +159,7 @@ backward pass.
158159
gradients.
159160
160161
Example::
162+
161163
>> import torch.distributed.autograd as dist_autograd
162164
>> with dist_autograd.context() as context_id:
163165
>> t1 = torch.rand((3, 3), requires_grad=True)

‎torch/csrc/distributed/rpc/init.cpp‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,10 @@ PyObject* rpc_init(PyObject* /* unused */) {
4949
py::call_guard<py::gil_scoped_release>());
5050

5151
auto pyRRef =
52-
shared_ptr_class_<PyRRef>(module, "RRef")
52+
shared_ptr_class_<PyRRef>(module, "RRef", R"(
53+
A class encapsulating a reference to a value of some type on a remote worker.
54+
This handle will keep the referenced remote value alive on the worker.
55+
)")
5356
.def(py::init<const py::object&>())
5457
.def(
5558
// not releasing GIL here to avoid context switch on getters
@@ -131,7 +134,7 @@ PyObject* rpc_init(PyObject* /* unused */) {
131134
});
132135

133136
module.def(
134-
"invoke_rpc_builtin",
137+
"_invoke_rpc_builtin",
135138
[](RpcAgent& agent,
136139
const WorkerInfo& dst,
137140
const std::string& opName,
@@ -141,7 +144,7 @@ PyObject* rpc_init(PyObject* /* unused */) {
141144
});
142145

143146
module.def(
144-
"invoke_rpc_python_udf",
147+
"_invoke_rpc_python_udf",
145148
[](RpcAgent& agent,
146149
const WorkerInfo& dst,
147150
std::string& pickledPythonUDF,
@@ -150,7 +153,7 @@ PyObject* rpc_init(PyObject* /* unused */) {
150153
});
151154

152155
module.def(
153-
"invoke_remote_builtin",
156+
"_invoke_remote_builtin",
154157
[](RpcAgent& agent,
155158
const WorkerInfo& dst,
156159
const std::string& opName,
@@ -160,7 +163,7 @@ PyObject* rpc_init(PyObject* /* unused */) {
160163
});
161164

162165
module.def(
163-
"invoke_remote_python_udf",
166+
"_invoke_remote_python_udf",
164167
[](RpcAgent& agent,
165168
const WorkerInfo& dst,
166169
std::string& pickledPythonUDF,

‎torch/distributed/autograd/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,11 @@ class context(object):
2222
2323
This is only needed in the "FAST" mode (as described in
2424
https://github.com/pytorch/pytorch/issues/23110) for distributed autograd,
25-
where we assume all RPC communication is would also be part of the backward
25+
where we assume all RPC communication in the forward pass would also be part of the backward
2626
pass.
2727
2828
Example::
29+
2930
>> import torch.distributed.autograd as dist_autograd
3031
>> with dist_autograd.context() as context_id:
3132
>> forward pass...

‎torch/distributed/rpc/api.py‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
from . import invoke_rpc_builtin, invoke_rpc_python_udf
2-
from . import invoke_remote_builtin, invoke_remote_python_udf
1+
from . import _invoke_rpc_builtin, _invoke_rpc_python_udf
2+
from . import _invoke_remote_builtin, _invoke_remote_python_udf
33
from . import _start_rpc_agent
44
from . import _destroy_rref_context, _cleanup_python_rpc_handler
55
from . import WorkerInfo
@@ -182,12 +182,12 @@ def remote(to, func, args=None, kwargs=None):
182182

183183
info = _to_worker_info(to)
184184
if qualified_name is not None:
185-
return invoke_remote_builtin(
185+
return _invoke_remote_builtin(
186186
_agent, info, qualified_name, *args, **kwargs)
187187
else:
188188
(pickled_python_udf, tensors) = _internal_rpc_pickler.serialize(
189189
PythonUDF(func, args, kwargs))
190-
return invoke_remote_python_udf(
190+
return _invoke_remote_python_udf(
191191
_agent, info, pickled_python_udf, tensors)
192192

193193

@@ -202,13 +202,13 @@ def _invoke_rpc(to, func, args=None, kwargs=None):
202202

203203
info = _to_worker_info(to)
204204
if qualified_name is not None:
205-
fut = invoke_rpc_builtin(
205+
fut = _invoke_rpc_builtin(
206206
_agent, info, qualified_name, *args, **kwargs
207207
)
208208
else:
209209
(pickled_python_udf, tensors) = _internal_rpc_pickler.serialize(
210210
PythonUDF(func, args, kwargs))
211-
fut = invoke_rpc_python_udf(
211+
fut = _invoke_rpc_python_udf(
212212
_agent, info, pickled_python_udf, tensors)
213213
return fut
214214

@@ -229,9 +229,10 @@ def rpc_sync(to, func, args=None, kwargs=None):
229229
invocation.
230230
231231
Returns:
232-
Returns the result of running ``func``on ``args`` and ``kwargs``.
232+
Returns the result of running ``func`` on ``args`` and ``kwargs``.
233233
234234
Example::
235+
235236
On worker 0:
236237
>>> import torch.distributed as dist
237238
>>> import torch.distributed.rpc as rpc

‎torch/distributed/rpc/backend_registry.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ def init_backend(backend, *args, **kwargs):
3939
return backend.value.init_backend_handler(*args, **kwargs)
4040

4141

42-
def process_group_init_backend_handler(
42+
def _process_group_init_backend_handler(
4343
store,
4444
self_name,
4545
self_rank,
@@ -91,4 +91,4 @@ def process_group_init_backend_handler(
9191

9292

9393

94-
register_backend("PROCESS_GROUP", process_group_init_backend_handler)
94+
register_backend("PROCESS_GROUP", _process_group_init_backend_handler)

0 commit comments

Comments
 (0)