1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
|
#pragma once
#include <torch/csrc/distributed/autograd/context/container.h>
#include <torch/csrc/distributed/autograd/engine/dist_engine.h>
namespace torch::distributed::autograd {
using torch::autograd::variable_list;
/// C++ API of Distributed Autograd that kicks off the distributed backward pass
/// using the provided roots. This currently implements the
/// :ref:`fast-mode-algorithm` which assumes all RPC messages sent in the same
/// distributed autograd context across workers would be part of the autograd
/// graph during the backward pass.
///
/// We use the provided roots to discover the autograd graph and compute
/// appropriate dependencies. This method blocks until the entire
/// autograd computation is done.
/// This function accumulates gradients in the leaves - you might need to zero
/// them before calling it.
///
/// \param context_id The autograd context id for which we should retrieve the
/// gradients.
/// \param roots Tensors which represent the roots of the autograd computation.
/// All the tensors should be scalars.
/// \param retain_graph If `false`, the graph used to compute the grad will be
/// freed. Note that in nearly all cases setting this
/// option to `true` is not needed and often can be worked
/// around in a much more efficient way. Usually, you need
/// to set this to `true` to run backward multiple times.
TORCH_API void backward(
int64_t context_id,
const variable_list& roots,
bool retain_graph = false);
} // namespace torch::distributed::autograd
|
