TorchOpt : An Efficient Library for Differentiable Optimization
Abstract
Recent years have witnessed the booming of various differentiable optimization algorithms.
These algorithms exhibit different execution patterns, and their execution needs massive computational resources that go beyond a single CPU and GPU.
Existing differentiable optimization libraries, however, cannot support efficient algorithm development and multi-CPU/GPU execution, making the development of differentiable optimization algorithms often cumbersome and expensive.
This paper introduces TorchOpt , a PyTorch-based efficient library for differentiable optimization.
TorchOpt provides a unified and expressive differentiable optimization programming abstraction.
This abstraction allows users to efficiently declare and analyze various differentiable optimization programs with explicit gradients, implicit gradients, and zero-order gradients.
TorchOpt further provides a high-performance distributed execution runtime.
This runtime can fully parallelize computation-intensive differentiation operations (e.g. tensor tree flattening) on CPUs / GPUs and automatically distribute computation to distributed devices.
Experimental results show that
TorchOpt achieves 5 .
2 × training time speedup on an 8-GPU server.
TorchOpt is available at: https://github.com/metaopt/torchopt .
1. Introduction
Recent years have witnessed the booming of differentiable optimization-based algorithms, including
MAML [11], OptNet [2] and MGRL [25].
One of the important parts of differentiable optimization is meta-gradient, which is the gradient term of outer-loop variables by differentiating through the inner-loop optimization process.
By leveraging meta-gradients, machine learning models can increase the sample efficiency [11] and the final performance [25].
Developing differentiable optimization algorithms poses several challenges.
First, developers need to realize different inner-loop optimization and implement algorithms with gradient flows on complex computational graphs.
Examples include explicit gradient computation of unrolled optimization [11, 25], implicit gradient for differentiable optimization [1, 2], evolutionary strategies for non-differentiable optimization [10], adjoint methods for differentiable ordinary differentiable equations [7], Gumbel-Softmax to differentiate through discrete distribution [14], and function interpolation for differentiable combinatorial solvers [21], etc.
Second, differentiable optimization is computation-intensive.
The meta-gradient computation requires heavy Hessian computation [11], high-dimensional linear equations [22], or large task-level batch size [20].
Such computation requirement often goes beyond what a single CPU and GPU can provide.
None of the existing differentiable optimization libraries can provide full support for efficient algorithm development and execution.
Most libraries target a limited number of differentiable optimizers [3, 4, 9, 13].
They cannot fully support implicit differentiation [6, 8, 12], zero-order gradient [8], and distributed training [3, 6], as shown in Table 1.
As a result, researchers have to implement algorithms in an ad-hoc and application-specific manner, making the development process cumbersome and expensive.
Further, critical system optimization techniques (e.g. GPU optimization and distributed execution) are tightly coupled with certain algorithms and they are hard to be enabled for all possible algorithms.
To address these issues, this paper introduces TorchOpt , a PyTorch library that makes it efficient to develop and execute differentiable optimization algorithms with multiple GPUs.
The design and implementation of TorchOpt make the following contributions:
(1) Unified and expressive differentiation mode for differentiable optimization.
TorchOpt provides a general set of low-level / high-level / functional / Object-Oriented (OO) API to help users flexibly enable differentiable optimization within the computational graphs produced by PyTorch.
Specifically, TorchOpt supports three differentiation modes for handling differentiable optimization problems: (i) Explicit gradient for unrolled optimization, (ii) implicit gradient for differentiable optimization, and (iii) zero-order gradient estimation for non-smooth/differentiable functions.
(2) High-performance and distributed execution runtime
TorchOpt aims to enable differentiable optimization algorithms to fully utilize CPUs and GPUs.
To achieve this, we design (i)
CPU/GPU accelerated optimizers (e.g., SGD, RMSProp, Adam) that realize the fusion of small differentiable operators and a full offloading of these operators to GPUs, (ii) parallel OpTree which can fully parallelize the nested structure flattening (Tree Operations), a key computation-intensive operation in differentiable optimization, on distributed CPUs, and (iii) a distributed auto-grad framework which can automatically identify the inner-loop tasks in differentiable optimizers and dispatch the execution of inner-loop tasks to distributed CPUs and GPUs.
Experimental results show that TorchOpt can reduce PyTorch optimizer forward/backward time, by 5 × to 10 × on CPU and 5 × to 20 × on GPU.
TorchOpt can reduce the training time of the MAML [11] algorithm by 5 . 2 × by distributing MAML computation to 8 GPUs.
2. TorchOpt Design and Implementation 2.1. Architecture Overview
Figure 1 gives an overview of the system architecture, TorchOpt consists of two different aspects, the unified and expressive differentiable optimization programming lets users easily implement differentiable optimization algorithms, we provide both high-level APIs and low-level APIs for three differentiation modes along with debugging tools, all of which are described in Sec. 2.2.
Then the high-performance and distributed execution runtime contains several accelerated solutions to support fast differentiation with different modes on GPU & CPU and distributed training features
[Table]
for multi-node multi-GPU scenario, which we demonstrate boost performance in Sec. 2.3.
Additionally, we offer OpTree to enable fast structure flatten and unflatten , which is specially designed for our functional programming implementation.
We use an optimized structure to avoid memory allocation if the sub-tree is small.
2.2. Programming Abstraction
TorchOpt aims to provide (i) high-level APIs that allow users to directly import differentiable optimizers, (ii) low-level APIs that enable automatic differentiation in different applications, and (iii) tools that allow users to analyze gradient flow in gigantic computational graphs.
The key challenge of consolidating these high-level and low-level APIs in a single library is that we must have a unified abstraction that allows different differentiable optimization algorithms to be easily declared.
To address this, we design a differentiable optimization updating scheme, which can be easily extended to realize various differentiable optimization processes.
As shown in Fig.
2, the scheme contains an outer level that has parameters 𝝓 that can be learned end-to-end through the inner level parameters solution 𝜽 ′ ( 𝝓 ) (treating solution 𝜽 ′ as a function of 𝜃𝜃𝜃𝜃𝜃𝜃′
Unrolled Inner Loop optimization 𝛻𝛻 𝜃𝜃 1 𝐿𝐿 𝑖𝑖𝑖𝑖 ( 𝜑𝜑 0 , 𝜃𝜃 1 ) 𝜑𝜑𝜃𝜃′ 𝜃𝜃
Non-trivial Transformation
Iterative optimization/ Fixed-point iteration
General Outer-loop 𝜃𝜃𝜃𝜃𝜃𝜃′ 𝜃𝜃′ 𝜃𝜃𝜃𝜃𝜃𝜃 𝑖𝑖
F F
F
Non-smooth/ Non-Differentiable Function Non-diff process … 𝜑𝜑
[Figure]
Parameter 𝜃𝜃
[Figure]
Parameter
Forward Pass Backward Pass 𝐿𝐿 𝑖𝑖𝑖𝑖 ( 𝜑𝜑 , 𝜃𝜃 ) :
Inner Loss F: Iterative Function 𝛻𝛻 𝜃𝜃 0 𝐿𝐿 𝑖𝑖𝑖𝑖 ( 𝜑𝜑 0 , 𝜃𝜃 0 )
[Figure]
𝝓 ) by using the best-response derivatives 𝜕 𝜽 ′ ( 𝝓 )/ 𝜕 𝝓 .
It can be seen that the key component of this algorithm is to calculate the best-response (BR) Jacobian.
From the BR-based perspective,
TorchOpt supports three differentiation modes: explicit gradient over unrolled optimization, implicit differentiation, and zero-order differentiation.
Explicit Gradient (EG) over unrolled optimization.
As shown in Fig.
2-2a, the idea of EG is to treat the gradient step as a differentiable function and try to backpropagate through the unrolled optimization path.
This differentiation mode is suitable for algorithms when the inner-level optimization solution is obtained by a few gradient steps, such as MAML [11], MGRL [25].
TorchOpt offers both functional and OOP API.
Refer to Listing 1 for the code snippet.
[Figure] API opt = torchopt.adam()
# Define meta and
[Table]
parameters meta
[Table] = ... fmodel, params = make_functional(model) # Initialize [Table] state state = opt.init(params) for iter in range(iter_times): loss = inner_loss(fmodel, params,
[Table]
Listing 1: TorchOpt code snippet for explicit gradient.
Implicit Gradient (IG).
As shown in Fig.
2-2b, by treating the solution 𝜽 ′ as an implicit function of 𝝓 , the idea of IG is to directly get analytical best-response derivatives 𝜕 𝜽 ′ ( 𝝓 )/ 𝜕 𝝓 by implicit function theorem [16].
This is suitable for algorithms when the inner-level optimal solution is achieved ( 𝜕𝐹 ( 𝜽 , 𝝓 ) 𝜕 𝜽 (cid:12)(cid:12)(cid:12) 𝜽 ′ = 0 ) or reaches some stationary conditions ( 𝐹 ( 𝜽 ′ , 𝝓 ) = 0 ), such as iMAML [22], DEQ [5].
TorchOpt offers functional/OOP API for supporting both conjugate gradient-based [22] and Neumann series [18] based method.
Refer to Listing 2 for the code snippet.
Zero-order Differentiation (ZD).
As shown in Fig.
2-2c, when the inner-loop process is nondifferentiable or one wants to eliminate the heavy computation burdens in the previous two modes
# Functional API for implicit gradient def stationary(params, meta_params, batch, labels): # Stationary condition construction ... return stationary condition @torchopt.diff.implicit.custom_root(stationary) def solve(params, meta_params, batch, labels): # Forward optimization process ... return optimal_params # OOP API class Module (torchopt.nn.
ImplicitMetaGradientModule): def __init__(self, meta_module, ...): ... def forward(self, x): # Forward process ... def optimality(self, batch, labels):
# Stationary condition construction ..
[Table]
[Table]
[Table] , labels): #
[Table]
optimization
[Table]
self Listing 2: TorchOpt code snippet for implicit gradient.
(brought by Hessian), one can choose ZD.
ZD typically gets gradients based on zero-order estimation, such as finite-difference, or Evolutionary Strategy (ES) [23].
ESMAML [24], and NAC [10], successfully solve the differentiable optimization problem based on ES.
Instead of optimizing the objective 𝐹 , ES optimize a smoothed objective defined as ˜ 𝑓 𝜎 ( 𝜽 )
=
𝔼 𝒛 ∼N( 0 , 𝑰 𝑑 ) [ 𝑓 ( 𝜽 + 𝜎 𝒛 )] , where 𝜎 denotes precision.
The gradient of such smoothed function is ∇ 𝜃 ˜ 𝑓 𝜎 ( 𝜃 ) = 1 𝜎 𝔼 𝒛 ∼N( 0 , 𝑰 𝑑 ) [ 𝑓 ( 𝜽 + 𝜎 𝒛 ) 𝒛 ] .
TorchOpt also offers functional and OOP API for the ES method.
Refer to Listing 3 for the code snippet.
[Table]
return output
[Table]
[Table]
return output
Listing 3: TorchOpt code snippet for zero-order differentiation.
Gradient graph visualization .
Complex gradient graph in meta-learning/differentiable optimization brings a great challenge for managing the gradient graph and debugging the code.
TorchOpt provides a visualization tool that draws variable (e.g. network parameters or meta parameters) names on the gradient graph for better analysis.
The visualization tool is modified from TorchViz [26].
Compared with TorchViz, TorchOpt fuses the operations within the optimization algorithm (such as Adam) to reduce the complexity and provide simpler visualization.
Refer to the visualization example in Appendix A.
2.3. High-performance and Distributed Runtime
CPU/GPU-accelerated optimizers.
We take the optimizer as a whole instead of separating it into several basic operators (e.g., sqrt and div ).
Therefore, by manually writing the forward and backward functions, we can perform the symbolic reduction.
In addition, we can store some intermediate data that can be reused during the back-propagation.
Our design reduces computation and also benefits numerical stability (by explicitly canceling some 0 / 0 cases in higher gradient computation).
We write the accelerated functions in C++ OpenMP and CUDA, bind them by pybind11 to allow Python can call them, and then we define the forward and backward behavior using torch.autograd.
Function .
The results in Fig. 3(a) and Fig. 3(b) show that our design
[Figure]
largely reduces the optimizer forward and backward time.
Refer to Appendix B for experimental results comparing TorchOpt and Higher [13] on the MAML example.
Memory-efficient and cache-friendly PyTree.
The tree operations (e.g., flatten and unflatten) are frequently called by the functional and Just-In-Time (JIT) components in TorchOpt .
To enable memory-efficient nested structure flattening, we implement the pytree utilities, named OpTree .
By optimizing their memory and cache performance (e.g., absl::InlinedVector ), TorchOpt can significantly improve the performance of differentiable optimization at scale.
Refer to Appendix
D for OpTree experimental results.
Distributed differentiable optimization.
TorchOpt allows users to reduce training time by using parallel GPUs.
Different from existing MPI-based synchronous training [19] and asynchronous model averaging [15] systems, TorchOpt adopts RPC as a flexible yet performance communication backend.
The distributed GPUs perform differentiable optimization tasks in parallel.
These
GPUs are coordinated by a chosen GPU device which realizes the synchronous execution of parallel
GPUs, thus guaranteeing the convergence of the model in a distributed training setting).
As shown in Fig.
6, TorchOpt distributes a differentiable optimization job across multiple
GPU workers and executes the workers in parallel.
TorchOpt users can wrap code in the distributed Autograd module and achieve substantial speedup in training time with only a few changes in existing training scripts.
Fig. 3(c) shows that TorchOpt can achieve linear speed-up with
MAML when increasing the number of GPU workers (more details in Appendix C).
3. Conclusion and Future work
This paper introduces TorchOpt , a novel efficient differentiable optimization library for PyTorch.
Experimental results show that TorchOpt can act as a user-friendly, high-performance, and scalable library when supporting challenging gradient computation with PyTorch.
In the future, we aim to support more complex differentiation modes and cover more non-trivial gradient computation problems, such as adjoint methods for the gradient of ODE solutions, RL or Gumbel-Softmax method for differentiating through discrete distribution, and differentiable combinatorial solvers.
TorchOpt has already been used for meta-gradient research problem [17] and we believe it can be served as an important auto-differentiation tool for more differentiable optimization problems.
Acknowledgments
We thank Vincent Moens for his constructive comments on TorchOpt .
Appendix A. Gradient Graph Visualization
Fig.
4 shows the visualization example of MAML.
We use red squares to represent what each part accomplishes separately.
Compared with TorchViz, TorchOpt fuses the operations within the
Adam together (orange) to reduce the complexity and provides a more straightforward visualization.
MseLossBackward
AddBackward
AddmmBackward
AddBackward0step1.fc.bias(1)
AccumulateGrad
AddmmBackward step0.fc.bias(1)
MulBackward
UpdatesOpBackward
MuOpBackward
ViewBackward
NuOpBackward
SumBackward
MseLossBackwardBackward
TBackward
AddBackward
TBackward
AccumulateGrad
AddBackward0step1.fc.weight(1, 5) step0.fc.weight(1, 5)
AccumulateGrad ()
AccumulateGrad (1) (1)
AccumulateGrad (1) (1) (1) (1)
TBackward
MulBackward
UpdatesOpBackward
MuOpBackward
TBackward
NuOpBackward
TBackward
MmBackward
AccumulateGrad (1, 5) (1, 5)
AccumulateGrad (1, 5) (1, 5) (1, 5) (1, 5) ()
MseLossBackward0 AddBackward
AddmmBackward
AddBackward
AccumulateGrad
AddmmBackward (1)
MulBackward0 DivBackward
DivBackward
AddBackward
MulBackward
AccumulateGrad (1)
ViewBackward
AddcmulBackward
SumBackward
MseLossBackwardBackward
TBackward
AddBackward
TBackward
AccumulateGrad
AddBackward (1, 5)
AccumulateGrad ()
AddBackward0 SqrtBackward
AddBackward0 DivBackward
MulBackward
AccumulateGrad
[Figure]
Forward
TorchOpt 2nd
Derivative 2nd
Derivative
OptimStep OptimStep
Forward
Appendix B. CPU/GPU-Accelerated Optimizers
Fig.
5 shows the meta-optimization time comparison with Higher [13] in the CPU and GPU settings.
Note that the meta-optimization process consists of extra computation beyond the optimizer, where we do not offer acceleration.
However, the acceleration is still significant (around %25 ) for the MLP model in the CPU setting and both Conv/MLP model in the GPU setting.
Appendix C. Distributed Training C.1. Distributed Framework
In Fig.
6 we show the overview of our distributed framework.
Inner Steps 0.
T i m e ( s )
Higher/Conv
Higher/MLP
TorchOpt/Conv
TorchOpt/MLP
[Figure]
Inner Steps 0.
T i m e ( s )
Higher/Conv
Higher/MLP
TorchOpt/Conv
TorchOpt/MLP
[Figure]
C.2. Distributed MAML Performance
In Fig.
7, we show the training accuracy and wall time comparison on the MAML
Omniglot example.
Distributed training achieves better performance and much higher computational efficiency.
[Figure]
Appendix D. OpTree Performance
In Table.
2 we show the Speedup ratios of tree operations with ResNet models comparing OpTree,
JAX XLA, PyTorch, and DM-Tree.
In Fig.
8, 9 and 10, we show the time cost of tree-flatten, treeunflatten, and tree-map trees in a different number of nodes comparing OpTree, JAX XLA, PyTorch, and DM-Tree.
OpTree achieves a large speedup compared with all baselines.
[Table]
J / O P / O
D / O
J / O P / O
D / O
J / O P / O
D / O
J / O P / O
D / O
Tree
[Table]
[Figure]
A l e x N e t
R e s N e
[Figure]
N
[Figure] [Figure]
[Figure]
T i m e ( m s )
A l e x N e t
R e s N e t R e s N e t
R e s N e t
S w i n - B
R e s N e t
[Figure]
V i T - H /
[Figure]
R e s N e t
[Figure]