Tensor program generation method and system based on large language model

Abstract

The invention discloses a tensor program generation method and system based on a large language model, and the method comprises the following steps: (1) generating a random ONNX computational graph through an AI compiler test tool nnsmith; (2) importing an ONNX calculation graph generated by nnsmith into a TVM, and converting the ONNX calculation graph into a RelayIR intermediate representation; (3) based on a model intermediate representation RelayIR in the TVM, adjusting and optimizing the model through an automatic adjusting and optimizing device Ansor in the TVM, generating an optimal tensor program, and outputting a JSON file measured during adjusting and optimizing; (4) continuously carrying out model tuning and tensor program measurement, generating a large number of measurement records during tuning, and forming a tensor program data set which comprises features of a tensor program and a corresponding optimization pass list; and (5) training a large language model according to the generated tensor program data set to realize understanding and generation of the tensor program. Through combination of an automatic tool and a large language model, efficient tensor program generation and optimization can be realized.

Description

Technical Field

[0001] The present invention belongs to the technical field of deep learning compilation and optimization, and in particular, to a tensor program generation method and system based on a large language model. Background Art

[0002] The training and reasoning process of deep learning models requires a large number of tensor calculations, and how to perform these calculations efficiently is an important issue. At present, deep learning compilers and optimization tools can improve the execution efficiency of models to a certain extent, but the manual optimization process is time-consuming and labor-intensive, and it is difficult to achieve optimal performance on different hardware platforms.

[0003] The training and reasoning process of deep learning models requires a large number of tensor calculations, and how to perform these calculations efficiently is an important issue. At present, deep learning compilers and optimization tools can improve the execution efficiency of models to a certain extent, but the manual optimization process is time-consuming and labor-intensive, and it is difficult to achieve optimal performance on different hardware platforms.

[0004] First, handwritten operator acceleration libraries (such as cuDNN) can provide efficient computing acceleration on specific hardware, but they lack versatility and flexibility. For example, handwritten operators often require a lot of reconstruction and adjustment for different deep learning models and data structures, resulting in low development efficiency. In addition, the use of handwritten operators also requires developers to have high professional knowledge to fully utilize the potential of the hardware, which limits its promotion in a wide range of application scenarios.

[0005] Secondly, although polyhedral compilation technology (such as the Polyhedral model) can optimize loops and memory access patterns through mathematical modeling to improve computing performance, its implementation complexity is relatively high and requires developers to have a high mathematical background. In addition, it has poor adaptability to different hardware architectures, especially in heterogeneous computing environments, and the effect of polyhedral compilation is often not as good as expected.

[0006] In terms of deep learning compilers, automated tuning tools such as AutoTVM and Ansor can automatically generate optimal operator implementations through search algorithms, thereby improving the execution efficiency of the model. However, these tools may not achieve satisfactory tuning results when the search space is large and the time cost is high, especially when facing complex models, the tuning results may not necessarily achieve the same results as manual optimization. In addition, the tuning process of AutoTVM requires a large number of training samples to guide the search, which may lead to inefficiency in some cases.

[0007] As deep learning compilation frameworks, TensorRT and MLIR each have their own advantages and disadvantages. TensorRT excels in inference performance and can be deeply optimized for NVIDIA hardware, but it is highly dependent on specific hardware and supports limited model frameworks. MLIR adopts a modular design concept and supports a variety of hardware platforms and models, but due to its relatively new technology stack, the developer community and ecosystem are not yet mature, resulting in some features not yet fully stable.

[0008] In addition, other technologies for deep learning compilers, such as one-shot tuner, adatune, and roller, provide different tuning strategies, but in practical applications they often face problems such as low search efficiency and complex implementation. Especially when migrating between large-scale models and multiple hardware platforms, the tuning effects of these tools become uncertain, and additional manual intervention is often required to ensure performance improvement.

[0009] AdaTune is a tuning tool based on adaptive learning. It dynamically adjusts hyperparameters by analyzing the characteristics of a specific model and performance feedback at runtime. The core idea of ​​AdaTune is to use machine learning methods to predict the optimal configuration, thereby reducing the workload of manual parameter tuning. The tool can learn and adjust in real time to adapt to the performance of different models on different hardware. However, the performance of AdaTune depends on the initial training samples and the characteristics of the model. It may not be stable enough for complex models, and the tuning effect may not be as expected when facing novel models.

[0010] Roller emphasizes a systematic automatic parameter tuning strategy. It divides the tuning process into multiple stages and combines heuristic algorithms and search strategies to gradually optimize the performance of the model. Roller can perform tuning on a variety of hardware platforms and has designed a gradient-based optimization method to improve search efficiency. Although Roller performs well in a variety of scenarios, its complex implementation and tuning process can still lead to high computational costs, especially when dealing with large-scale models.

[0011] One-Shot Tuner is an innovative tuning method that quickly finds the optimal configuration by training multiple candidate models at once and evaluating their performance. The advantage of this method is that it can significantly shorten the tuning time and reduce the demand for computing resources. One-Shot Tuner utilizes the shared information between models and can transfer knowledge between different models, thereby improving the tuning efficiency. However, the effect of One-Shot Tuner depends largely on the similarity between models. For models with large structural differences, the tuning effect may be affected.

[0012] In summary, although the current deep learning compilation and optimization technology has made some progress, it still has many shortcomings in terms of versatility, flexibility, and efficiency. More efficient and intelligent solutions are urgently needed to cope with the growing computing needs of deep learning. Summary of the invention

[0013] The present invention provides a method and system for generating a tensor program based on a large language model, which can realize efficient tensor program generation and optimization through the combination of automated tools and the large language model.

[0014] A method for generating a tensor program based on a large language model comprises the following steps: (1) Generate a random ONNX computation graph through the AI ​​compiler testing tool nnsmith; (2) Import the ONNX computation graph generated by nnsmith into the machine learning compilation framework TVM and convert it into the RelayIR intermediate representation; (3) Based on the model intermediate representation RelayIR inside TVM, the model is tuned through the automatic tuner Ansor in TVM to generate the optimal tensor program and output the JSON file measured during tuning; (4) Continuously perform model tuning and tensor program measurement, generate a large number of tuning measurement records, and form a tensor program dataset. The dataset contains the characteristics of the tensor program and the corresponding list of scheduling primitives required for optimization, namely the pass list. (5) Train a large language model based on the generated tensor program dataset to achieve the understanding and generation of tensor programs.

[0015] In step (1), after generating a random ONNX computation graph, set the input constraints of the ONNX computation graph. The input constraints include the dimension, data type, and range of the input tensor.

[0016] In step (2), the ONNX front-end interface of TVM is used to parse the ONNX computational graph. The steps of parsing the ONNX computational graph include computational graph rewriting, computational graph optimization, operator optimization, and code generation, and finally exporting it into a deployable module.

[0017] In step (2), during the conversion to the RelayIR intermediate representation, graph optimization operations are performed, including constant folding and operator merging, to improve the execution efficiency of the intermediate representation.

[0018] RelayIR is an intermediate representation in the TVM framework, which aims to provide a higher level of abstraction for the compilation and optimization of deep learning models, making it easier for the compiler to optimize. RelayIR represents the calculation as a directed acyclic graph, which facilitates a series of optimizations and transformations, and supports multiple backends: With RelayIR, users can compile models to a variety of different platforms and hardware, including CPU, GPU, and other custom hardware.

[0019] In step (3), the model is tuned using the automatic tuner Ansor in TVM, including setting the search space and target hardware for tuning and executing Ansor’s automatic tuning process. Ansor uses automatic search algorithms to optimize the execution path and implementation of tensor operations and generate the optimal tensor program.

[0020] The search space for optimization includes but is not limited to different tensor operation implementations, hardware execution paths, etc., to ensure that a variety of possible optimization options can be covered. The target hardware can include different types of computing platforms such as CPU and GPU, to ensure that the generated tensor program has efficient execution performance on different hardware platforms.

[0021] In step (3), the output JSON file records the performance data of each tensor operation implementation scheme during the tuning process, including key parameters such as execution time and scheduling primitive list, providing data support for subsequent analysis and optimization.

[0022] In step (5), the specific process of training the large language model is as follows: Extract samples from the generated tensor program dataset, including the characteristics of each program and the corresponding optimized passlist; perform data cleaning and formatting to ensure input and output alignment; Using data parallel or model parallel strategies, the data set is distributed to multiple computing devices for synchronous or asynchronous updates. The loss function is used to evaluate the degree of match between the pass list generated by the large language model and the real list, and parameters are optimized.

[0023] A tensor program generation system based on a large language model includes a memory and one or more processors, wherein the memory stores executable code, and when the one or more processors execute the executable code, they are used to implement the above-mentioned tensor program generation method.

[0024] Compared with the prior art, the present invention has the following beneficial effects: 1. The present invention uses a large language model to learn the optimization mode of existing tensor programs and directly generates high-performance tensor program code, avoiding the cumbersome code writing and debugging process of traditional methods and greatly improving the generation efficiency of tensor programs.

[0025] 2. By learning from a large number of tuning records and understanding the target hardware and application scenarios, the present invention can generate tensor programs with stronger targeting and better performance. Compared with codes written manually or generated by traditional compilers, the execution speed and efficiency of the program can be significantly improved.

[0026] 3. The present invention uses randomly generated ONNX computational graphs and rich tuning records for training, so that the large language model can learn optimization strategies for different types of tensor operations and different hardware platforms, has stronger generalization capabilities, and can be applied to a wider range of scenarios and tasks.

[0027] 4. The present invention automates the process of generating tensor programs. Users only need to provide simple tensor operation descriptions or model structure information to automatically generate high-performance program codes without in-depth understanding of the underlying hardware and compilation optimization technology. This greatly reduces the threshold for tensor program development and enables more developers to participate in the field of high-performance computing. BRIEF DESCRIPTION OF THE DRAWINGS

[0028] Figure 1 The present invention is a flowchart of a method for generating a tensor program based on a large language model according to an embodiment of the present invention.

[0029] Figure 2 The figure is a schematic diagram of the process of using nnsmith to generate a computational graph in an embodiment of the present invention.

[0030] Figure 3 Schematic diagram of the process of importing a computation graph into RelayIR using TVM in an embodiment of the present invention.

[0031] Figure 4 Schematic diagram of the process of automatic tuning using Ansor in an embodiment of the present invention. DETAILED DESCRIPTION

[0032] The present invention is further described in detail below in conjunction with the accompanying drawings and embodiments. It should be pointed out that the embodiments described below are intended to facilitate the understanding of the present invention and do not have any limiting effect on the present invention.

[0033] like Figure 1 As shown, a tensor program generation method based on a large language model includes the following steps: S1. Generate a random ONNX computation graph through the AI ​​compiler testing tool nnsmith and apply constraints to control the input.

[0034] S2. Import the ONNX computation graph generated by nnsmith into TVM and convert it to RelayIR intermediate representation.

[0035] S3. Based on the model intermediate representation RelayIR in TVM, the model is tuned through the TVM automatic tuner Ansor, and a JSON file of the measurements during tuning is output.

[0036] S4. Continue to perform model tuning and tensor program measurement, generate a large number of tuning measurement records, and form a tensor program dataset to assist large language models in generating high-performance tensor programs.

[0037] S5. Based on the tensor program dataset generated by S4, use methods such as large-scale parallel training to train a large language model to achieve understanding and generation of tensor programs.

[0038] like Figure 2 As shown, in step S1, nnsmith generates different ONNX computational graphs in three dimensions: computational graph validity, model diversity, and numerical validity. nnsmith can randomly generate structures such as multilayer perceptrons, convolutional neural networks, recursive neural networks, and the generated computational graphs can include different layer types, activation functions, and connection methods. This helps to test the generalization ability of the compiler. By comparing the generated model with the existing deep learning framework, nnsmith can verify whether the output of the compiler is consistent with the calculation results of the standard library. Set the input shape constraints of the generated ONNX computational graph to meet specific application requirements. The constraints may include, but are not limited to, the dimensions, data types, and ranges of the input tensors to ensure that the generated computational graphs are practical and effective in specific application scenarios. By setting the input shape constraints, a computational graph that better meets actual needs can be generated, thereby improving the applicability and performance of the generated tensor program.

[0039] In step S2, the import step includes using TVM's ONNX front-end interface to parse the ONNX computation graph and convert it into a RelayIR intermediate representation. Figure 3 As shown in the figure, the main steps of TVM parsing ONNX computational graphs are computational graph rewriting, computational graph optimization, operator optimization, and code generation, which can eventually be exported as deployable modules. Basic graph optimization operations are performed during the conversion process, including but not limited to constant folding, operator merging, etc., to improve the execution efficiency of the intermediate representation. Through this step, it is ensured that the model can be further optimized and tuned under the TVM framework, providing a foundation for subsequent steps. RelayIR is an intermediate representation in the TVM framework, which aims to provide a higher level of abstraction for the compilation and optimization of deep learning models, making it easier for the compiler to optimize. RelayIR represents the calculation as a directed acyclic graph, which facilitates a series of optimizations and conversions, and supports multiple backends: With RelayIR, users can compile models to a variety of different platforms and hardware, including CPU, GPU, and other custom hardware.

[0040] In step S3, the tuning step includes setting the tuning search space and target hardware, and executing Ansor's automatic tuning process. Ansor optimizes the execution path and implementation method of tensor operations through automatic search algorithms, and generates high-performance program code. Figure 4 As shown in the figure, the tuning process of Ansor is mainly implemented through task scheduler, program sampler, performance tuner and performance measurer. The search space of Ansor tuning includes but is not limited to different tensor operation implementation schemes, hardware execution paths, etc., to ensure that a variety of possible optimization options can be covered. The target hardware can include different types of computing platforms such as CPU and GPU to ensure that the generated tensor program has efficient execution performance on different hardware platforms. By executing Ansor's automatic tuning process, the optimal tensor operation implementation scheme can be automatically discovered and selected, thereby improving the execution efficiency and performance of the program. The output JSON file records the detailed performance data of each implementation scheme during the tuning process, including key parameters such as execution time and scheduling primitive list, providing data support for subsequent analysis and optimization.

[0041] In step S4, the data set contains a variety of different types of tensor operation instances and corresponding tuning records, providing rich training data. Through the accumulation and analysis of a large amount of data, the accuracy and efficiency of the large language model in generating high-performance tensor programs are improved. The step of generating a large number of measurement records when tuning further includes multiple tunings of different types of tensor operations and different input shapes, which can generate a wide range of tuning records to ensure the richness and diversity of the data set. After aggregating all the tuning records, a training data set containing rich tensor operation instances is formed, providing sufficient sample data for the training of the large language model, thereby improving the generalization ability of the model and the ability to generate high-performance tensor programs.

[0042] The generated data set is used to train the large language model so that the model can learn and master the characteristics of different types of tensor operations and their optimization solutions. Through the trained large language model, high-performance tensor operation program code can be automatically generated, significantly improving the execution efficiency and performance of the program. This step ensures that the generated tensor program not only has efficient execution performance, but also can adapt to different application scenarios and requirements, thereby improving the performance and applicability of the overall system.

[0043] In step S5, samples are extracted from the tensor program dataset generated in step S4, including the features of each program and the corresponding optimized pass list. Data cleaning and formatting are performed to ensure input and output alignment for model learning. Using the llama3 large language model, the input and output structure of the model are configured, with the input being the feature representation of the tensor program and the output being the optimized pass list. The training process is accelerated using large-scale parallel training techniques. Data parallelism or model parallelism strategies are adopted. The dataset is distributed to multiple computing devices for synchronous or asynchronous updates. The model is trained using tensor programs and their corresponding optimized pass lists. The loss function is used to evaluate the degree of match between the pass list generated by the model and the real list, and parameters are optimized. The model parameters are gradually adjusted so that it can accurately predict the appropriate pass list. Regularization techniques are applied to prevent overfitting and ensure that the model performs well on unseen data. The model is pre-trained on related tasks through transfer learning, and then fine-tuned to a specific pass list generation task. The pass list generated by the model is tested on the validation set to evaluate its effectiveness in optimizing the tensor program. The generated optimized pass list is compared with the manually adjusted list to verify its performance improvement effect. Deploy the trained model to generate optimization pass lists in real time and achieve fast tensor program tuning.

[0044] Based on the same inventive principle, an embodiment of the present invention also provides a tensor program generation system based on a large language model, including a memory and one or more processors, wherein executable code is stored in the memory, and when one or more processors execute the executable code, it is used to implement the tensor program generation method of the above embodiment.

[0045] The embodiments described above provide a detailed description of the technical solutions and beneficial effects of the present invention. It should be understood that the above are only specific embodiments of the present invention and are not intended to limit the present invention. Any modifications, supplements and equivalent substitutions made within the scope of the principles of the present invention should be included in the protection scope of the present invention.

Claims

1. A method for generating tensor programs based on a large language model, characterized in that: The following steps are involved: (1) Generate a random ONNX computation graph through the AI ​​compiler testing tool nnsmith; (2) Import the ONNX computation graph generated by nnsmith into the machine learning compilation framework TVM and convert it into the RelayIR intermediate representation; (3) Based on the model intermediate representation RelayIR inside TVM, the model is tuned through the automatic tuner Ansor in TVM to generate the optimal tensor program and output the JSON file measured during tuning; (4) Continuously perform model tuning and tensor program measurement, generate a large number of tuning measurement records, and form a tensor program dataset. The dataset contains the characteristics of the tensor program and the corresponding list of scheduling primitives required for optimization, namely the passlist; (5) Train a large language model based on the generated tensor program dataset to achieve the understanding and generation of tensor programs.

2. The method for generating a tensor program based on a large language model according to claim 1, characterized in that: In step (1), after generating a random ONNX computation graph, set the input constraints of the ONNX computation graph. The input constraints include the dimension, data type, and range of the input tensor.

3. The method for generating a tensor program based on a large language model according to claim 1, characterized in that: In step (2), the ONNX front-end interface of TVM is used to parse the ONNX computational graph. The steps of parsing the ONNX computational graph include computational graph rewriting, computational graph optimization, operator optimization, and code generation, and finally exporting it into a deployable module.

4. The method for generating a tensor program based on a large language model according to claim 1, characterized in that: In step (2), during the conversion to the RelayIR intermediate representation, graph optimization operations are performed, including constant folding and operator merging, to improve the execution efficiency of the intermediate representation.

5. The method for generating a tensor program based on a large language model according to claim 1, characterized in that: In step (3), the model is tuned using the automatic tuner Ansor in TVM, including setting the search space and target hardware for tuning and executing Ansor’s automatic tuning process. Ansor uses automatic search algorithms to optimize the execution path and implementation of tensor operations and generate the optimal tensor program.

6. The method for generating a tensor program based on a large language model according to claim 1, characterized in that: In step (3), the output JSON file records the performance data of each tensor operation implementation scheme during the tuning process.

7. The method for generating a tensor program based on a large language model according to claim 1, characterized in that: In step (5), the specific process of training the large language model is as follows: Extract samples from the generated tensor program dataset, including the characteristics of each program and the corresponding optimization pass list; perform data cleaning and formatting to ensure input and output alignment; Using data parallel or model parallel strategies, the data set is distributed to multiple computing devices for synchronous or asynchronous updates. The loss function is used to evaluate the degree of match between the pass list generated by the large language model and the real list, and parameters are optimized.

8. A tensor program generation system based on a large language model, characterized in that: It comprises a memory and one or more processors, wherein the memory stores executable code, and when the one or more processors execute the executable code, it is used to implement the tensor program generation method described in any one of claims 1 to 7.

Images