Fusion system for automatic activation value checkpoint search and automatic tensor parallel search
By integrating automatic activation value checkpoint search and tensor parallel search, the memory usage of large-scale deep learning models is optimized, solving the problems of memory pressure and long solution time, and achieving more efficient model training.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- BEIJING LUCHEN TECH CO LTD
- Filing Date
- 2022-12-31
- Publication Date
- 2026-06-09
AI Technical Summary
Existing automatic activation value checkpoint search and tensor parallel search systems have failed to effectively combine, resulting in inaccurate memory estimation, poor optimization results, and excessively long solution times, making it difficult to cope with the memory pressure problem of large-scale deep learning models.
A fusion system employing automatic activation value checkpoint search and automatic tensor parallel search is proposed. Through a two-stage solver consisting of network linearization, metadata extraction, automatic activation value checkpoint solver, and automatic tensor parallel solver, the computational graph is optimized by combining metadata and memory budget to reduce memory overhead and runtime.
It effectively mitigates the memory wall problem in deep learning training, improves model training performance, and ensures reasonable memory usage and computational efficiency.
Smart Images
Figure CN116050512B_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the field of deep learning technology, specifically to a fusion system of automatic activation value checkpoint search and automatic tensor parallel search. Background Technology
[0002] As deep learning models grow larger to achieve higher accuracy, the number of parameters in these models has reached tens or even hundreds of billions. The Zhiyuan Wudao 2.0 model even boasts 1.75 trillion parameters.
[0003] In this context, employing various techniques to alleviate GPU memory pressure is crucial for training large deep learning models. Activation value checkpointing and tensor parallelism are two effective methods, and many related works have also presented policy searches based on dynamic programming and integer programming to provide better optimization strategies. However, no system has yet organically combined these two approaches.
[0004] The current mainstream automatic activation value checkpoint search schemes are Rotor and Checkmate.
[0005] rotor: This search system provides a complete framework for automatic activation value checkpoint search of linearized networks.
[0006] System Assumptions
[0007] Network linearization: This system assumes that the network can be viewed as a linearized sequence of executions, meaning that during forward propagation, each network layer is only related to the output of the previous layer. Memory consistency: Activation values that are determined to be saved are retained in memory and will not be discarded midway before being used in reverse computation.
[0008] System Process
[0009] The overhead of each layer in the computation sequence was obtained through testing (memory overhead from storing activation values, computation time).
[0010] Under the linearization assumption, discretizing the memory overhead allows us to model the arrangement of optimal activation checkpoints as a dynamic programming problem solvable in polynomial time. The paper can be found at: https: / / proceedings.neurips.cc / paper / 2021 / hash / c8461bf13fca8a2b9912ab2eb1668e4b-Abstract.html
[0011] checkmate: This system provides the optimal solution for activation checkpoints in any network (directed acyclic graph). The system assumes...
[0012] The network is an arbitrary directed acyclic graph, which blurs the decomposition of forward and backward propagation and considers all computations holistically.
[0013] The system flow involves testing to obtain the overhead of each layer in the computational sequence (memory overhead and computation time for storing activation values). Considering the entire computational graph and the sequence of activation value checkpoints, the problem can be modeled as a standard integer programming problem and solved using a standard integer programming solver.
[0014] Current mainstream automatic tensor parallel search solutions include Alpa
[0015] Alpa: This system provides detailed modeling for automatic tensor parallelism and automatic pipelined parallelism. Given the cost model of the network, it uses integer programming to solve for the optimal strategy.
[0016] Paper link: https: / / arxiv.org / pdf / 2201.12023.pdf
[0017] Flexflow: Flexflow is a system in SOAP space that automatically searches for the optimal parallel strategy based on the execution simulator cost model for any DNN model and device topology.
[0018] However, Rotor, based on real tensor computation, requires significant time to estimate memory and computation time, sometimes resulting in out-of-memory errors. The memory estimation is not very accurate, leading to problems in the final optimization results. Existing linearization requires manual model tuning, which is inconvenient. Linearization significantly reduces the search space.
[0019] checkmate: Calculates node costs using actual tensors. Integer programming results in very long solution times.
[0020] Alpa: It does not consider automatic activation value checkpoint search and does not have a PyTorch version. To address this, we propose a fusion system of automatic activation value checkpoint search and automatic tensor parallel search. Summary of the Invention
[0021] To address the shortcomings of existing technologies, this invention provides a fusion system of automatic activation value checkpoint search and automatic tensor parallel search, which solves the problems mentioned in the background.
[0022] To achieve the above objectives, this invention provides the following technical solution: a fusion system of automatic activation value checkpoint search and automatic tensor parallel search, comprising network linearization, metadata extraction, an automatic activation value checkpoint solver, an automatic tensor parallel solver, a two-stage solver, and an optimized computational graph.
[0023] It also includes the following steps:
[0024] S1: Initialize the model.
[0025] S2: Use meta-information extraction to provide the necessary memory and computational overhead information for the automatic tensor parallel solver;
[0026] S3: The automatic tensor parallel solver solves the problem based on the given memory budget and metadata.
[0027] S4: The two-stage solver linearizes a series of tensor parallel strategies into a network and passes the linearized network to the automatic activation value checkpoint solver.
[0028] S5: The automatic activation value checkpoint solver keeps memory overhead below the actual memory budget, obtains a series of strategies, and selects the strategy with the shortest estimated running time.
[0029] S6: Optimize the computation graph based on the given composite strategy, then compile and execute.
[0030] Preferably, the network linearization includes the following steps:
[0031] Meta-tracing is used to obtain the computational graph of the network, finding general-purpose nodes in the graph based on user annotations (e.g., attention masks for language models) and node characteristics.
[0032] Based on node dependencies (ignoring general nodes), obtain the partition points of the linearized network computation graph.
[0033] Preferably, the metadata extraction, in the absence of actual execution, uses PyTorch's __torch_dispatch__ mechanism to obtain the underlying operators of the tensor when it is run on the desired device (GPU, CPU), and thereby estimate the runtime memory overhead and computational load.
[0034] Use the __torch_dispatch__ mechanism to obtain the underlying operators of each operator;
[0035] Based on the acquired operators, the computational and memory overhead of the model is abstracted into mathematical formulas that depend on the shapes of the input and output, facilitating the automatic tensor parallel acquisition of operators under different partitioning methods.
[0036] Preferably, the automatic activation value checkpoint solver is based on the Rotor algorithm, with modified modeling to consider tensor parallel multi-card communication scenarios, and provides the optimal solution considering tensor parallel communication overhead within the user-input memory limit.
[0037] Preferably, the automatic tensor parallel solver is an automatic tensor parallel solver built on PyTorch based on Alpa, which performs tensor parallel strategy solving given meta-information.
[0038] Preferably, the two-stage solver includes the following steps: given an automatic tensor parallel solver and an automatic activation value checkpoint solver, inputting a set of virtual memory budgets into the automatic tensor parallel solver based on the user's actual available memory budget, providing a series of tensor parallel strategies, then using the automatic activation value checkpoint solver to provide activation value checkpoint strategies to control memory below the actual memory budget, comparing a series of strategies, and selecting the strategy with the shortest estimated running time.
[0039] Preferably, the optimized computation graph reconstructs the original PyTorch computation graph based on the optimal solution of the dynamic programming solver, generating a new computation graph so that the activation value checkpoints obtained by automatic search are ultimately used, thereby significantly reducing training memory overhead.
[0040] This invention provides a fusion system of automatic activation value checkpoint search and automatic tensor parallel search, which has the following beneficial effects:
[0041] 1. This fusion system of automatic activation value checkpoint search and automatic tensor parallel search provides the automatic tensor parallel solver with the necessary memory and computational overhead information through metadata extraction. The automatic tensor parallel solver solves the problem based on the given memory budget and metadata. The two-stage solver linearizes a series of tensor parallel strategies into a network and passes the linearized network to the automatic activation value checkpoint solver. The automatic activation value checkpoint solver controls the memory overhead within the actual memory budget, obtains a series of strategies, selects the strategy with the shortest estimated running time, optimizes the computation graph according to the given composite strategy, compiles and executes the solution, and combines the strategies of automatic tensor parallel search and automatic activation value checkpoint search to further alleviate the memory wall problem encountered in deep learning training. At the same time, the two-stage solver can also ensure that the memory problem is solved while maintaining the performance of model training. Attached Figure Description
[0042] Figure 1 This is a schematic diagram of the structure of the present invention;
[0043] Figure 2 This is a schematic diagram of the network linearization structure of the present invention;
[0044] Figure 3This is a schematic diagram of the structure for extracting metadata in this invention;
[0045] Figure 4 This is a schematic diagram of the two-stage solver of the present invention. Detailed Implementation
[0046] The technical solutions of the present invention will be clearly and completely described below with reference to the accompanying drawings of the embodiments of the present invention. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments.
[0047] Please see Figures 1 to 4 This invention provides a technical solution: 1. A fusion system of automatic activation value checkpoint search and automatic tensor parallel search, characterized in that it includes network linearization, metadata extraction, an automatic activation value checkpoint solver, an automatic tensor parallel solver, a two-stage solver, and an optimized computation graph.
[0048] It also includes the following steps:
[0049] S1: Initialize the model.
[0050] S2: Use meta-information extraction to provide the necessary memory and computational overhead information for the automatic tensor parallel solver;
[0051] S3: The automatic tensor parallel solver solves the problem based on the given memory budget and metadata.
[0052] S4: The two-stage solver linearizes a series of tensor parallel strategies into a network and passes the linearized network to the automatic activation value checkpoint solver.
[0053] S5: The automatic activation value checkpoint solver keeps memory overhead below the actual memory budget, obtains a series of strategies, and selects the strategy with the shortest estimated running time.
[0054] S6: Optimize the computation graph based on the given composite strategy, then compile and execute.
[0055] Combining the strategies of automatic tensor parallelism and automatic activation value checkpointing further mitigates the memory wall problem encountered in deep learning training. At the same time, the two-stage solver can ensure that the memory problem is solved while maintaining the performance of model training.
[0056] Network linearization includes the following steps:
[0057] Meta-tracing is used to obtain the computational graph of the network, finding general-purpose nodes in the graph based on user annotations (e.g., attention masks for language models) and node characteristics.
[0058] Based on node dependencies (ignoring general nodes), obtain the partition points of the linearized network computation graph.
[0059] In the case of non-real execution, the metadata extraction uses PyTorch's __torch_dispatch__ mechanism to obtain the underlying operators of the tensor when it is run on the desired device (GPU, CPU), and thereby estimate the runtime memory overhead and computational load.
[0060] Use the __torch_dispatch__ mechanism to obtain the underlying operators of each operator;
[0061] Based on the acquired operators, the computational and memory overhead of the model is abstracted into mathematical formulas that depend on the shapes of the input and output, facilitating the automatic tensor parallel acquisition of operators under different partitioning methods.
[0062] The automatic activation value checkpoint solver is based on the Rotor algorithm, with modified modeling to consider tensor parallel multi-card communication scenarios. Under the user-input memory limit, it provides the optimal solution considering the overhead of tensor parallel communication.
[0063] The automatic tensor parallel solver is built on PyTorch based on Alpa, and performs tensor parallel solutions given meta-information.
[0064] The two-stage solver includes the following steps: given an automatic tensor parallel solver and an automatic activation value checkpoint solver, a set of virtual memory budgets is input into the automatic tensor parallel solver based on the user's actual available memory budget, a series of tensor parallel strategies are given, and then the automatic activation value checkpoint solver is used to give activation value checkpoint strategies to control memory below the actual memory budget. The series of strategies are compared, and the strategy with the shortest estimated running time is selected.
[0065] The optimized computation graph is based on the optimal solution of the dynamic programming solver. The original PyTorch computation graph is reconstructed to generate a new computation graph, so that the activation value checkpoints obtained by automatic search are finally used, thereby greatly reducing the training memory overhead.
[0066] Example 2:
[0067] Taking the composite strategy search of OPT as an example:
[0068] First, an OPT model needs to be created, and for this model, a set of input tensors in the form of meta-tensors need to be created.
[0069] Based on the input meta-tensor, the automatic tensor parallel solver can use meta-information extraction to obtain the memory and computational overhead of operators under various partitioning schemes. Based on this information, and given a set of virtual memory budgets, the automatic tensor parallel solver solves for a series of tensor parallel strategies.
[0070] After obtaining a series of tensor-parallel strategies, the two-stage solver linearizes the computation graph corresponding to these strategies and hands it over to the automatic activation value checkpoint solver. The automatic activation value checkpoint solver provides corresponding activation value checkpoint strategies based on this series of linearized networks, keeping the estimated runtime memory below the actual memory overhead specified by the user.
[0071] Ultimately, the solver selects the strategy with the shortest estimated runtime from a series of costs, and compiles and executes the computation graph according to this strategy.
[0072] In summary, this fusion system of automatic activation value checkpoint search and automatic tensor parallel search, when in use, initializes the model and uses metadata extraction to provide the automatic tensor parallel solver with the necessary memory and computational overhead information. The automatic tensor parallel solver solves the problem based on the given memory budget and metadata. The two-stage solver linearizes a series of tensor parallel strategies into a network and passes the linearized network to the automatic activation value checkpoint solver. The automatic activation value checkpoint solver controls the memory overhead to within the actual memory budget, obtains a series of strategies, selects the strategy with the shortest estimated running time, optimizes the computational graph according to the given composite strategy, and then compiles and executes the solution.
[0073] The above description is only a preferred embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any equivalent substitutions or modifications made by those skilled in the art within the scope of the technology disclosed in the present invention, based on the technical solution and inventive concept of the present invention, should be covered within the scope of protection of the present invention.
Claims
1. A fusion system of automatic activation value checkpoint search and automatic tensor parallel search, characterized in that: This includes network linearization, metadata extraction, automatic activation value checkpoint solver, automatic tensor parallel solver, two-stage solver, and optimized computation graph. It also includes the following steps: S1: Initialize the model. S2: Use meta-information extraction to provide the necessary memory and computational overhead information for the automatic tensor parallel solver; S3: The automatic tensor parallel solver solves the problem based on the given memory budget and metadata. S4: The two-stage solver linearizes a series of tensor parallel strategies into a network and passes the linearized network to the automatic activation value checkpoint solver. S5: The automatic activation value checkpoint solver keeps memory overhead below the actual memory budget, obtains a series of strategies, and selects the strategy with the shortest estimated running time. S6: Optimize the computation graph based on the given composite strategy, compile and execute; The network linearization includes the following steps: Meta-tracing is used to obtain the computational graph of the network and to find general-purpose nodes in the graph based on user annotations and node characteristics. Based on the dependencies between nodes, obtain the partition points of the linearized network computation graph. The metadata extraction, without actual execution, uses PyTorch's __torch_dispatch__ mechanism to obtain the underlying operators of the tensor during runtime on the device, and thereby estimates the runtime memory overhead and computational load. Use the __torch_dispatch__ mechanism to obtain the underlying operators of each operator; Based on the acquired operators, the computational and memory overhead of the model is abstracted into mathematical formulas that depend on the shapes of the input and output, facilitating the automatic tensor parallel acquisition of operators under different partitioning methods. The automatic activation value checkpoint solver is based on the Rotor algorithm, with modified modeling to consider tensor parallel multi-card communication scenarios. Within the user-input memory limit, it provides the optimal solution considering the overhead of tensor parallel communication. The automatic tensor parallel solver is built on PyTorch based on Alpa, and performs tensor parallel solution under given meta-information. The two-stage solver includes the following steps: given an automatic tensor parallel solver and an automatic activation value checkpoint solver, a set of virtual memory budgets is input to the automatic tensor parallel solver based on the user's actual available memory budget, a series of tensor parallel strategies are given, and then the automatic activation value checkpoint solver is used to give activation value checkpoint strategies to control memory below the actual memory budget. The series of strategies are compared, and the strategy with the shortest estimated running time is selected.
2. The fusion system of automatic activation value checkpoint search and automatic tensor parallel search according to claim 1, characterized in that: The optimized computation graph reconstructs the original PyTorch computation graph based on the optimal solution of the dynamic programming solver, generating a new computation graph so that the activation value checkpoints obtained by automatic search are finally used, thereby significantly reducing training memory overhead.