Memory management method, artificial intelligence processing system and computer program product

Abstract

The invention provides a memory management method, an artificial intelligence processing system and a computer program product, and the memory management method comprises the steps: obtaining a plurality of architecture parameters of an artificial intelligence model; obtaining a memory configuration size used in a training process of the artificial intelligence model according to the plurality of architecture parameters; according to the memory configuration size, a target memory area is configured in the random access memory, and multiple training processing stages in the training process of the artificial intelligence model commonly use the target memory area; and temporarily storing relay data corresponding to each training processing stage in the training process of the artificial intelligence model in the target memory area, wherein the relay data are transmitted among the acceleration processor module, the random access memory and the storage device. Therefore, according to the invention, the optimal memory configuration size can be automatically calculated, and the problem of memory fragmentation is effectively improved through the design of sharing the memory space.

Description

Technical Field

[0001] The present invention relates to a memory management method, an artificial intelligence processing system and a computer program product, and in particular to a method, an artificial intelligence processing system and a computer program product that can automatically configure and manage memory space during the training process of an artificial intelligence model. Background Art

[0002] With the rapid development of artificial intelligence technology, the demand for computing resources for training large artificial intelligence models is also increasing. In traditional artificial intelligence training systems, when the amount of training data exceeds the memory capacity of the acceleration processor module, the system must temporarily store part of the data in the storage device. In this process, the data transmission must be mediated by the random access memory.

[0003] However, the prior art has the following problems: First, the user needs to manually set the memory configuration parameters based on experience, which not only increases the difficulty of use, but also easily affects the system performance due to improper settings. Secondly, in traditional artificial intelligence training systems, different derivative products (also called relay data, such as weight parameters, gradient parameters, optimization parameters, etc.) generated during the training of artificial intelligence models will be configured in different memory spaces. This approach leads to inefficient memory usage, because the derivative products are actually used in different training processing stages of the training process of the artificial intelligence model. If independent memory space is configured, the memory space of the training processing stage that is not used will be idle.

[0004] In addition, due to the repeated configuration and release of the memory, it is easy to cause memory fragmentation. The memory fragmentation not only reduces the efficiency of memory use, but also may cause system performance degradation. Especially when conducting large-scale artificial intelligence model training, these problems will be more obvious and may even affect the stability of training. Summary of the invention

[0005] In view of the above problems, the present invention provides a memory management method, an artificial intelligence processing system and a computer program product. The present invention analyzes multiple architecture parameters of the artificial intelligence model, and automatically calculates the memory space requirements of different training processing stages before executing the training process, sets the memory configuration size accordingly, and implements a memory space sharing mechanism to improve the problems of low memory utilization efficiency and memory fragmentation in the prior art, thereby avoiding the delay problem caused by reallocating memory space for reading and writing.

[0006] An exemplary embodiment of the present invention provides a memory management method, which is applicable to an artificial intelligence processing system having a processor, an acceleration processor module, a random access memory, and a storage device. The memory management method includes: obtaining multiple architecture parameters of an artificial intelligence model; obtaining a memory configuration size used in a training process of the artificial intelligence model according to the multiple architecture parameters; configuring a target memory area in the random access memory according to the memory configuration size, wherein the target memory area is commonly used by multiple training processing stages in the training process of the artificial intelligence model; and temporarily storing relay data corresponding to each training processing stage in the training process of the artificial intelligence model in the target memory area, wherein the relay data is transmitted between the acceleration processor module, the random access memory, and the storage device.

[0007] In an exemplary embodiment of the present invention, the method also includes: obtaining data from the artificial intelligence model executed by the accelerated processor module, and writing the data to the target memory area as the relay data; and writing the relay data from the target memory area to the storage device.

[0008] In an exemplary embodiment of the present invention, the step of obtaining the memory configuration size used in the training process of the artificial intelligence model includes: obtaining a first memory usage in a back propagation phase of the training process of the artificial intelligence model; obtaining a second memory usage in a parameter update phase of the training process of the artificial intelligence model; and determining the memory configuration size based on the first memory usage and the second memory usage.

[0009] In an exemplary embodiment of the present invention, determining the memory allocation size according to the first memory usage and the second memory usage includes:

[0010] The maximum of the first memory usage and the second memory usage is selected as the memory configuration size for the training process of the artificial intelligence model.

[0011] In an exemplary embodiment of the present invention, the target memory area includes a first memory usage space configured for the back-propagation phase and a second memory usage space configured for the parameter update phase, wherein the first memory usage space and the second memory usage space at least partially overlap.

[0012] In an exemplary embodiment of the present invention, the multiple architecture parameters include: model architecture information, including the parameter quantity and input data dimension of each model layer of the M model layers of the artificial intelligence model; continuous operation segment width, used to set the total number of multiple model layers executed simultaneously; batch size, indicating the data processing volume of the training process; input data specifications, including image size or text sequence length; number of operation units, indicating the total number of multiple target acceleration processors in the acceleration processor module used to execute the artificial intelligence model; and parameter precision configuration, which defines the precision type and number of bytes of weight parameters, gradient parameters and optimization parameters.

[0013] In an exemplary embodiment of the present invention, the step of obtaining the first memory usage of the back propagation stage of the training process of the artificial intelligence model includes: calculating the first parameter memory requirement of each model layer according to the first parameter amount of each model layer of the M model layers in the back propagation stage, the first number of bytes defined by the parameter precision configuration and the number of operation units; calculating the activation value memory requirement of each model layer according to the batch size, the input data specification and the input data dimension and the number of activation value bytes defined by the parameter precision configuration; calculating the activation value memory requirement of each model layer according to the first parameter storage The memory requirement and the activation value memory requirement are used to calculate the first model layer data volume corresponding to each model layer; according to the continuous operation segment width and the first model layer data volume corresponding to each model layer, multiple first continuous operation segment data volumes are obtained, wherein the i-th first continuous operation segment data volume is the sum of the first model layer data volumes from the i-th model layer to the j-th model layer, wherein i is 1 to (M-(N-1)), N is the continuous operation segment width, and j is i+(N-1); and the maximum value of the multiple first continuous operation segment data volumes is used as the first memory usage.

[0014] In an exemplary embodiment of the present invention, the step of obtaining the second memory usage in the parameter update phase of the training process of the artificial intelligence model includes: calculating the second parameter memory requirement of each model layer according to the second parameter amount of each model layer of the M model layers in the parameter update phase, the second number of bytes defined by the parameter precision configuration and the number of operation units; obtaining the second model layer data amount corresponding to each model layer according to the second parameter memory requirement; obtaining multiple second continuous operation segment data amounts according to the continuous operation segment width and the second model layer data amount corresponding to each model layer, wherein the i-th second continuous operation segment data amount is the sum of the second model layer data amounts from the i-th model layer to the j-th model layer, where i is 1 to (M-(N-1)), N is the continuous operation segment width, and j is i+(N-1); and taking the maximum value of the multiple second continuous operation segment data amounts as the second memory usage.

[0015] In an exemplary embodiment of the present invention, the method also includes: determining whether the training process of the artificial intelligence uses a hybrid stage corresponding to the backpropagation stage and the parameter update stage; and in response to determining to use the hybrid stage: obtaining multiple hybrid continuous operation segment data amounts according to the continuous operation segment width, the first model layer data amount and the second model layer data amount corresponding to each model layer, wherein the i-th hybrid continuous operation segment data amount is the sum of the first model layer data amount and the second model layer data amount from the i-th model layer to the j-th model layer, where i is 1 to (M-(N-1)), N is the continuous operation segment width, and j is i+(N-1); and using the maximum value of the multiple hybrid continuous operation segment data amounts as the third memory usage of the hybrid stage, and using the third memory usage as the memory configuration size of the training process of the artificial intelligence model.

[0016] In an exemplary embodiment of the present invention, the method also includes: executing an intermediate program code, wherein the intermediate program code is used to: retrieve the multiple architectural parameters of the artificial intelligence model; perform operations on the memory configuration size; perform configuration of the target memory area; and manage the transmission of the relay data between the acceleration processor module, the random access memory, and the storage device.

[0017] In an exemplary embodiment of the present invention, the back propagation stage includes: based on a predefined data set, comparing the output result of the artificial intelligence model with the corresponding target result in the data set to calculate an error value; and according to the error value, performing gradient calculation on the parameters of each model layer in the artificial intelligence model using a back propagation algorithm to obtain the gradient values ​​corresponding to the multiple parameters.

[0018] In an exemplary embodiment of the present invention, the parameter updating stage includes: calculating the parameter updating direction of the artificial intelligence model according to a preset optimization algorithm and the gradient value calculated in the back propagation stage; and adjusting the parameter value in the artificial intelligence model according to the parameter updating direction.

[0019] An exemplary embodiment of the present invention provides an artificial intelligence processing system, comprising: a processor; an acceleration processor module for executing an artificial intelligence model; a random access memory; and a storage device. The processor is electrically connected to the acceleration processor module, the random access memory, and the storage device. The processor is configured by executing an intermediate program code to: obtain multiple architecture parameters of the artificial intelligence model; obtain a memory configuration size used in a training process of the artificial intelligence model according to the multiple architecture parameters; configure a target memory area in the random access memory according to the memory configuration size, wherein the target memory area is commonly used by multiple training processing stages in the training process of the artificial intelligence model; and temporarily store relay data corresponding to each training processing stage in the training process of the artificial intelligence model in the target memory area, wherein the relay data is transmitted between the acceleration processor module, the random access memory, and the storage device.

[0020] An exemplary embodiment of the present invention provides a computer program product, including an intermediary program code. The intermediary program code is executed by a processor of an artificial intelligence processing system to: obtain multiple architecture parameters of an artificial intelligence model, wherein the artificial intelligence model is executed by an acceleration processor module of the artificial intelligence processing system; obtain a memory configuration size used in the training process of the artificial intelligence model according to the multiple architecture parameters; configure a target memory area in a random access memory of the artificial intelligence processing system according to the memory configuration size, wherein multiple training processing stages in the training process of the artificial intelligence model use the target memory area together; and temporarily store relay data corresponding to each training processing stage in the training process of the artificial intelligence model in the target memory area, wherein the relay data is transmitted between the acceleration processor module, the random access memory and a storage device of the artificial intelligence processing system.

[0021] Based on the above, the memory management method, artificial intelligence processing system and non-temporary computer program product provided by the exemplary embodiments of the present invention can automatically calculate the optimal memory configuration size according to multiple architecture parameters of the artificial intelligence model, including information such as model architecture, continuous operation segment width, batch size, input data specifications, number of operation units and parameter precision configuration. By configuring the target memory area in the random access memory and allowing each training processing stage in the artificial intelligence model training process to use the target memory area together, the memory usage can be effectively reduced. In addition, by temporarily storing the relay data corresponding to each training processing stage in the training process of the artificial intelligence model in the target memory area, and transmitting data between the acceleration processor module, the random access memory and the storage device, the memory fragmentation problem traditionally caused by repeated configuration and release of memory space can be improved. BRIEF DESCRIPTION OF THE DRAWINGS

[0022] The accompanying drawings are included to provide a further understanding of the present invention and are incorporated in and constitute a part of this specification. The drawings illustrate embodiments of the present invention and together with the description serve to explain the principles of the present invention.

[0023] Figure 1 is a block diagram of an artificial intelligence processing system and a storage device according to an exemplary embodiment of the present invention;

[0024] Figure 2 is a schematic diagram of the operation of an artificial intelligence processing system and a storage device using a memory management method according to an exemplary embodiment of the present invention;

[0025] Figure 3A is a schematic diagram of continuous operation sections of a forward propagation stage according to an exemplary embodiment of the present invention;

[0026] Figure 3B is a schematic diagram of continuous operation sections of a back-propagation phase according to an exemplary embodiment of the present invention;

[0027] Figure 3C is a schematic diagram of continuous operation sections of a parameter updating phase according to an exemplary embodiment of the present invention;

[0028] Figure 3D is a schematic diagram of continuous operation sections of a hybrid stage according to an exemplary embodiment of the present invention;

[0029] Figure 4 is a flowchart of a memory management method according to an exemplary embodiment of the present invention.

[0030] Description of Figure Numbers

[0031] 1: Artificial Intelligence Computer

[0032] 10: Artificial Intelligence Processing System

[0033] 20: Storage device

[0034] 110: Processor

[0035] 120: Host memory

[0036] 130: Data transmission interface circuit

[0037] 140: Acceleration processor module

[0038] 210: Memory controller

[0039] 211: Processor

[0040] 212: Data management circuit

[0041] 213: memory interface control circuit

[0042] 214: Buffer memory

[0043] 220: Rewritable non-volatile memory module

[0044] 230: Connecting interface circuit

[0045] SD: Architecture parameters

[0046] TD: relay data

[0047] MA: Artificial Intelligence Model

[0048] A20: Arrow (indicates that the AI ​​model is executed on the acceleration processor module)

[0049] A21: Arrow (indicates obtaining multiple architecture parameters from the artificial intelligence model)

[0050] A22: Arrow (indicates calculation of memory configuration size and configuration of target memory area)

[0051] A23: Arrow (indicates writing relay data to the target memory area)

[0052] A24: Arrow (indicates writing relay data from the target memory area to the storage device)

[0053] A25: Arrow (indicates reading relay data from the storage device to the target memory area)

[0054] WT1~WTM: weight parameters

[0055] ACT1~ACTM: Activation value

[0056] GRD1~GRDM: Gradient parameters

[0057] OPT1~OPTM: Optimization parameters

[0058] LR1~LRM: Model layer

[0059] W1~WJ: continuous operation section

[0060] DS0~DS3: Continuous operation segment data volume

[0061] A30: Arrow (indicates the data volume of the 0th continuous operation segment corresponding to the continuous operation segment W1)

[0062] A31: Arrow (indicates the first continuous operation segment data volume corresponding to the continuous operation segment W1)

[0063] A32: Arrow (indicates the amount of data in the second continuous operation segment corresponding to the continuous operation segment W1)

[0064] A33: Arrow (indicates the amount of data in the third continuous operation segment corresponding to the continuous operation segment W1)

[0065] S410: Step (obtaining multiple architecture parameters of the artificial intelligence model)

[0066] S420: Step (obtaining a memory configuration size used in a training process of the artificial intelligence model according to the plurality of architecture parameters)

[0067] S430: Step (configuring a target memory area in a random access memory according to the memory configuration size, wherein the target memory area is commonly used by multiple training processing stages in the training process of the artificial intelligence model)

[0068] S440: Step (temporarily storing the relay data corresponding to each training processing stage of the training process of the artificial intelligence model in the target memory area, wherein the relay data is transmitted between the acceleration processor module, the random access memory and the storage device) DETAILED DESCRIPTION

[0069] Reference will now be made in detail to exemplary embodiments of the present invention, examples of which are illustrated in the accompanying drawings. Whenever possible, the same reference numerals are used in the drawings and the description to refer to the same or like parts.

[0070] Generally speaking, a storage device (also known as a memory storage system) includes a rewritable non-volatile memory module and a controller (also known as a memory controller). The storage device can be used together with an artificial intelligence processing system so that the artificial intelligence processing system can write data to the storage device or read data from the storage device.

[0071] Figure 1 FIG. 1 is a block diagram of an artificial intelligence processing system and a storage device according to an embodiment of the present invention. Figure 1 , the artificial intelligence computer (AIPC) 1 includes an artificial intelligence processing system 10 and a storage device 20. The artificial intelligence processing system 10 includes a processor (Processor) 110 (also called a first processor), a host memory (Host Memory) 120, a data transfer interface circuit (Data Transfer Interface Circuit) 130 and an acceleration processor module 140. In this embodiment, the processor 110 is coupled (also called electrically connected) to the host memory 120, the data transfer interface circuit 130 and the acceleration processor module. In another embodiment, the processor (Processor) 110, the host memory 120, the data transfer interface circuit 130 and the acceleration processor module 140 are electrically connected to each other using a system bus (System Bus). In this embodiment, the processor 110, the host memory 120 and the data transfer interface circuit 130 can be set on the motherboard of the artificial intelligence processing system 10.

[0072] In one embodiment, the artificial intelligence processing system 10 (also referred to as an artificial intelligence model execution device) includes: a processor 110, such as i5-14500 CPU; an accelerator module 140, such as an NVIDIA RTX 4060Ti graphics processor, having 16GB of video memory; a host memory 120, such as a DDR5 4800 memory with a capacity of 32GB x 2. A storage device 20 having 2TB of operating system disk space (which can be a HDD or SSD); and an intermediary program code (aiDAPTIVLink) for managing data transmission between the accelerator module 140, the host memory 120 and the storage device 20.

[0073] In one embodiment, the artificial intelligence processing system is, for example, aiDAPTIV+AI TPC launched by Phison. More specifically, the artificial intelligence processing system may include:

[0074] (1) Hardware layer: Cost-effective GPU; solid-state drive.

[0075] (2) Middleware layer (eg, middleware code): a middleware library that assists in managing data transmission between the unit 10 and the unit 20.

[0076] (3) Framework layer: artificial intelligence framework, such as PyTorch, TensorFlow. In one embodiment, the intermediary program code can provide target memory area information to the framework layer. This target memory area information includes the starting address, size and access permissions of the target memory area. More specifically, the processor 110 can obtain relevant information about the target memory area by executing the intermediary program code, including the starting memory address of the area in the host memory 120, the total configured capacity, and read or write access permissions. The processor 110 will provide this information to the framework layer, wherein the framework layer is an artificial intelligence computing development framework such as PyTorch. The processor 110 executes the framework layer to provide a program development interface to the artificial intelligence model.

[0077] (4) Application layer: artificial intelligence applications.

[0078] In one embodiment, the artificial intelligence processing system 10 is suitable for executing multiple artificial intelligence models, including: Large Language Model (LLM), such as llama-3.1-8B or Llama-2-13B; Image Generation Model (Image Generator), such as Stablediffusion SDXL; Speech Recognition Model, such as Whisper.

[0079] In one embodiment, as the top layer of the artificial intelligence training system 10, the application layer can execute a variety of artificial intelligence applications. These applications are mainly used to develop and train artificial intelligence models in specific fields. For example: (1) In the field of natural language processing, developers can use the interface provided by the application layer to train large language models such as Llama-2-13B, and adjust the model through domain-specific data sets so that it can better understand the professional terms and context of a specific industry. (2) In the field of computer vision, developers can use image generation models such as Stable Diffusion SDXL as a basis to train models based on the company's internal product gallery to generate product diagrams that match the company's image.

[0080] In one embodiment, the computing tasks supported by the artificial intelligence processing system 10 include: inference computing: executing a trained model to perform prediction or generation tasks; and domain training (also known as model training process): model training or fine-tuning for a specific application domain.

[0081] Among them, the model training process includes multiple training processing stages: back propagation stage: calculating the gradient of the loss function to the parameters; parameter update stage: updating the model parameters; hybrid stage: combining the operation mode of back propagation and parameter update. During this training process, the artificial intelligence processing system 10 uses the memory management method of the present invention to allow different training processing stages to share memory space, thereby effectively reducing memory usage. Even when using a smaller capacity host memory 120, the artificial intelligence processing system 10 can still support larger-scale artificial intelligence model training and can be applied to artificial intelligence teaching, research and other fields.

[0082] The storage device 20 includes a storage controller 210, a rewritable non-volatile memory module 220 and a connection interface circuit 230. The storage controller 210 includes a processor 211, a data management circuit 212 and a memory interface control circuit 213.

[0083] In this embodiment, the artificial intelligence processing system 10 is electrically connected to the storage device 20 through the data transmission interface circuit 130 and the connection interface circuit 230 of the storage device 20 to perform data access operations. For example, the artificial intelligence processing system 10 can store data to the storage device 20 or read data from the storage device 20 via the data transmission interface circuit 130.

[0084] In this embodiment, the number of the data transmission interface circuit 130 can be one or more. Through the data transmission interface circuit 130, the motherboard can be electrically connected to the storage device 20 via a wired or wireless manner. The storage device 20 can be, for example, a USB storage device, a memory, a solid state drive (SSD) or a wireless memory storage device. The wireless memory storage device can be, for example, a near field communication (NFC) memory storage device, a wireless fax (WiFi) memory storage device, a Bluetooth memory storage device or a low power consumption Bluetooth memory storage device (e.g., iBeacon) and other memory storage devices based on various wireless communication technologies. In addition, the motherboard can also be electrically connected to various I / O devices such as a global positioning system (GPS) module, a network interface card, a wireless transmission device, a keyboard, a screen, a speaker, etc. through a system bus.

[0085] In this embodiment, the data transmission interface circuit 130 and the connection interface circuit 230 are interface circuits compatible with the high-speed peripheral component interconnect interface (Peripheral Component Interconnect Express, PCI Express) standard or other types of standard interface circuits. In addition, the data transmission interface circuit 130 and the connection interface circuit 230 use the fast non-volatile memory interface standard (Non-Volatile Memory express, NVMe) communication protocol or other types of standard communication protocols to transmit data.

[0086] Furthermore, in another embodiment, the connection interface circuit 230 and the memory controller 210 may be packaged in one chip, or the connection interface circuit 230 may be disposed outside a chip including the memory controller 210 .

[0087] In one embodiment, the acceleration processor module 140 may include multiple graphics processing units (GPUs), tensor processing units (TPUs), neural network processors (NPUs), artificial intelligence accelerators (AI Accelerators) or other processors suitable for performing artificial intelligence operations. The acceleration processor module 140 is used to execute artificial intelligence models. These acceleration processor modules 140 have parallel computing capabilities and are particularly suitable for processing large-scale parallel computing tasks such as matrix operations and convolution operations in artificial intelligence models.

[0088] In one embodiment, the artificial intelligence model executed by the acceleration processor module 140 includes multiple model layers, each of which has one or more nodes (the nodes of two model layers can be connected to each other), and each model layer includes specific parameter quantities and input data dimensions corresponding to the one or more nodes. The training process of the artificial intelligence model can be mainly divided into multiple training processing stages that run independently, including: a forward operation (Forward) stage, a back propagation (Backward) stage, and a parameter update (Update) stage.

[0089] In one embodiment, the training process stage includes a forward operation stage, in which the artificial intelligence model performs inference operations based on the input data to generate model prediction results. Specifically, in the forward operation stage, the artificial intelligence model will process from the first model layer to the last model layer in sequence, and each model layer will perform operations on the input data according to the weight parameters of the model layer, and use the operation results as input data for the next layer, and finally generate the prediction results of the model.

[0090] In one embodiment, the training process stage includes a back propagation stage. In the back propagation stage, the value of the loss function (Loss Function) (also known as the error value) is first calculated based on the difference between the target output result (Target) defined in advance in the data set and the prediction result (Prediction) of the artificial intelligence model. Then, based on the value of the loss function, the system uses a back propagation algorithm to calculate the gradient (Gradient) corresponding to the parameters (Weight / Parameter) of each model layer, starting from the last model layer of the artificial intelligence model. The gradient represents the direction in which the model parameters need to be adjusted to reduce the value of the loss function.

[0091] In one embodiment, the training process stage includes a parameter update stage. In the parameter update stage, the system uses the gradient calculated in the back propagation stage and calculates the correction direction (Update direction) (also called parameter update direction) of the model parameters according to a pre-set optimization algorithm (e.g., Stochastic Gradient Descent (SGD) or AdamW algorithm). Then, the system updates the parameters of the model according to the correction direction, so that the artificial intelligence model can produce a prediction result that is closer to the target result in the next training iteration.

[0092] In some embodiments, the back propagation phase and the parameter update phase can adopt a hybrid operation mode (also called a hybrid phase), that is, a specific model layer can simultaneously perform the back propagation phase and the parameter update phase. Each training processing phase in the training process will perform corresponding data processing operations for each model layer in the artificial intelligence model in sequence according to its corresponding operation target.

[0093] In the forward operation phase, the AI ​​model will perform operations from the first layer to the last layer in sequence to generate prediction results. In the back propagation phase, the system will calculate the gradient of the loss function to the parameters of each layer. In the parameter update phase, the system will update the model parameters according to the optimization algorithm.

[0094] In one embodiment, each training processing stage generates relay data that needs to be temporarily stored. The relay data may include one or more of the following:

[0095] (1) Model parameters: including weight parameters stored in mixed precision (such as BF16) and weight parameters stored in high precision (such as FP32). In convolutional neural networks, these weight parameters may be the filter weights of the convolution layer, for example, a 3×3 convolution kernel may contain 9 weight values; in the Transformer model, these weight parameters may be the weights of the query, key, and value matrices in the self-attention mechanism.

[0096] (2) Gradient parameters: record the gradient value of the loss function for each layer parameter. For example, in an image classification task, when there is an error between the model's predicted output and the actual label, the system calculates the partial differential of the loss function for each weight parameter. If cross entropy is used as the loss function, the gradient parameter means adjusting the direction and magnitude of each weight parameter to reduce the prediction error.

[0097] (3) Optimization parameters: including momentum parameters and variance parameters. Taking the AdamW optimization algorithm as an example, the momentum parameter records the moving average of the weight parameter updates in the past, which is used to maintain the inertia of parameter updates; the variance parameter records the moving average of the square of the gradient, which is used to adaptively adjust the learning rate of each parameter. These parameters help the model converge faster and avoid falling into local minima.

[0098] (4) Activation value: records the output results of each layer. In the convolutional neural network for image processing, the activation value may be the feature map processed by the ReLU (Rectified Linear Unit) function. For example, from the original 224×224 pixel input image, 64 112×112 feature maps may be obtained after the first layer of convolution. In the language model, the activation value may be the word vector sequence output by each layer of the Transformer encoder, and its dimension may be 512 or 1024, etc.

[0099] These relay data can be used interchangeably in different training processing stages of the training process. For example, weight parameters are used in the forward operation and back-propagation stages, gradient parameters are not used in the forward propagation stage, and optimization parameters are used in the parameter update stage, etc. By analyzing these usage patterns, different stages can share the same memory space, thereby improving memory usage efficiency. For example, when processing an image classification model with 1,000 categories, the last fully connected layer may have more than 1 million weight parameters. These parameters can be stored in BF16 format to save memory space during forward operation, generate corresponding gradient parameters during back-propagation, and finally participate in the operation together with momentum parameters and variance parameters during parameter update to update the model.

[0100] In one embodiment, taking a large language model with 180B (billion) parameters as an example, if the AdamW optimization algorithm is used for training, in the mixed precision training mode:

[0101] (1) The weight parameters are stored in BF16 format, and each parameter requires 2 bytes.

[0102] (2) Gradient parameters are stored in FP32 format, and each parameter requires 4 bytes.

[0103] (3) Momentum and variance parameters are stored in FP32 format, requiring 4 bytes per parameter.

[0104] In this case:

[0105] During the back-propagation phase, the system needs to access weight parameters in BF16 format and gradient parameters in FP32 format and generate activation values.

[0106] During the parameter update phase, the system needs to access weight parameters, gradient parameters, momentum parameters, and variance parameters in FP32 format.

[0107] If a hybrid operation mode is adopted, the system will immediately perform parameter updates of each model layer after completing the back-propagation calculation of the model layer.

[0108] Through the memory management method of the present invention, the amount of DRAM required for training large language models can be greatly reduced. For example, for a model with 180B parameters, the traditional method may require 1TB to 2TB of DRAM, but after adopting the method of the present invention, only 64GB or 128GB of DRAM is needed to complete the training. This reduction in memory usage mainly comes from the design of sharing memory space for relay data in the back propagation stage and the parameter update stage.

[0109] In the present embodiment, the host memory 120 is used to temporarily store instructions or data executed by the processor 110. For example, in the present embodiment, the host memory 120 may be a dynamic random access memory (DRAM), a static random access memory (SRAM), etc. However, it must be understood that the present invention is not limited thereto, and the host memory 120 may also be other suitable memories. The host memory 120 will be configured with a target memory area to temporarily store the relay data TD during the training process of the artificial intelligence model.

[0110] The memory controller 210 is used to execute a plurality of logic gates or control instructions implemented in hardware or firmware and to perform operations such as writing, reading and erasing data in the rewritable non-volatile memory module 220 according to the instructions of the artificial intelligence processing system 10. In more detail, the processor 211 in the memory controller 210 is hardware with computing capabilities, which is used to control the overall operation of the memory controller 210.

[0111] On the other hand, the processor 110 in the artificial intelligence processing system 10 is also hardware with computing capabilities, which is used to control the overall operation of the artificial intelligence processing system 10. Specifically, the processor 110 is programmed by a plurality of control instructions / program codes, and when the storage device 10 operates, these control instructions / program codes are executed to implement the memory management method provided by the present invention. In other embodiments, the control instructions / program codes corresponding to the memory management method can be further implemented as a circuit unit in the form of hardware to implement the memory management method provided by the present invention.

[0112] It is worth mentioning that in the present embodiment, the processor 110 and the processor 211 are, for example, a central processing unit (CPU), a microprocessor, or other programmable processing units (Microprocessor), a digital signal processor (DSP), a programmable controller, an application specific integrated circuit (ASIC), a programmable logic device (PLD) or other similar circuit elements, but the present invention is not limited thereto.

[0113] In this embodiment, the memory controller 210 further includes a data management circuit 212 and a memory interface control circuit 213 .

[0114] The data management circuit 212 is electrically connected to the processor 211, the memory interface control circuit 213 and the connection interface circuit 230. The data management circuit 212 is used to receive instructions from the processor 211 to transmit data. For example, data is read from the artificial intelligence processing system 10 (e.g., the host memory 120) via the connection interface circuit 230, and the read data is written to the rewritable non-volatile memory module 220 via the memory interface control circuit 213 (e.g., the write operation is performed according to the write instruction from the artificial intelligence processing system 10). For another example, data is read from one or more physical units of the rewritable non-volatile memory module 220 via the memory interface control circuit 213 (the data can be read from one or more storage units in the one or more physical units), and the read data is written to the artificial intelligence processing system 10 (e.g., the host memory 120) via the connection interface circuit 230 (e.g., the read operation is performed according to the read instruction from the artificial intelligence processing system 10). In another embodiment, the data management circuit 212 can also be integrated into the processor 211.

[0115] The memory interface control circuit 213 is used to receive instructions from the processor 211 and cooperate with the data management circuit 212 to perform a write (also called programming) operation, a read operation or an erase operation on the rewritable non-volatile memory module 220 .

[0116] In addition, the data to be written to the rewritable non-volatile memory module 220 will be converted into a format acceptable to the rewritable non-volatile memory module 220 via the memory interface control circuit 213. Specifically, if the processor 211 wants to access the rewritable non-volatile memory module 220, the processor 211 will transmit a corresponding instruction sequence to the memory interface control circuit 213 to instruct the memory interface control circuit 213 to perform a corresponding operation. For example, these instruction sequences may include a write instruction sequence indicating writing data (such as relay data), a read instruction sequence indicating reading data (such as relay data), an erase instruction sequence indicating erasing data, and corresponding instruction sequences for indicating various memory operations. These instruction sequences may include one or more signals, or data on the bus. These signals or data may include instruction codes or program codes. For example, in the read instruction sequence, information such as the read identification code, memory address, and physical address may be included.

[0117] In one embodiment, the memory controller 210 further includes a buffer memory 214. The buffer memory is electrically connected to the processor 211 and is used to temporarily store data and instructions from the artificial intelligence processing system 10, data from the rewritable non-volatile memory module 220, or other system data for managing the storage device 20, so that the processor 211 can quickly access the data, instructions or system data from the buffer memory 214.

[0118] The rewritable non-volatile memory module 220 is electrically connected to the memory controller 210 (memory interface control circuit 213) and is used to store data written by the artificial intelligence processing system 10. The data is, for example, the architecture data of the artificial intelligence model or the neural network model itself and / or the relay data TD related to the artificial intelligence model or the neural network model executed by the processor 110.

[0119] In this embodiment, the rewritable non-volatile memory module 220 includes a NAND flash memory, but the present invention is not limited thereto. For example, in other embodiments, the rewritable non-volatile memory module 220 may be a NOR flash memory, a phase change memory (PCM), a resistive random-access memory (ReRAM), or other types of non-volatile memory.

[0120] Figure 2 2 is a schematic diagram of the operation of an artificial intelligence processing system and a storage device using a memory management method according to an exemplary embodiment of the present invention.

[0121] Please refer to Figure 2 , Figure 2The diagram shows the data flow and interaction between the processor 110 , the host memory 120 , the accelerated processor module 140 and the storage device 20 in the artificial intelligence processing system 10 .

[0122] In one embodiment, as indicated by arrow A20 , the artificial intelligence model MA is executed on the acceleration processor module 140 .

[0123] As shown by arrow A21, the processor 110 obtains multiple architecture parameters SD of the artificial intelligence model MA from the acceleration processor module 140. These architecture parameters SD include: model architecture information, continuous operation segment width, batch size, input data specification, number of operation units and parameter precision configuration. In other embodiments, the architecture parameters SD can also be obtained from a database corresponding to the artificial intelligence model MA (for example, the database can be stored in the storage device 20).

[0124] In one embodiment, the multiple architecture parameters obtained by the processor 110 from the artificial intelligence model MA include, for example:

[0125] (1) Model architecture information: the number of parameters and input data dimensions of each of the M model layers of the artificial intelligence model. For example, in a language model with a Transformer architecture, each Transformer layer may include a self-attention mechanism and a feedforward neural network, where the number of parameters of the self-attention mechanism depends on the hidden layer dimension (e.g., 1024), while the input data dimension may be the dimension of the word embedding vector (e.g., 768).

[0126] (2) Continuous operation segment width: The continuous operation segment width is used to set the total number of multiple model layers to be executed simultaneously. For example, in a large language model with 48 layers, the continuous operation segment width may be set to 4, which means that the system will load and execute 4 consecutive transformer layers at the same time, and after completing the operation of these 4 layers, it will load the next 4 layers for operation.

[0127] (3) Batch size: The batch size indicates the amount of data processed during the training process. For example, in an image classification task, the batch size may be set to 32, which means that the system will process 32 images at a time; in language model training, it may be set to 16, which means that 16 text sequences will be processed at a time.

[0128] (4) Input data specifications: The input data specifications include image size or text sequence length. For example, in image processing tasks, the input specification may be an RGB image of 224×224 pixels; in language processing tasks, it may be a text sequence of length 512, each sequence containing 512 tokens.

[0129] (5) Number of computing units: The number of computing units represents the total number of multiple target accelerators in the accelerator module 140 used to execute the artificial intelligence model. For example, the system may use 8 image processors for distributed training, each image processor is responsible for processing 1 / 8 of the overall model parameters, or use 4 image processors for data parallel processing, each GPU processes different data batches.

[0130] (6) Parameter precision configuration: The parameter precision configuration defines the precision type and number of bytes of weight parameters, gradient parameters, and optimization parameters. For example, in mixed precision training, the weight parameters of forward operation and back propagation use the BF16 format (2 bytes) to save memory space, while gradient accumulation and parameter update use the FP32 format (4 bytes) to ensure numerical stability. Optimization parameters such as momentum and variance all use the FP32 format to provide sufficient numerical precision.

[0131] (7) Update strategy: SGD model or AdamW model. For example, in the traditional training process, the system will first complete the backpropagation operation of all layers, and then perform the parameter update operation; in the hybrid operation mode, the system will immediately perform the parameter update of each layer after completing the backpropagation operation of each layer. Taking the mixed precision AdamW optimization algorithm as an example, when performing parameter update, each model layer will require five types of data: (1) weight parameters in BF16 format (2 bytes), (2) weight parameters in FP32 format (4 bytes), (3) gradient parameters in FP32 format (4 bytes), (4) momentum parameters in FP32 format (4 bytes), and (5) variance parameters in FP32 format (4 bytes). If calculated based on the number of parameters per layer, these parameters require a total of 18 bytes (2+4+4+4+4=18) of storage space. Through the hybrid operation mode, the system can release part of the memory space of the layer immediately after the backpropagation is completed, further improving the memory usage efficiency.

[0132] As shown by arrow A22, the processor 110 calculates the memory configuration size and configures the target memory area 121 in the host memory 120 based on the multiple architecture parameters. Specifically, the processor 110 first obtains the first memory usage of the back propagation stage of the training process of the artificial intelligence model, and the second memory usage of the parameter update stage. Then, the processor 110 selects the largest one from the first memory usage and the second memory usage as the memory configuration size of the training process of the artificial intelligence model. In the target memory area 121, a first memory usage space configured as the back propagation stage and a second memory usage space configured as the parameter update stage are included, wherein the first memory usage space and the second memory usage space at least partially overlap.

[0133] For example, in one embodiment, the overlapping design of the first memory usage space and the second memory usage space is based on the usage sequence of different types of parameters in different stages of the artificial intelligence model training process. In the back propagation stage, the first memory usage space is mainly used to store: (1) weight parameters stored in BF16 format, and (2) gradient parameters stored in FP32 format. In the parameter update stage, the second memory usage space is mainly used to store: (1) weight parameters stored in FP32 format, (2) gradient parameters stored in FP32 format, (3) momentum parameters stored in FP32 format, and (4) variance parameters stored in FP32 format.

[0134] Taking the AdamW optimization algorithm of mixed precision training as an example, assuming that a certain model layer has 1 million parameters, in the back propagation stage: the weight parameters in BF16 format require 2MB of space (1 million × 2 bytes); the gradient parameters in FP32 format require 4MB of space (1 million × 4 bytes), a total of 6MB of first memory usage space is required.

[0135] In the parameter update stage: the weight parameters in FP32 format require 4MB of space (1 million × 4 bytes); the gradient parameters in FP32 format require 4MB of space (1 million × 4 bytes); the momentum parameters in FP32 format require 4MB of space (1 million × 4 bytes); the variance parameters in FP32 format require 4MB of space (1 million × 4 bytes), a total of 16MB of second memory space is required.

[0136] Since the back propagation phase and the parameter update phase are not executed at the same time, the two phases can share some memory space. For example, the space of the gradient parameter space used in the back propagation phase (4MB) can be reused for the gradient parameters in the parameter update phase (also requiring 4MB); the BF16 weight parameter space (2MB) in the back propagation phase can be used as part of the weight parameter space (4MB) in the parameter update phase.

[0137] With this overlapping design, for this model layer with 1 million parameters, only about 16MB of memory space is actually allocated, instead of 22MB (6MB + 16MB) which is the sum of the memory requirements of the two stages. This design can significantly reduce memory usage.

[0138] As shown by arrow A23, the data generated by the acceleration processor module 140 when executing the artificial intelligence model will be written into the target memory area 121 as relay data TD. For example, in the back propagation stage, when the acceleration processor module 140 performs gradient calculation, it will access the stored weight parameters (e.g., in BF16 format), and write the calculated gradient parameters (e.g., in FP32 format) into the first memory usage space; in the parameter update stage, when the acceleration processor module 140 performs parameter update operation, it will access the stored gradient parameters, and write the updated weight parameters, momentum parameters and variance parameters into the second memory usage space. Since the first memory usage space and the second memory usage space at least partially overlap, and these parameters are used alternately in different stages, the memory space can be effectively utilized to avoid wasting memory resources.

[0139] As shown by arrow A24, when the target memory area 121 needs to store the relay data TD in the storage device 20, the processor 110 controls the relay data to be written from the target memory area 121 to the storage device 20. Conversely, as shown by arrow A25, when the acceleration processor module 140 needs to use the relay data TD in the storage device 20, the processor 110 controls the relay data TD to be read back from the storage device 20 to the target memory area 121 so that the acceleration processor module 140 can use the relay data TD. This data transmission mechanism allows the acceleration processor module 140 to perform training of large artificial intelligence models even when the memory capacity is limited.

[0140] Through the above architecture, this embodiment provides an effective memory management method. First, by analyzing the multiple architecture parameters, the memory configuration size is automatically calculated, avoiding the problem of the user having to manually set the memory parameters. Secondly, by allowing the first memory usage space in the back propagation phase and the second memory usage space in the parameter update phase to share the target memory area 121, the memory usage efficiency is improved. Finally, by establishing an effective data transmission mechanism between the acceleration processor module 140, the target memory area 121 and the storage device 20, the system can support the training needs of large artificial intelligence models.

[0141] Figure 3A FIG. 4 is a schematic diagram of continuous operation sections of a forward propagation stage according to an exemplary embodiment of the present invention.

[0142] Please refer to Figure 3A , assuming that the artificial intelligence model MA has M model layers (LR1~LRM), in which in the forward propagation stage which is one of multiple training processing stages, each model layer contains two types of relay data: weight parameters (WT1~WTM) and activation values ​​(ACT1~ACTM).

[0143] In one embodiment, the processor 110 calculates the zeroth memory usage (also called forward memory usage) of the forward propagation stage. In one embodiment, the step of obtaining the zeroth memory usage of the forward propagation stage of the training process of the artificial intelligence model includes: calculating the zeroth parameter memory requirement (also called forward parameter memory requirement) of each model layer according to the zeroth parameter amount (also called forward parameter amount) of each model layer of the M model layers in the forward propagation stage and the zeroth byte number (also called forward byte number) defined by the parameter precision configuration and the number of operation units; calculating the activation value memory requirement of each model layer according to the batch size, the input data specification and the input data dimension and the number of activation value bytes defined by the parameter precision configuration; calculating the activation value memory requirement of each model layer according to the zeroth parameter memory; calculating the activation value memory requirement of each model layer according to the batch size, the input data specification and the input data dimension and the number of activation value bytes defined by the parameter precision configuration; calculating the activation value memory requirement of each model layer according to the zeroth parameter memory The memory requirement and the activation value memory requirement are used to calculate the zeroth model layer data volume (also called the forward model layer data volume) corresponding to each model layer; according to the continuous operation segment width and the zeroth model layer data volume corresponding to each model layer, multiple zeroth continuous operation segment data volumes (also called the forward continuous operation segment data volumes) are obtained, wherein the i-th zeroth continuous operation segment data volume is the sum of the zeroth model layer data volumes from the i-th model layer to the j-th model layer, wherein i is 1 to (M-(N-1)), N is the continuous operation segment width, and j is i+(N-1); and the maximum value of the multiple zeroth continuous operation segment data volumes is used as the zeroth memory usage.

[0144] For example, first, the processor 110 calculates the memory requirement of the zeroth parameter of each model layer. Specifically, for each model layer (e.g., model layer LR1), the memory requirement of its weight parameter (WT1) is calculated by multiplying the parameter amount of the model layer by the number of bytes defined by the parameter precision configuration (e.g., the weight parameter in BF16 format is 2 bytes), and then dividing it by the number of computing units (because the parameters can be distributed on multiple computing units in the acceleration processor module).

[0145] Next, the processor 110 calculates the memory requirement of the activation value of each model layer. Taking the activation value (ACT1) of the model layer LR1 as an example, its memory requirement is calculated by multiplying the batch size (indicating the amount of data processed simultaneously) by the input data specification (such as image size or text sequence length) and the input data dimension, and then multiplying it by the number of activation value bytes defined by the parameter precision configuration (for example, 4 bytes for FP32 format).

[0146] Afterwards, the processor 110 adds the zeroth parameter memory requirement and the activation value memory requirement of each model layer to obtain the zeroth model layer data volume corresponding to the model layer. Figure 3A As shown, the continuous operation segment width (N) is 3, which means that the system will load and operate 3 consecutive model layers at the same time. The processor 110 will calculate multiple zeroth continuous operation segment data volumes (W1-WJ) according to this continuous operation segment width. As shown by arrow A30, in this example, DS0 represents the zeroth continuous operation segment data volume corresponding to the first continuous operation segment W1.

[0147] For example, the first continuous operation segment W1 (the continuous operation segment width is 3) includes the sum of the first model layer data volumes of model layers LR1, LR2 and LR3; the second continuous operation segment W2 includes the sum of the first model layer data volumes of model layers LR2, LR3 and LR4, and so on. The starting model layer number i of each continuous operation segment ranges from 1 to (M-(N-1)), where M is the total number of model layers and N is the continuous operation segment width (3 in this example, but the present invention is not limited to this). Each continuous operation segment includes model layers starting from the i-th layer to the j-th layer (j=i+(N-1)).

[0148] Finally, the processor 110 selects the maximum value from the data amounts of these continuous operation segments as the zeroth memory usage in the forward propagation stage.

[0149] For another example, please refer to Figure 3A , when calculating the zeroth memory usage in the forward propagation phase, the processor 110 performs the following steps in sequence:

[0150] (1) Calculate the parameter memory requirement for each model layer (corresponding to Figure 3A Weight parameters WT1~WTM in: parameter memory requirement = [number of model layer parameters] × [number of parameter precision bytes] ÷ [number of computing units]. For example, assuming that a model layer has 1 million parameters, the weight parameters use the BF16 format (2 bytes), and 4 accelerated processors are used for distributed computing, then the parameter memory requirement of the layer is: (1,000,000) × 2 ÷ 4 = 500,000 bytes.

[0151] (2) Calculate the input data memory requirement for each model layer: Input data memory requirement = [batch size] × [input data specification] × [number of bytes of precision]. For example, if the batch size is 32, the input is a text sequence of length 512 (input data specification), and the FP32 format (4 bytes) is used, then the input data memory requirement is: 32 × 512 × 4 = 65,536 bytes.

[0152] (3) Calculate the activation value memory requirement for each model layer (corresponding to Figure 3A ACT1~ACTN in ): Activation value memory requirement = [batch size] × [output data dimension] × [number of precision bytes]. For example, if the batch size is 32, the output dimension is 1024, and the FP32 format (4 bytes) is used, the activation value memory requirement is: 32×1024×4=131,072 bytes.

[0153] (4) Calculate the zeroth model layer data volume of each model layer: zeroth model layer data volume = parameter memory requirement + input data memory requirement + activation value memory requirement. Continuing with the above example, the zeroth model layer data volume of the layer is: 500,000 + 65,536 + 131,072 = 696,608 bytes.

[0154] (5) Calculate the data volume of the continuous operation segment: Figure 3A As shown, the continuous operation segment width N is 3, and the system will calculate the data volume of multiple continuous operation segments (W1~WJ). The calculation method of the zeroth continuous operation segment data volume DS0 of each continuous operation segment Wi is: DS0 = the zeroth model layer data volume of the i-th layer + the zeroth model layer data volume of the (i+1)th layer + the zeroth model layer data volume of the (i+2)th layer.

[0155] As shown by arrow A30, DS0 represents the data amount of the zeroth continuous operation segment corresponding to the continuous operation segment W1. For example, assuming that each layer has a different data amount: LR1 (self-attention layer): 696,608 bytes; LR2 (feedforward neural network layer): 835,930 bytes; LR3 (self-attention layer): 696,608 bytes; LR4 (feedforward neural network layer): 835,930 bytes; LR5 (self-attention layer): 696,608 bytes.

[0156] The data volume of each continuous operation segment is calculated as follows (taking three segments as an example): W1 data volume = LR1 + LR2 + LR3 = 696,608 + 835,930 + 696,608 = 2,229,146 bytes; W2 data volume = LR2 + LR3 + LR4 = 835,930 + 696,608 + 835,930 = 2,368,468 bytes; W3 data volume = LR3 + LR4 + LR5 = 696,608 + 835,930 + 696,608 = 2,229,146 bytes.

[0157] (6) Obtain the zeroth memory usage: zeroth memory usage = MAX(W1 data size, W2 data size, W3 data size) = MAX(2,229,146,2,368,468,2,229,146) = 2,368,468 bytes. In this example, since W2 contains two feedforward neural network layers (LR2 and LR4), its data size is the largest. This value (zeroth memory usage, 2,368,468 bytes) will be used to determine the amount of memory space allocated to the forward propagation stage in the target memory area.

[0158] It is worth noting that compared with the back propagation stage ( Figure 3B ), the forward propagation stage does not need to store gradient parameters, so its memory requirement is usually small. Therefore, in one embodiment, the processor 110 does not need to calculate the zeroth memory usage in the forward propagation stage (because it can be replaced by the first memory usage in the backward propagation stage).

[0159] Figure 3B FIG. 4 is a schematic diagram of consecutive operation sections of a back-propagation stage according to an exemplary embodiment of the present invention.

[0160] Please refer to Figure 3B , assuming that the artificial intelligence model MA has M model layers (LR1~LRM), each of which contains three types of relay data in the back-propagation stage: weight parameters (WT1~WTM), activation values ​​(ACT1~ACTM) and gradient parameters (GRD1~GRDM).

[0161] In one embodiment, the processor 110 calculates a first memory usage of the back-propagation phase. In one embodiment, the step of obtaining the first memory usage of the back propagation stage of the training process of the artificial intelligence model includes: calculating the first parameter memory requirement of each model layer according to the first parameter amount of each model layer of the M model layers in the back propagation stage, the first number of bytes defined by the parameter precision configuration and the number of operation units; calculating the activation value memory requirement of each model layer according to the batch size, the input data specification and the input data dimension and the number of activation value bytes defined by the parameter precision configuration; calculating the first model layer data amount corresponding to each model layer according to the first parameter memory requirement and the activation value memory requirement; obtaining multiple first continuous operation segment data amounts according to the continuous operation segment width and the first model layer data amount corresponding to each model layer, wherein the i-th first continuous operation segment data amount is the sum of the first model layer data amounts from the i-th model layer to the j-th model layer, wherein i is 1 to (M-(N-1)), N is the continuous operation segment width, and j is i+(N-1); and taking the maximum value of the multiple first continuous operation segment data amounts as the first memory usage.

[0162] For example, first, the processor 110 calculates the first parameter memory requirement of each model layer. Specifically, for each model layer (e.g., model layer LR1), the memory requirement of its weight parameter (WT1) and gradient parameter (GRD1) is calculated by multiplying the parameter amount of the model layer by the number of bytes defined by the parameter precision configuration (e.g., the weight parameter in BF16 format is 2 bytes, and the gradient parameter in FP32 format is 4 bytes), and then dividing by the number of computing units (because the parameters can be distributed on multiple computing units of the acceleration processor module).

[0163] Next, the processor 110 calculates the memory requirement of the activation value of each model layer. Taking the activation value (ACT1) of the model layer LR1 as an example, its memory requirement is calculated by multiplying the batch size (indicating the amount of data processed simultaneously) by the input data specification (such as image size or text sequence length) and the input data dimension, and then multiplying it by the number of activation value bytes defined by the parameter precision configuration (for example, 4 bytes for FP32 format).

[0164] Afterwards, the processor 110 adds the first parameter memory requirement and the activation value memory requirement of each model layer to obtain the first model layer data volume corresponding to the model layer. Figure 3BAs shown, the continuous operation segment width (N) is 3, which means that the system will load and operate 3 consecutive model layers at the same time. The processor 110 will calculate multiple first continuous operation segment data volumes (corresponding to W1-WJ) according to the continuous operation segment width. As shown by arrow A31, in this example, DS1 represents the first continuous operation segment data volume corresponding to the continuous operation segment W1.

[0165] For example, the first continuous operation segment W1 (the continuous operation segment width is 3) includes the sum of the first model layer data volumes of model layers LR1, LR2 and LR3; the second continuous operation segment W2 includes the sum of the first model layer data volumes of model layers LR2, LR3 and LR4, and so on. The starting model layer number i of each continuous operation segment ranges from 1 to (M-(N-1)), where M is the total number of model layers and N is the continuous operation segment width (3 in this example, but the present invention is not limited to this). Each continuous operation segment includes model layers starting from the i-th layer to the j-th layer (j=i+(N-1)).

[0166] Finally, the processor 110 selects the maximum value from the data amounts of these continuous operation segments as the first memory usage of the back propagation stage. This first memory usage will be used to determine the size of the memory space allocated to the back propagation stage in the target memory area. For example, the size of the memory space allocated to the back propagation stage can be the first memory usage.

[0167] For another example, please refer to Figure 3B , when calculating the first memory usage in the back-propagation phase, the processor 110 performs the following steps in sequence:

[0168] (1) Calculate the first parameter memory requirement of each model layer (corresponding to Figure 3B Weight parameters WT1~WTM and gradient parameters GRD1~GRDM in the model layer): the first parameter memory requirement = 2×[number of model layer parameters]×[number of bytes of parameter precision]÷[number of computing units]. For example, assuming that a model layer has 1 million parameters, the weight parameters use the BF16 format (2 bytes), the gradient parameters use the FP32 format (4 bytes), and 4 acceleration processors are used for distributed computing, then the first parameter memory requirement of the layer is: 2×(1,000,000)×(2+4)÷4=3,000,000 bytes.

[0169] The reason for multiplying this step by 2 is that when training a neural network, in addition to storing the weights of the model (that is, the parameters learned by the model), it also needs to store the gradients used for optimization. The so-called parameter quantity refers to the total number of all adjustable parameters in this neural network layer. The storage space size of these parameters will be affected by the numerical precision we choose. For example, when using the BF16 format (a more memory-saving floating-point format), each parameter only takes up 2 bytes. The reason for dividing by the number of GPUs at the end is that when we use multiple GPUs to train the model, these parameters will be distributed to different GPUs to speed up the calculation. Therefore, each GPU only needs to store a part of the parameters.

[0170] (2) Calculate the activation value memory requirement for each model layer (corresponding to Figure 3B ACT1~ACTM in ): Activation value memory requirement = [batch size] × [input data size] × [input data dimension] × [activation value precision byte number]. For example, if the batch size is 32, the input is a text sequence of length 512, the embedding dimension of each token is 1024, and the activation value is stored in FP32 format (4 bytes), then the memory requirement for the activation value of the layer is: 32×512×1024×4=67,108,864 bytes.

[0171] (3) Calculate the first model layer data size of each model layer: first model layer data size = first parameter memory requirement + (activation value memory requirement × [offload flag]). The offload flag is a Boolean value (0 or 1) indicating whether the activation value needs to be offloaded to external storage. If the offload flag is 1, the memory space requirement for the activation value needs to be considered more; otherwise (0), only the memory space requirement for the weight and gradient is considered.

[0172] Continuing with the above example, if the activation value needs to be stored ([Uninstall Flag] = 1), the first model layer data size of the layer is: 3,000,000 + (67,108,864 × 1) = 70,108,864 bytes

[0173] (4) Calculate the data volume of the continuous operation segment: Figure 3BAs shown, the continuous operation segment width N is 3, and the system will calculate the data volume of each of the multiple continuous operation segments (corresponding to W1~WJ). The calculation method of the continuous operation segment data volume DS1 of each continuous operation segment Wi is: DS1 = the first model layer data volume of the i-th layer + the first model layer data volume of the (i+1)-th layer + the first model layer data volume of the (i+(N-1))-th layer. Among them, i ranges from 1 to (M-2), and M is the total number of model layers. As shown by arrow A31, DS1 represents the first continuous operation segment data volume corresponding to the continuous operation segment W1.

[0174] For example, suppose that each model layer in the transformer model has a different amount of first block data due to different structures:

[0175] LR1 (self-attention layer): the first model layer has 70,108,864 bytes of data; LR2 (feedforward neural network layer): the first model layer has 85,324,800 bytes of data; LR3 (self-attention layer): the first model layer has 70,108,864 bytes of data; LR4 (feedforward neural network layer): the first model layer has 85,324,800 bytes of data; LR5 (self-attention layer): the first model layer has 70,108,864 bytes of data.

[0176] The data volume of each continuous operation segment is: W1=70,108,864+85,324,800+70,108,864=225,542,528 bytes (including LR1, LR2, LR3); W2=85,324,800+70,108,864+85,324,800=240,758,464 bytes (including LR2, LR3, LR4); W3=70,108,864+85,324,800+70,108,864=225,542,528 bytes (including LR3, LR4, LR5).

[0177] (5) Obtaining the first memory usage: first memory usage = MAX (W1 data volume, W2 data volume, ..., WJ data volume). In this example, the data volume corresponding to W2 is the largest (reflecting the peak memory requirement of the entire back-propagation process). Therefore, the processor 110 selects the W2 data volume, 240,758,464 bytes, as the first memory usage of the back-propagation stage.

[0178] Through the above calculation process, the system can accurately evaluate the maximum memory usage required to load N consecutive model layers at any time during the back-propagation phase, thereby optimizing the configuration of memory space. This calculation method takes into account the parameter memory requirements, activation value storage strategy, and the characteristics of distributed computing to effectively configure memory space.

[0179] Figure 3C FIG. 4 is a schematic diagram of continuous operation sections of a parameter updating phase according to an exemplary embodiment of the present invention.

[0180] Please refer to Figure 3C , assuming that the artificial intelligence model MA has M model layers (LR1~LRN), each of which contains three types of relay data in the parameter update phase: weight parameters (WT1~WTM), gradient parameters (GRD1~GRDM) and optimization parameters (OPT1~OPTM). It should be noted that each layer may contain multiple optimization parameters. For example, when using the AdamW optimization algorithm, the optimization parameters of each layer will include momentum parameters (such as OPT1 corresponding to model layer 1LR1) and variance parameters (such as OPT1.1 corresponding to model layer 1LR1 (not shown)).

[0181] In one embodiment, the processor 110 calculates the second memory usage in the parameter updating phase.

[0182] In one embodiment, the step of obtaining the second memory usage in the parameter update phase of the training process of the artificial intelligence model includes:

[0183] First, the processor 110 calculates the second parameter memory requirement of each model layer. Specifically, for each model layer (e.g., model layer LR1), the second parameter quantity of the model layer (including the parameter quantities of the FP32 format weight parameter WT1, gradient parameter GRD1, momentum parameter, and variance parameter of the model layer) is multiplied by the second byte number defined by the parameter precision configuration (e.g., 4 bytes in FP32 format), and then divided by the number of computing units (because the parameters can be distributed on multiple computing units of the acceleration processor module).

[0184] Next, the processor 110 sets the second parameter memory requirement of each model layer to the second model layer data volume corresponding to the model layer. Since the activation value does not need to be stored during the parameter update phase, the second model layer data volume only includes the second parameter memory requirement. Figure 3C As shown, the continuous operation segment width (N) is 3, which means that the system will load and operate 3 consecutive model layers at the same time. The processor 110 will calculate multiple second continuous operation segment data volumes (W1-WJ) according to the continuous operation segment width. As shown by arrow A32, in this example, DS2 represents the second continuous operation segment data volume corresponding to the continuous operation segment W1.

[0185] For example, the first continuous operation segment W1 includes the sum of the second model layer data of model layers LR1, LR2 and LR3; the second continuous operation segment W2 includes the sum of the second model layer data of model layers LR2, LR3 and LR4, and so on. The starting model layer number i of each continuous operation segment ranges from 1 to (M-(N-1)), where M is the total number of model layers and N is the width of the continuous operation segment (3 in this example). Each continuous operation segment includes model layers starting from the i-th layer to the j-th layer (j=i+(N-1)).

[0186] Finally, the processor 110 selects the maximum value from the data amounts of the second continuous operation segments as the second memory usage of the parameter update phase. The second memory usage will be used to determine the size of the memory space allocated to the parameter update phase in the target memory area. For example, the size of the memory space allocated to the parameter update phase can be the second memory usage.

[0187] Please refer to Figure 3C When calculating the second memory usage in the parameter update phase, the processor 110 sequentially performs the following steps:

[0188] (1) Calculate the amount of data required for each update block. Taking the mixed-precision AdamW optimization algorithm as an example, each update block requires five types of data: weight parameters in BF16 format (2 bytes); weight parameters in FP32 format (4 bytes); gradient parameters in FP32 format (4 bytes); momentum parameters in FP32 format (4 bytes); and variance parameters in FP32 format (4 bytes).

[0189] The second parameter memory requirement of each model layer is calculated as follows: second parameter memory requirement = [number of model layer parameters] × [total number of required precision bytes] ÷ [number of computing units]. Wherein, total number of required precision bytes = 2 + 4 + 4 + 4 + 4 = 18. For example, assuming that a model layer has 1 million parameters and 4 accelerated processors are used for distributed computing, the second parameter memory requirement of the layer is: (1,000,000) × 18 ÷ 4 = 4,500,000 bytes.

[0190] (2) Calculate the continuous operation segment data volume of each continuous operation segment. Figure 3C As shown, when the continuous operation segment width (N) is 3, starting from the i-th model layer of the model, the memory requirements of 3 continuous model layers are accumulated. The second continuous operation segment data volume DS2 of each continuous operation segment Wi is calculated as follows: the second continuous operation segment data volume = the second parameter memory requirement of the i-th layer + the second parameter memory requirement of the (i+1)-th layer + the second parameter memory requirement of the (i+(N-1))-th layer.

[0191] For example, suppose in the Transformer model: LR1 (self-attention layer): 4,500,000 bytes; LR2 (feed-forward neural network layer): 5,400,000 bytes; LR3 (self-attention layer): 4,500,000 bytes.

[0192] As indicated by arrow A32, DS2 represents the data volume of the second continuous operation segment corresponding to the continuous operation segment W1: DS2 = 4,500,000 + 5,400,000 + 4,500,000 = 14,400,000 bytes.

[0193] This calculation process is repeated until all layers of the entire artificial intelligence model MA are covered to find the maximum continuous operation segment data volume. Each time a layer is moved, the continuous operation segment data volume in the new continuous operation segment is calculated to ensure that all possible memory usage scenarios that may occur during the update process are taken into account.

[0194] Through the above method, we can obtain the data volume of multiple second continuous operation segments (taking M=5 as an example, assuming a total of 5 model layers). For example: the data volume of W1 = 14,400,000 bytes (including LR1, LR2, LR3); the data volume of W2 = 15,300,000 bytes (including LR2, LR3, LR4); the data volume of W3 = 14,400,000 bytes (including LR3, LR4, LR5).

[0195] As shown by arrow A32, DS2 represents the amount of data of the second continuous operation segment corresponding to the first continuous operation segment W1, that is, 14,400,000 bytes. In the parameter update phase, the artificial intelligence model MA needs to process five different parameter data at the same time: weight parameters in BF16 format (2 bytes), weight parameters in FP32 format (4 bytes), gradient parameters in FP32 format (4 bytes), momentum parameters in FP32 format (4 bytes) and variance parameters in FP32 format (4 bytes). Therefore, for each parameter in the model, 18 bytes (2+4+4+4+4=18) of memory space are actually required to store all the data required in the parameter update phase.

[0196] It is worth noting that feedforward neural network layers (e.g., LR2, LR4) usually have larger parameters, so their memory requirements (5,400,000 bytes) are higher than the self-attention layer (4,500,000 bytes). This difference is reflected in the data size of the continuous operation segment, for example, the data size of W2 (15,300,000 bytes) is larger than the data size of W1 and W3 (14,400,000 bytes), because W2 contains two feedforward neural network layers (LR2 and LR4).

[0197] Through the above calculation process, the system can accurately evaluate the memory usage required to load N consecutive model layers at any time during the parameter update phase. This calculation method takes into account the parameter requirements of different precision formats and the characteristics of distributed computing, thereby effectively configuring the memory space. It is worth noting that in the parameter update phase, due to the need to simultaneously maintain the weight parameters in BF16 format (2 bytes), weight parameters in FP32 format (4 bytes), gradient parameters (4 bytes), momentum parameters (4 bytes) and variance parameters (4 bytes), each parameter position requires 18 bytes of memory space. In contrast, in the forward propagation phase, each parameter only requires the weight parameters in BF16 format (2 bytes), and the memory requirement is smaller. In addition, although the forward propagation phase also needs to store activation values ​​(4 bytes), the calculation of activation values ​​is not related to the number of parameters, but depends on the batch size, input data specifications and input data dimensions.

[0198] Figure 3D FIG. 4 is a schematic diagram of consecutive operation sections of a hybrid stage according to an exemplary embodiment of the present invention.

[0199] Please refer to Figure 3D , assuming that the artificial intelligence model MA has M model layers (LR1~LRM), each of which contains four types of relay data in the hybrid stage: weight parameters (WT1~WTM), activation values ​​(ACT1~ACTM), gradient parameters (GRD1~GRDM) and optimization parameters (OPT1~OPTM). These relay data reflect that in the hybrid stage, the system will perform backpropagation operations and parameter update operations at the same time.

[0200] In one embodiment, the processor 110 first determines whether the training process of the artificial intelligence model uses a hybrid stage. If it is determined to use a hybrid stage, the processor 110 calculates the third memory usage of the hybrid stage. In one embodiment, the step of obtaining the third memory usage of the hybrid stage of the training process of the artificial intelligence model includes: adding the first model layer data volume of each model layer in the back propagation stage and the second model layer data volume in the parameter update stage to obtain the model layer data volume of the model layer in the hybrid stage; calculating multiple hybrid continuous operation segment data volumes based on the continuous operation segment width and the hybrid stage model layer data volume of each model layer.

[0201] When calculating the third memory usage of the hybrid stage, the processor 110 executes the following steps in sequence:

[0202] (1) Calculate the first model layer data size of each model layer in the back propagation phase: first model layer data size = (2×[model layer parameter size]×[parameter precision byte number]÷[number of computing units])+([batch size]×[input data specification]×[input data dimension]×[activation value precision byte number]). In one embodiment, this step takes into account, for example: weight parameter memory requirements in BF16 format; gradient parameter memory requirements in FP32 format; activation value memory requirements.

[0203] (2) Calculate the second model layer data volume of each model layer in the parameter update phase: second model layer data volume = [model layer parameter volume] × [total precision required bytes (e.g., 18 bytes)] ÷ [number of computing units]. In one embodiment, this step may take into account, for example: weight parameters in BF16 format (e.g., 2 bytes); weight parameters in FP32 format (e.g., 4 bytes); gradient parameters in FP32 format (e.g., 4 bytes); momentum parameters in FP32 format (e.g., 4 bytes); variance parameters in FP32 format (e.g., 4 bytes).

[0204] (3) Calculate the model layer data volume of each model layer in the hybrid stage: hybrid model layer data volume = first model layer data volume + second model layer data volume.

[0205] For example, suppose in the transformer model, for the self-attention layer: the number of model layer parameters is 1 million, the batch size is 32, the input sequence length is 512, the input dimension is 1024, and 4 accelerated processors are used for distributed computing.

[0206] The calculation of the data volume of the first model layer is: (a) BF16 weights and FP32 gradients: 2×1,000,000×(2+4)÷4=3,000,000 bytes; (b) activation values: 32×512×1024×4=67,108,864 bytes. The total data volume of the first model layer is: 70,108,864 bytes.

[0207] The data volume of the second model layer is calculated as: 1,000,000×18÷4=4,500,000 bytes.

[0208] Therefore, the amount of model layer data in the hybrid stage of the self-attention layer is: 70,108,864+4,500,000=74,608,864 bytes

[0209] Similarly, for the feedforward neural network layer (assuming the number of parameters is 1.2 million): the first model layer data volume is: (a) BF16 weights and FP32 gradients: 2×1,200,000×(2+4)÷4=3,600,000 bytes; (b) activation values: 32×512×1024×4=67,108,864 bytes, totaling: 70,708,864 bytes. In addition, the second model layer data volume is: 1,200,000×18÷4=5,400,000 bytes.

[0210] Therefore, the model layer data volume of the feedforward neural network layer in the hybrid stage is: 70,708,864+5,400,000=76,108,864 bytes.

[0211] (4) Calculate the data volume of the continuous operation segment: Figure 3D As shown, the continuous operation segment width N is 3, and the system will calculate the data volume of multiple continuous operation segments (W1~WJ). The calculation method of the hybrid continuous operation segment data volume DS3 of each continuous operation segment Wi is: DS3 = the hybrid model layer data volume of the i-th layer + the hybrid model layer data volume of the (i+1)th layer + the hybrid model layer data volume of the (i+2)th layer.

[0212] For example, assuming that each layer has a different amount of hybrid model layer data due to different structures (taking M=5 as an example, assuming a total of 5 model layers), LR1 (self-attention layer): 74,608,864 bytes; LR2 (feedforward neural network layer): 76,108,864 bytes; LR3 (self-attention layer): 74,608,864 bytes; LR4 (feedforward neural network layer): 76,108,864 bytes; LR5 (self-attention layer): 74,608,864 bytes.

[0213] The data volume of the mixed continuous operation segments of each continuous operation segment is: W1 data volume = LR1 + LR2 + LR3 = 74,608,864 + 76,108,864 + 74,608,864 = 225,326,592 bytes; W2 data volume = LR2 + LR3 + LR4 = 76,108,864 + 74,608,864 + 76,108,864 = 226,826,592 bytes; W3 data volume = LR3 + LR4 + LR5 = 74,608,864 + 76,108,864 + 74,608,864 = 225,326,592 bytes.

[0214] (5) Obtain the third memory usage (taking M=5 as an example): the third memory usage=MAX(W1 data volume, W2 data volume, W3 data volume)=MAX(225,326,592,226,826,592,225,326,592)=226,826,592 bytes. This value (the third memory usage, 226,826,592 bytes) will be used to determine the size of the memory space allocated to the hybrid stage in the target memory area.

[0215] It is worth noting that the memory requirement of the hybrid stage is significantly higher than that of the pure back-propagation stage or parameter update stage, because the system needs to maintain all the relay data required for the back-propagation operation and the parameter update operation at the same time. Although this design increases the real-time memory requirement, it can reduce the consumption of data transmission resources between different training processing stages and improve training efficiency.

[0216] Figure 4 is a flowchart of a memory management method according to an exemplary embodiment of the present invention.

[0217] Please refer to Figure 4 In step S410, the processor 110 obtains multiple architecture parameters of an artificial intelligence model.

[0218] Then, in step S420, the processor 110 obtains the memory configuration size used in the training process of the artificial intelligence model based on the multiple architecture parameters. Specifically, the processor calculates the first memory usage in the back propagation phase, the second memory usage in the parameter update phase, or the third memory usage in the hybrid phase according to the training strategy of the artificial intelligence model, and uses it as the memory configuration size. For example, for a model with 180 billion parameters, traditional training methods may require 1TB to 2TB of DRAM resources, but after adopting the memory management method of the present invention, only 64GB or 128GB of DRAM resources are required to complete the training.

[0219] Next, in step S430, the processor 110 configures a target memory area in the machine access memory according to the memory configuration size, wherein the target memory area is commonly used by multiple training processing stages in the training process of the artificial intelligence model. By allowing different training processing stages to share the same memory space, memory usage can be effectively reduced and memory fragmentation can be avoided. This design is suitable for training environments with limited memory resources, such as AI computers using a single accelerated processor (such as NVIDIA RTX 4070Ti SUPER).

[0220] Finally, in step S440, the processor 110 temporarily stores the relay data corresponding to each training processing stage in the training process of the artificial intelligence model in the target memory area, wherein the relay data is transmitted between the acceleration processor module, the random access memory and the storage device. For example, when the acceleration processor module needs to write relay data to the storage device, the relay data will first be written to the target memory area, and then written to the storage device from the target memory area; when the acceleration processor module needs to read the relay data in the storage device, the relay data will first be read from the storage device to the target memory area, and then provided to the acceleration processor module from the target memory area. Through this design, even on a system using a smaller capacity random access memory, it can effectively support the training of large AI models, such as the training requirements of models such as Llama-3.1-8B or Llama-2-13B.

[0221] However, Figure 4 The steps in the above are described in detail, so I will not repeat them here. It is worth noting that Figure 4 Each step in the above process can be implemented as multiple program codes or circuits, and the present invention is not limited thereto. Figure 4 The method can be used in conjunction with the above exemplary embodiments or can be used alone, and the present invention is not limited thereto.

[0222] In one embodiment, the processor of the artificial intelligence processing system 10 manages the memory by executing an intermediary program code (middleware). The intermediary program code is a software layer that connects the storage device hardware and the artificial intelligence processing system in series, and its main functions include: capturing and analyzing the architecture information of the artificial intelligence model, including the number of model layers, parameter quantity, computing requirements, etc.; automatically calculating the optimized memory configuration scheme based on the analysis results; coordinating the data transmission between the acceleration processor, the random access memory and the storage device; providing a standardized program interface for use by upper-level applications; dynamically managing the configuration and release of system resources to avoid resource waste and memory fragmentation.

[0223] Specifically, the intermediary program code first extracts the multiple architecture parameters from the artificial intelligence model, including: model architecture information (such as the number of parameters and input data dimensions of each model layer), continuous operation segment width, batch size, input data specifications, number of operation units, and parameter precision configuration, etc. For example, when the user chooses to train the Llama-2-13B model, the intermediary program code will automatically analyze the structure of the model and obtain the required architecture parameters.

[0224] Next, the intermediate program code performs a memory allocation size calculation based on the multiple architecture parameters. For example, if the user uses the mixed precision AdamW optimization algorithm for training, the intermediate program code calculates the memory space required to process the weight parameters, gradient parameters, momentum parameters, and variance parameters simultaneously in the hybrid stage.

[0225] Then, the intermediary program code will execute the configuration of the target memory area in the random access memory. For example, when the system is equipped with 64GB of DDR5 memory, the intermediary program code will configure a target memory area of ​​appropriate size according to the calculation results, so that the area can be shared by different training processing stages. It is worth mentioning that in another implementation example, the target memory area can also be configured in a high-bandwidth memory. Specifically, in one embodiment, the target memory area can be configured in a high-bandwidth memory (HBM) of the acceleration processor module 140. High-bandwidth memory has a higher data access bandwidth. In this case, the target memory area can be configured in a specific segment in this high-bandwidth memory, so that the artificial intelligence model MA executed in the acceleration processor module 140 can access the relay data TD more quickly during the training process.

[0226] In one embodiment, in order to improve data access efficiency, the target memory area can be configured in both the host memory 120 and the high-bandwidth memory of the acceleration processor module 140. This is because during the training process of the artificial intelligence model, data needs to be frequently transmitted between the acceleration processor module 140 and the host memory 120. By configuring the target memory area of ​​the same size in the high-bandwidth memory of the acceleration processor module 140, a corresponding memory space can be established for temporarily storing the relay data TD during the training process. This design can reduce the delay of data transmission and improve training efficiency.

[0227] Finally, the intermediary program code is responsible for managing the transmission of relay data between the acceleration processor module, the random access memory and the storage device. When the artificial intelligence model needs to access relay data during training, the intermediary program code coordinates the read and write operations of the data. For example, when the gradient parameters stored in the storage device need to be accessed during the parameter update phase, the intermediary program code manages the data transmission path to ensure that the data can flow effectively between different devices.

[0228] Through these functions of the intermediate program code, the artificial intelligence processing system 10 can automate the memory management process, reduce the user's configuration burden, and improve the memory usage efficiency of the artificial intelligence processing system 10.

[0229] In one embodiment, the computer program product of the present invention includes an intermediary program code to perform memory management through the processor 110 of the artificial intelligence processing system 10. When the user installs and executes the computer program product on the artificial intelligence processing system 10, the intermediary program code will be loaded and executed by the processor of the artificial intelligence processing system 10. The intermediary program code can be stored in a storage medium or the Internet, and can be read and executed to implement the memory management method of the present invention. Through this design, the computer program product can reduce the temporary memory resources required to implement artificial intelligence operations, so that the training needs of large AI models can be effectively supported in a resource-constrained training environment.

[0230] In summary, the memory management method, artificial intelligence processing system and computer program product provided by the exemplary embodiments of the present invention can automatically calculate the optimal memory configuration size based on multiple architecture parameters of the artificial intelligence model, including information such as model architecture, continuous operation segment width, batch size, input data specifications, number of operation units and parameter precision configuration. By configuring the target memory area in the random access memory and allowing each stage in the artificial intelligence model training process to use the target memory area together, the memory usage can be effectively reduced. In addition, by temporarily storing the relay data corresponding to each stage of the training process of the artificial intelligence model in the target memory area and transmitting data between the acceleration processor module, the random access memory and the storage device, the memory fragmentation problem traditionally caused by repeated configuration and release of memory space can be improved.

[0231] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, rather than to limit it. Although the present invention has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the aforementioned embodiments, or replace some or all of the technical features therein with equivalents. However, these modifications or replacements do not cause the essence of the corresponding technical solutions to deviate from the scope of the technical solutions of the embodiments of the present invention.

Claims

1. A memory management method, applicable to an artificial intelligence processing system having a processor, an accelerated processor module, a random access memory and a storage device, characterized in that: The memory management method comprises: Get multiple architectural parameters of the AI ​​model; According to the plurality of architecture parameters, obtaining a memory configuration size used in a training process of the artificial intelligence model; According to the memory configuration size, a target memory area is configured in the random access memory, wherein the target memory area is commonly used by multiple training processing stages in the training process of the artificial intelligence model, wherein in each training processing stage, a corresponding data processing operation is performed on each model layer in the artificial intelligence model according to a processing order of the training processing stage; and The relay data corresponding to each training processing stage in the training process of the artificial intelligence model is temporarily stored in the target memory area, wherein the relay data is transmitted between the acceleration processor module, the random access memory and the storage device.

2. The memory management method according to claim 1, characterized in that: The method further comprises: Acquire data from the artificial intelligence model executed by the accelerated processor module, and write the data to the target memory area as the relay data; and The relay data is written from the target memory area to the storage device.

3. The memory management method according to claim 1, characterized in that: The multiple training processing stages include a back propagation stage and a parameter update stage, and the step of obtaining the memory configuration size used in the training process of the artificial intelligence model includes: Obtaining a first memory usage of the back-propagation phase of the training process of the artificial intelligence model; Obtaining a second memory usage in the parameter update phase of the training process of the artificial intelligence model; and The memory allocation size is determined according to the first memory usage and the second memory usage.

4. The memory management method according to claim 3, characterized in that: Determining the memory configuration size according to the first memory usage and the second memory usage includes: The maximum of the first memory usage and the second memory usage is selected as the memory configuration size for the training process of the artificial intelligence model.

5. The memory management method according to claim 3, characterized in that: The target memory area includes a first memory usage space configured for the back-propagation phase and a second memory usage space configured for the parameter update phase, wherein the first memory usage space and the second memory usage space at least partially overlap.

6. The memory management method according to claim 3, characterized in that: The plurality of architectural parameters include: Model architecture information, including the parameter quantity and input data dimension of each model layer of the M model layers of the artificial intelligence model; The continuous operation section width is used to set the total number of multiple model layers to be executed simultaneously; Batch size, which represents the amount of data that needs to be processed during the training process; Input data specifications, including image size or text sequence length; The number of computing units represents the total number of multiple target acceleration processors in the acceleration processor module for executing the artificial intelligence model; and Parameter precision configuration, which defines the precision type and number of bytes of weight parameters, gradient parameters, and optimization parameters.

7. The memory management method according to claim 6, characterized in that: The step of obtaining the first memory usage in the back propagation stage of the training process of the artificial intelligence model includes: Calculate the first parameter memory requirement of each model layer according to the first parameter amount of each model layer of the M model layers in the back propagation stage, the first number of bytes defined by the parameter precision configuration and the number of operation units; Calculate the activation value memory requirement of each model layer according to the batch size, the input data specification, the input data dimension, and the number of activation value bytes defined by the parameter precision configuration; Calculating the first model layer data volume corresponding to each model layer according to the first parameter memory requirement and the activation value memory requirement; According to the continuous operation segment width and the first model layer data volume corresponding to each model layer, a plurality of first continuous operation segment data volumes are obtained, wherein the i-th first continuous operation segment data volume is the sum of the first model layer data volumes from the i-th model layer to the j-th model layer, wherein i is 1 to (M-(N-1)), N is the continuous operation segment width, and j is i+(N-1); and The maximum value among the data amounts of the plurality of first continuous operation segments is used as the first memory usage.

8. The memory management method according to claim 7, characterized in that: The step of obtaining the second memory usage in the parameter updating stage of the training process of the artificial intelligence model comprises: Calculate the second parameter memory requirement of each model layer according to the second parameter amount of each model layer of the M model layers in the parameter updating phase, the second number of bytes defined by the parameter precision configuration and the number of operation units; According to the second parameter memory requirement, obtaining the second model layer data volume corresponding to each model layer; According to the continuous operation segment width and the second model layer data volume corresponding to each model layer, a plurality of second continuous operation segment data volumes are obtained, wherein the i-th second continuous operation segment data volume is the sum of the second model layer data volumes from the i-th model layer to the j-th model layer, wherein i is 1 to (M-(N-1)), N is the continuous operation segment width, and j is i+(N-1); and The maximum value among the data amounts of the plurality of second continuous operation segments is used as the second memory usage.

9. The memory management method according to claim 8, characterized in that: The method further comprises: Determining whether the training process of the artificial intelligence uses a hybrid phase corresponding to the back-propagation phase and the parameter update phase; and React to the decision using the hybrid stage: According to the continuous operation segment width, the first model layer data volume and the second model layer data volume corresponding to each model layer, a plurality of hybrid continuous operation segment data volumes are obtained, wherein the i-th hybrid continuous operation segment data volume is the sum of the first model layer data volume and the second model layer data volume from the i-th model layer to the j-th model layer, wherein i is 1 to (M-(N-1)), N is the continuous operation segment width, and j is i+(N-1); and The maximum value among the data amounts of the multiple hybrid continuous operation segments is used as the third memory usage of the hybrid stage, and the third memory usage is used as the memory configuration size of the training process of the artificial intelligence model.

10. The memory management method according to claim 1, characterized in that: The method further comprises: Executing a mediator code, wherein the mediator code is configured to: Retrieving the plurality of architecture parameters of the artificial intelligence model; Execute the memory configuration size operation; performing configuration of the target memory region; and The transmission of the relay data between the acceleration processor module, the random access memory and the storage device is managed.

11. The memory management method according to claim 10, characterized in that: The method further comprises: The intermediary program code provides target memory area information to the framework layer, wherein the target memory area information includes: The starting address of the target memory area; The size of the target memory area; and access rights to the target memory area; The framework layer is an artificial intelligence computing development framework, and the processor executes the framework layer to provide a program development interface to the artificial intelligence model.

12. The memory management method according to claim 10, characterized in that: The target memory area may be further configured in a high bandwidth memory of the acceleration processor module.

13. The memory management method according to claim 3, characterized in that: The back propagation phase includes: Based on a predefined data set, comparing the output result of the artificial intelligence model with the corresponding target result in the data set to calculate an error value; and According to the error value, a back propagation algorithm is used to perform gradient calculation on the parameters of each model layer in the artificial intelligence model to obtain gradient values ​​corresponding to the multiple parameters.

14. The memory management method according to claim 3, characterized in that: The parameter updating phase includes: Calculating the parameter update direction of the artificial intelligence model according to the preset optimization algorithm and the gradient value calculated in the back propagation stage; and According to the parameter update direction, the parameter values ​​in the artificial intelligence model are adjusted.

15. An artificial intelligence processing system, characterized in that: include: processor; Accelerated processor modules to execute artificial intelligence models; Random Access Memory; a storage device, wherein the processor is electrically connected to the acceleration processor module, the random access memory, and the storage device, wherein the processor is configured to: Obtaining multiple architecture parameters of the artificial intelligence model; According to the plurality of architecture parameters, obtaining a memory configuration size used in a training process of the artificial intelligence model; According to the memory configuration size, a target memory area is configured in the random access memory, wherein the target memory area is commonly used by multiple training processing stages in the training process of the artificial intelligence model, wherein in each training processing stage, a corresponding data processing operation is performed on each model layer in the artificial intelligence model according to the processing order of the training processing stage; as well as The relay data corresponding to each training processing stage in the training process of the artificial intelligence model is temporarily stored in the target memory area, wherein the relay data is transmitted between the acceleration processor module, the random access memory and the storage device.

16. The artificial intelligence processing system according to claim 15, characterized in that: The processor is further configured to: acquiring data from the artificial intelligence model, and writing the data to the target memory area as the relay data; and The relay data is written from the target memory area to the storage device.

17. The artificial intelligence processing system according to claim 15, characterized in that: The multiple training processing stages include a back propagation stage and a parameter update stage, and the step of obtaining the memory configuration size used in the training process of the artificial intelligence model includes: The processor obtains a first memory usage of the back-propagation phase of the training process of the artificial intelligence model; The processor obtains a second memory usage in the parameter update phase of the training process of the artificial intelligence model; and The processor determines the memory allocation size according to the first memory usage and the second memory usage.

18. The artificial intelligence processing system according to claim 17, characterized in that: Determining the memory configuration size according to the first memory usage and the second memory usage includes: The processor selects the largest of the first memory usage and the second memory usage as the memory configuration size for the training process of the artificial intelligence model.

19. The artificial intelligence processing system according to claim 17, characterized in that: The target memory area includes a first memory usage space configured for the back-propagation phase and a second memory usage space configured for the parameter update phase, wherein the first memory usage space and the second memory usage space at least partially overlap.

20. The artificial intelligence processing system according to claim 17, characterized in that: The plurality of architectural parameters include: Model architecture information, including the parameter quantity and input data dimension of each model layer of the M model layers of the artificial intelligence model; The continuous operation section width is used to set the total number of multiple model layers to be executed simultaneously; Batch size, which represents the amount of data that needs to be processed during the training process; Input data specifications, including image size or text sequence length; The number of computing units represents the total number of multiple target acceleration processors in the acceleration processor module for executing the artificial intelligence model; and Parameter precision configuration, which defines the precision type and number of bytes of weight parameters, gradient parameters, and optimization parameters.

21. The artificial intelligence processing system according to claim 20, characterized in that: The step of obtaining the first memory usage in the back propagation stage of the training process of the artificial intelligence model includes: The processor calculates the first parameter memory requirement of each model layer according to the first parameter amount of each model layer of the M model layers in the back propagation stage, the first number of bytes defined by the parameter precision configuration and the number of operation units; The processor calculates the activation value memory requirement of each model layer according to the batch size, the input data specification and the input data dimension and the number of activation value bytes defined by the parameter precision configuration; The processor calculates the first model layer data volume corresponding to each model layer according to the first parameter memory requirement and the activation value memory requirement; The processor acquires a plurality of first continuous operation segment data amounts according to the continuous operation segment width and the first model layer data amount corresponding to each model layer, wherein the i-th first continuous operation segment data amount is the sum of the first model layer data amounts from the i-th model layer to the j-th model layer, wherein i is 1 to (M-(N-1)), N is the continuous operation segment width, and j is i+(N-1); and The processor uses a maximum value among the data amounts of the plurality of first continuous operation segments as the first memory usage.

22. The artificial intelligence processing system according to claim 21, characterized in that: The step of obtaining the second memory usage in the parameter updating stage of the training process of the artificial intelligence model comprises: The processor calculates the second parameter memory requirement of each model layer according to the second parameter amount of each model layer of the M model layers in the parameter updating phase, the second number of bytes defined by the parameter precision configuration and the number of operation units; The processor obtains the second model layer data volume corresponding to each model layer according to the second parameter memory requirement; The processor acquires a plurality of second continuous operation segment data amounts according to the continuous operation segment width and the second model layer data amount corresponding to each model layer, wherein the i-th second continuous operation segment data amount is the sum of the second model layer data amounts from the i-th model layer to the j-th model layer, wherein i is 1 to (M-(N-1)), N is the continuous operation segment width, and j is i+(N-1); and The processor uses a maximum value among the data amounts of the plurality of second continuous operation segments as the second memory usage.

23. The artificial intelligence processing system according to claim 22, characterized in that: The processor is further configured to: Determining whether the training process of the artificial intelligence uses a hybrid phase corresponding to the back-propagation phase and the parameter update phase; and React to the decision using the hybrid stage: According to the continuous operation segment width, the first model layer data volume and the second model layer data volume corresponding to each model layer, a plurality of hybrid continuous operation segment data volumes are obtained, wherein the i-th hybrid continuous operation segment data volume is the sum of the first model layer data volume and the second model layer data volume from the i-th model layer to the j-th model layer, wherein i is 1 to (M-(N-1)), N is the continuous operation segment width, and j is i+(N-1); and The maximum value among the data amounts of the multiple hybrid continuous operation segments is used as the third memory usage of the hybrid stage, and the third memory usage is used as the memory configuration size of the training process of the artificial intelligence model.

24. The artificial intelligence processing system according to claim 15, characterized in that: The intermediary program code is used to: Retrieving the plurality of architecture parameters of the artificial intelligence model; Execute the memory configuration size operation; performing configuration of the target memory region; and The transmission of the relay data between the acceleration processor module, the random access memory and the storage device is managed.

25. The artificial intelligence processing system according to claim 24, characterized in that: Also includes: The intermediary program code provides target memory area information to the framework layer, wherein the target memory area information includes: The starting address of the target memory area; The size of the target memory area; and access rights to the target memory area; The framework layer is an artificial intelligence computing development framework, and the processor executes the framework layer to provide a program development interface to the artificial intelligence model.

26. The artificial intelligence processing system according to claim 24, characterized in that: The target memory area may be further configured in a high bandwidth memory of the acceleration processor module.

27. The artificial intelligence processing system according to claim 17, characterized in that: The back propagation phase includes: Based on a predefined data set, comparing the output result of the artificial intelligence model with the corresponding target result in the data set to calculate an error value; and According to the error value, a back propagation algorithm is used to perform gradient calculation on the parameters of each model layer in the artificial intelligence model to obtain gradient values ​​corresponding to the multiple parameters.

28. The artificial intelligence processing system according to claim 17, characterized in that: The parameter updating phase includes: Calculating the parameter update direction of the artificial intelligence model according to the preset optimization algorithm and the gradient value calculated in the back propagation stage; and According to the parameter update direction, the parameter values ​​in the artificial intelligence model are adjusted.

29. A computer program product, characterized in that The method comprises: comprising: comprising: executing, via a processor of an artificial intelligence processing system, an intermediary program code to: Obtaining a plurality of architectural parameters of an artificial intelligence model, wherein the artificial intelligence model is executed by an accelerated processor module of the artificial intelligence processing system; According to the plurality of architecture parameters, obtaining a memory configuration size used in a training process of the artificial intelligence model; Allocate a target memory area in a random access memory of the artificial intelligence processing system according to the memory configuration size, wherein the target memory area is commonly used by multiple training processing stages in the training process of the artificial intelligence model; as well as The relay data corresponding to each training processing stage in the training process of the artificial intelligence model is temporarily stored in the target memory area, wherein the relay data is transmitted between the acceleration processor module, the random access memory and the storage device of the artificial intelligence processing system.

Images