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 [24].
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 [24].
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], 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 [20], etc.
Second, differentiable optimization is computation-intensive.
The meta-gradient computation requires heavy Hessian computation [11], high-dimensional linear equations [21], or large task-level batch size [19].
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. 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 let 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 θ ⋆ 𝜃𝜃𝜃𝜃𝜃𝜃′ 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 … 𝜑𝜑
Outer Parameter 𝜃𝜃
Inner Parameter
Forward Pass
Backward Pass 𝐿𝐿 𝑖𝑖𝑖𝑖 ( 𝜑𝜑 , 𝜃𝜃 ) :
Inner Loss F: Iterative Function 𝛻𝛻 𝜃𝜃 0 𝐿𝐿 𝑖𝑖𝑖𝑖 ( 𝜑𝜑 0 , 𝜃𝜃 0 )
[Figure]
over ϕ ) 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 over unrolled optimization (EG) (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 [24].
TorchOpt offers both functional and OOP API to fit different user applications.
Refer to Listing 1 for the code snippet.
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table] [Figure] [Table] for iter in
[Table] [Table]
Listing 1: TorchOpt Code Snippet for Explicit Gradient.
Implicit gradient (IG) (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 ( ∂F ( θ,ϕ ) ∂θ (cid:12)(cid:12)(cid:12) θ ⋆ = 0 ) or reaches some stationary conditions ( F ( θ, ϕ ) = 0 ), such as iMAML [21], DEQ [5].
TorchOpt offers functional/OOP API for supporting both conjugate gradient-based [21] and Neumann series [17] based IG method.
Refer to Listing 2 for the code snippet.
# 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 ... def solve(self, batch, labels): # Forward optimization process ... return self
Listing 2: TorchOpt code snippet for Implicit Gradient.
Zero-order Differentiation (ZD) (Fig. 2-2c).
When the inner-loop process is non-differentiable or one wants to eliminate the heavy computation burdens in the previous two modes (brought by
Hessian), one can choose ZD.
ZD typically gets gradients based on zero-order estimation, such as finite-difference, or Evolutionary Strategy (ES) [22].
ESMAML [23], and NAC [10], successfully solve the differentiable optimization problem based on ES.
Instead of optimizing the objective F , ES optimize a smoothed objective defined as ˜
f
σ ( θ )
=
E z ∼N ( 0 , I d ) [ f ( θ + σ z )] .
The gradient of such smoothed function is ∇ θ ˜ f
σ ( θ ) =
1 σ E z ∼N ( 0 , I d ) [ f ( θ + σ z ) z ] .
TorchOpt also offers functional and OOP API for the ES method.
Refer to Listing 3 for the code snippet.
[Table]
(params, batch, labels): # Forward process return output # OOP API class ESModule (torchopt.nn. ZeroOrderGradientModule): def sample(self, batch, labels, sample_shape): # Customize the noise sampling function in ES ... return sample_noise def forward(self, batch, labels): # Forward process ... return output
Listing 3: TorchOpt code snippet for Zero-order Differentiation.
Finally, TorchOpt offers Gradient Graph Visualization as a debugging tool.
Complex gradient flow in meta-learning/differentiable optimization brings a great challenge for managing the gradient flow 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 [25].
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 [Figure] using torch.autograd.
Function .
The results in Fig. 3(a) and Fig. 3(b) show that our design largerly reduce 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 supports distributed training based on the PyTorch RPC module, thus allowing users to reduce training time by adding multiple GPUs.
Different from the de-facto MPI-based synchronous training [18] or asynchronous model averaging [15], we propose to use the point-to-point communication: RPC as the communication backend for distributed differentiable optimization. The point-to-point communication is coordinated by a selected worker (i.e., a GPU device), thus enabling TorchOpt to achieve synchronous training (which guarantees the convergence of the model in a distributed training setting).
As shown in Appendix
C Fig.
6, TorchOpt distributes a differentiable optimization job across multiple workers and executes the workers in parallel.
Users can wrap their code with the distributed Autograd and obtain instant speedup with minimal code changes.
Fig. 3(c) presents the experimental results of MAML, which illustrates the linear relationship between the speed-up ratio and the number of workers.
Refer to Appendix C for more experimental results.
3. Conclusion and Future work
This paper presents TorchOpt a high-performance advanced optimization library for PyTorch.
TorchOpt provides a user-friendly and efficient framework to fill the vacancy of complex gradient computation problems in PyTorch.
In the future, TorchOpt aims to support more complex differentiation modes and cover more non-trivial gradient computation problems, such as the adjoint method for the gradient of ODE solution, RL or Gumbel-Softmax method for differentiating through discrete distribution, and differentiable combinatorial solver.
We believe TorchOpt can be served as an important auto-differentiation tool for related machine learning research and applications. Acknowledgments
We thank Vincent Moens for his constructive comments on TorchOpt .
Appendix A. Visualization
Fig.
4 show 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 optimizer
Fig.
5 shows the meta-optimization time comparison with Higher [13] in the CPU and GPU setting.
Note that the meta-optimization process consists of extra computation beyond optimizer, where we do not offer acceleration.
However, the acceleration is still significant (around %25 ) for 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 experimental result
In Fig.
7 we show the training accuracy and wall time comparison on 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 tree in different number of nodes comparing OpTree, JAX XLA, PyTorch, and DM-Tree.
[Table]
J
[Table]
O
P / O
D / O
J / O
P / O
D / O
J / O
P / O
D / O
J / O
P / O
D / O
[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]