WO2026001038A1 - 人工智能ai模型的训练方法、装置及计算设备 - Google Patents

人工智能ai模型的训练方法、装置及计算设备

Info

Publication number
WO2026001038A1
WO2026001038A1 PCT/CN2025/079244 CN2025079244W WO2026001038A1 WO 2026001038 A1 WO2026001038 A1 WO 2026001038A1 CN 2025079244 W CN2025079244 W CN 2025079244W WO 2026001038 A1 WO2026001038 A1 WO 2026001038A1
Authority
WO
WIPO (PCT)
Prior art keywords
module
computation
model
task
reverse
Prior art date
Legal status (The legal status is an assumption and is not a legal conclusion. Google has not performed a legal analysis and makes no representation as to the accuracy of the status listed.)
Pending
Application number
PCT/CN2025/079244
Other languages
English (en)
French (fr)
Inventor
李星辰
徐华
江健宇
Current Assignee (The listed assignees may be inaccurate. Google has not performed a legal analysis and makes no representation or warranty as to the accuracy of the list.)
Huawei Technologies Co Ltd
Original Assignee
Huawei Technologies Co Ltd
Priority date (The priority date is an assumption and is not a legal conclusion. Google has not performed a legal analysis and makes no representation as to the accuracy of the date listed.)
Filing date
Publication date
Application filed by Huawei Technologies Co Ltd filed Critical Huawei Technologies Co Ltd
Publication of WO2026001038A1 publication Critical patent/WO2026001038A1/zh
Pending legal-status Critical Current
Anticipated expiration legal-status Critical

Links

Classifications

    • GPHYSICS
    • G06COMPUTING OR CALCULATING; COUNTING
    • G06NCOMPUTING ARRANGEMENTS BASED ON SPECIFIC COMPUTATIONAL MODELS
    • G06N3/00Computing arrangements based on biological models
    • G06N3/02Neural networks
    • G06N3/04Architecture, e.g. interconnection topology
    • GPHYSICS
    • G06COMPUTING OR CALCULATING; COUNTING
    • G06NCOMPUTING ARRANGEMENTS BASED ON SPECIFIC COMPUTATIONAL MODELS
    • G06N3/00Computing arrangements based on biological models
    • G06N3/02Neural networks
    • G06N3/08Learning methods
    • G06N3/082Learning methods modifying the architecture, e.g. adding, deleting or silencing nodes or connections

Definitions

  • This application relates to the field of artificial intelligence (AI) technology, and in particular to a training method, apparatus and computing device for an artificial intelligence (AI) model.
  • AI artificial intelligence
  • Module A is taken as the module that needs to be recomputed.
  • the accelerator card sequentially executes the forward computation of modules A and B.
  • the accelerator card discards the intermediate activation values of module A and stores the output activation values of module A.
  • the accelerator card performs the backward computation of module B based on the output activation values of module A stored during the forward computation process.
  • the forward computation of module A is re-executed to obtain the intermediate activation values and output activation values of module A (i.e., the recomputation of module A is performed).
  • the backward computation of module A is performed based on the intermediate activation values of module A.
  • the intermediate activation values of module A are discarded, reducing the memory usage of the accelerator card.
  • the accelerator card since the recalculation and reverse calculation of module A need to be executed together, and the reverse calculation of module A is executed after the reverse calculation of module B, in order to ensure that the accelerator card can obtain the output activation value of module A when performing the reverse calculation of module B, the accelerator card must retain the output activation value of module A when performing the forward calculation of module A and retain it until after the reverse calculation of module B, resulting in a waste of accelerator card memory resources.
  • This application provides a training method, apparatus, and computing device for an artificial intelligence (AI) model, which can save memory resources of the accelerator card during the training process of the AI model.
  • AI artificial intelligence
  • this application provides a method for training an AI model, the method comprising:
  • forward computation is performed on multiple modules of the AI model.
  • These modules include a first module and a second module.
  • the first module is the module that needs to be recomputed during the backward computation process
  • the second module is the module that needs to use the output data of the first module during the forward computation process.
  • the AI model performs reverse computation on multiple modules. Before performing reverse computation on the second module, the first module is recomputed. Based on the data obtained from recomputing the first module, the second module is then reverse computed.
  • the output data of the first module is discarded after the forward computation of the second module is completed.
  • the recomputation of the first module is scheduled to be performed before the backward computation of the second module. In this way, when performing the backward computation of the second module, it can rely on the data obtained by recomputing the first module.
  • the second module is the module following the first module among a plurality of modules.
  • the AI model includes multiple network layers, with at least one network layer executed on each accelerator card.
  • the at least one network layer includes multiple modules executed sequentially, with the second module being the last of the multiple modules.
  • the AI model is a Transformer model
  • the network layer is a Transformer layer
  • this application provides a method for training an artificial intelligence (AI) model, the method comprising:
  • multiple modules of the AI model are identified, including a first module to be recomputed.
  • a recomputation task is sent to the accelerator card. There is a dependency between the first reverse computation task and the recomputation task.
  • training requests are executed by N accelerator cards in a pipelined parallel manner, where N is a positive integer;
  • a recomputation task is sent to the accelerator card, including: before sending the first reverse computation task to the first accelerator card among N accelerator cards, a recomputation task is sent to the first accelerator card, the first reverse computation task is the reverse computation task of the first module, and the first accelerator card is used to execute the computation tasks related to the first module.
  • the method further includes:
  • the second accelerator card is used to execute computation tasks related to the third module.
  • the third module is the next module after the first module.
  • the second reverse computation task is the reverse computation task of the third module.
  • a first reverse computation task is sent to the first accelerator card.
  • the first reverse computation task instructs the first module to be reverse computed based on the execution result of the second reverse computation task.
  • training requests are executed by M accelerator cards in parallel via a virtual pipeline, where M is a positive integer.
  • multiple modules of the AI model are determined, including: in response to the training request, the first module is determined based on the number of virtual pipelines, the number of pipelines, and the number of M accelerator cards;
  • a recomputation task is sent to the accelerator card, including: before sending the first reverse computation task to the third accelerator card among the M accelerator cards, a recomputation task is sent to the third accelerator card.
  • the first reverse computation task is the reverse computation task of the first module
  • the third accelerator card is used to execute the computation task related to the first module.
  • the method also includes: sending a forward computation task of the first module to the third accelerator card, wherein the forward computation task of the first module instructs that after performing forward computation on the first module, the intermediate activation value of the first module is discarded.
  • embodiments of this application provide an AI model training apparatus, which includes at least one functional unit for executing the AI model training method provided in the first aspect or any possible implementation thereof.
  • embodiments of this application provide an AI model training apparatus, which includes at least one functional unit for executing the AI model training method provided in the second aspect or any possible implementation thereof.
  • this application provides a computing device including a processor and a memory, wherein the processor is configured to execute at least a segment of program code stored in the memory to enable the computing device to implement the AI model training method provided by the first aspect or any possible implementation thereof, or to implement the AI model training method provided by the second aspect or any possible implementation thereof.
  • embodiments of this application provide a computing system including a host and at least one accelerator card.
  • the computing system is used to implement the training method for an AI model provided by the first aspect or any possible implementation of the first aspect, or to implement the training method for an AI model provided by the second aspect or any possible implementation of the second aspect.
  • this application provides a computer program product that, when run on a computing device, causes the computing device to implement the AI model training method provided by the first aspect or any possible implementation thereof, or to implement the AI model training method provided by the second aspect or any possible implementation thereof.
  • the computer program product can be a software installation package; when the aforementioned AI model training method needs to be implemented, the computer program product can be downloaded and executed on the computing device.
  • Figure 1 is a schematic diagram illustrating the application scenario of a related technology and this application;
  • Figure 2 is a schematic diagram of an implementation environment provided in an embodiment of this application.
  • Figure 3 is a schematic diagram of the structure of a computing device provided in an embodiment of this application.
  • Figure 4 is a flowchart of an AI model training method provided in an embodiment of this application.
  • FIG. 6 is a flowchart of another AI model training method provided in an embodiment of this application.
  • FIG. 7 is a flowchart of another AI model training method provided in an embodiment of this application.
  • Figure 11 is a schematic diagram of another training method for an AI model based on virtual pipeline parallelism provided in an embodiment of this application;
  • Figure 13 is a schematic diagram of the structure of another AI model training device provided in an embodiment of this application.
  • AI models are a class of mathematical algorithm models that use machine learning concepts to solve practical problems.
  • AI models include a large number of parameters and calculation formulas (or calculation rules).
  • An operator is a computational unit or function that runs on a computing device.
  • neural network layers and even the entire model are composed of operators, which correspond to the computational logic within the neural network layers.
  • a convolutional layer is an operator
  • the weight summation process in a fully-connected layer (FC layer) is also an operator.
  • Acceleration cards also known as acceleration devices, accelerators, or acceleration chips
  • acceleration devices are a type of specialized hardware accelerator or computer system designed to accelerate AI applications, especially neural networks, machine vision, and machine learning. For example, they provide computational power to perform calculations on matrices and vectors, thereby accelerating the computation of AI models.
  • acceleration devices include graphics processing units (GPUs), neural network processing units (XPUs), intelligent processing units (IPUs), tensor processing units (TPUs), domain-specific architecture (DSA) chips, and so on; this application is not limited to these.
  • peak memory usage refers to the maximum amount of memory consumed by the accelerator card during AI model training.
  • the magnitude of peak memory usage depends on several factors, such as model size (number of layers, number of parameters), training data size, batch size, and storage of intermediate computation results (e.g., activation values generated from forward computation).
  • model size number of layers, number of parameters
  • training data size e.g., training data size
  • batch size e.g., training data size
  • intermediate computation results e.g., activation values generated from forward computation.
  • activation values generated from forward computation need to be reserved for backward computation. Therefore, before backward computation begins, the number of activation values generated from forward computation increases significantly, leading to excessive memory consumption and potentially reaching peak memory usage, resulting in insufficient memory.
  • peak memory usage is generally reached during forward computation, and after backward computation begins, memory usage gradually decreases as backward computation progresses.
  • Recomputation refers to the process of training an AI model by not retaining the intermediate activation values of each module during forward computation. Instead, the forward computation of the corresponding module is re-executed before the backward computation of each module to generate new intermediate activation values. After the backward computation of the module is completed, these intermediate activation values are discarded, thereby alleviating the problem of excessively high memory peaks on the accelerator card.
  • Pipeline parallelism is a model parallel training method that splits different layers of the model onto different devices (such as accelerator cards), with each device responsible for training a few layers. This reduces the memory consumption of a single device, thereby enabling large-scale model training.
  • a model might have 12 network layers. Accelerator card 0 executes layers 0-3, accelerator card 1 executes layers 4-7, and accelerator card 2 executes layers 8-11. The accelerator cards exchange data via a communication link.
  • card 0 performs forward computation on layers 0-3 and sends the result to card 1.
  • Card 1 performs forward computation on layers 4-7 and sends the result to card 2.
  • Card 2 then performs forward computation on layers 8-11.
  • the backward computation process is similar, starting from card 2 to obtain gradient information, which is then passed back to card 1 and card 0.
  • this pipeline parallelism is also called naive pipeline parallelism.
  • VPP Virtual pipeline parallelism
  • pipeline parallelism while keeping the number of devices constant, it divides the training into more pipeline stages.
  • all layers of the model are divided into vpp parts. After each part is completed, the result from the last device is sent to the first device, and the pipeline parallelism of the next part continues until all parts are completed.
  • Card 0 executes layers 0-1, card 1 executes layers 2-3, card 2 executes layers 4-5, card 0 executes layers 6-7, card 1 executes layers 8-9, and card 2 executes layers 10-11.
  • Micro-batch pipeline parallelism is an extended training method of pipeline parallelism.
  • the samples used in one iteration form a batch (global batch).
  • the gradients generated based on these samples are accumulated and used for weight updates.
  • the utilization rate of equipment in pipeline parallelism can be improved, and the idle time of the equipment can be reduced.
  • the accelerator card executes forward or backward computations of a few layers based on a micro-batch each time.
  • a bubble is a phenomenon in parallel pipelines where a device may be idle, waiting for the result from another device before continuing its computation. This idle state is considered a waste of computing power.
  • FIG. 1 is a schematic diagram illustrating a related technology and the application scenario of this application.
  • Module A is taken as the module that needs to be recomputed.
  • the accelerator card sequentially executes the forward computation of module A and module B.
  • the accelerator card executes the forward computation of module A, it discards the intermediate activation value of module A and stores the output activation value Output A of module A.
  • the accelerator card executes the reverse computation of module B based on the output activation value of module A stored during the forward computation process.
  • module A After that, it re-executes the forward computation of module A to obtain the intermediate activation value and output activation value Output A of module A (that is, to execute the recomputation of module A). Then, it executes the reverse computation of module A based on the intermediate activation value of module A. After executing the reverse computation of module A, the intermediate activation value of module A is discarded to reduce the memory occupation of the accelerator card.
  • the accelerator card since the recalculation and reverse computation of module A need to be executed together, and the reverse computation of module A is executed after the reverse computation of module B, in order to ensure that the accelerator card can obtain the output activation value Output A of module A when executing the reverse computation of module B, the accelerator card must retain the output activation value Output A of module A when executing the forward computation of module A, and retain it until after the reverse computation of module B, resulting in a waste of the accelerator card's memory resources.
  • this application provides an implementation method for using recomputation during AI model training, which can effectively reduce the peak memory usage of the accelerator card and reduce memory resource waste.
  • the recomputation and backward computation of module A are separated into two functions. That is, the recomputation is extracted and implemented as a separate function. In this way, the forward computation, backward computation, and recomputation involved in the AI model training process can be scheduled independently.
  • the accelerator card sequentially executes the forward computations of modules A and B.
  • the accelerator card discards the intermediate activation values of module A and stores the output activation value Output A.
  • the accelerator card After executing the forward computation of module B, the actual storage of Output A is discarded, but Output A is retained in the logical view. This ensures the accuracy of the computational logic while freeing up accelerator card memory space.
  • the accelerator card first recompiles module A to obtain the intermediate and output activation values Output A. Then, based on these values, it sequentially executes the backward computations of modules B and A. That is, when executing the backward computation of module B, the accelerator card can rely on the recomputed output activation value Output A of module A. Furthermore, after executing the backward computation of module A, the accelerator card discards the intermediate activation values of module A, reducing memory usage. Therefore, the technical solution provided in this application increases the flexibility of recomputation, reduces the peak memory usage of the accelerator card, and improves model training performance.
  • Figure 1 above is an example using two modules in the model.
  • the model includes module A, module B, and module C, and module A is the module to be recomputed
  • the recomputation of module A can be scheduled between the reverse computation of module C and the reverse computation of module B.
  • the reverse computation of module B depends on the output activation value of module A
  • scheduling the recomputation of module A before the reverse computation of module B, or scheduling the recomputation of module A after the accelerator card reaches its memory peak can minimize the memory usage of the accelerator card by the output activation value of module A, thereby improving the model training performance.
  • Figure 2 is a schematic diagram of an implementation environment provided by an embodiment of this application. As shown in Figure 2, the implementation environment includes a computing system, which includes a host 100 and at least one accelerator card 200, and the host 100 and the at least one accelerator card 200 are connected in communication.
  • a computing system which includes a host 100 and at least one accelerator card 200, and the host 100 and the at least one accelerator card 200 are connected in communication.
  • the host 100 is used to provide AI services.
  • the host 100 can control the accelerator card 200 to execute computational tasks involved in the AI model training process.
  • the host 100 sends forward computation tasks, backward computation tasks, and recomputation tasks of each module in the AI model to the accelerator card 200, thereby performing forward computation, backward computation, and recomputation on each module.
  • This process can also be understood as loading the computational tasks involved in the AI model training process into the accelerator card 200 for execution.
  • the number of hosts 100 can be one or more, and this application does not limit this.
  • Accelerator card 200 provides computing power for the training process of AI models, thereby accelerating the training process.
  • accelerator card 200 can be a GPU, XPU, IPU, TPU, DSA chip, etc., and this application is not limited to these.
  • accelerator card 200 receives AI model computation tasks (such as forward computation tasks, backward computation tasks, recomputation tasks, etc.) sent by host 100, calls the corresponding operator, processes the input data of the operator, and obtains the output data, that is, the execution result of the computation task.
  • the number of accelerator cards 200 can be one or more, and this application does not limit this.
  • Pipeline parallelism includes, but is not limited to, naive pipelined parallelism, virtual pipelined parallelism, micro-batch pipelined parallelism, etc.
  • the aforementioned host 100 and accelerator card 200 can be integrated into a single server or configured separately; this application does not impose any limitations on this. Illustratively, taking the integration of host 100 and accelerator card 200 into a single server as an example, host 100 and accelerator card 200 are connected via a peripheral component interconnect express (PCIe) link, and data exchange between host 100 and accelerator card 200 occurs through the PCIe link.
  • the aforementioned server can be a standalone physical server, a server cluster composed of multiple physical servers, a distributed file system, or a cloud server providing basic cloud computing services such as cloud services, cloud databases, cloud computing, cloud functions, cloud storage, network services, cloud communication, middleware services, domain name services, security services, content delivery networks (CDNs), and big data and artificial intelligence platforms.
  • CDNs content delivery networks
  • the computing device can also be referred to as a cloud platform (i.e., short for cloud computing platform), which refers to services based on hardware and software resources, providing computing, network, and storage capabilities.
  • Cloud computing platforms leverage the internet to remotely process and analyze massive amounts of data before returning it to users. This approach offers advantages such as large scale, distributed computing, virtualization, high availability, scalability, on-demand service, and security. Cloud platforms can rapidly deploy and distribute configurable computing resources with minimal management overhead and low interaction complexity between users and service providers.
  • the networks mentioned above include, but are not limited to, data center networks, storage area networks (SANs), local area networks (LANs), metropolitan area networks (MANs), wide area networks (WANs), mobile, wired or wireless networks, private networks, or virtual private networks (VPNs).
  • technologies and/or formats including Hypertext Markup Language (HTML) and Extensible Markup Language (XML), are used to represent data exchanged over the network.
  • conventional encryption technologies such as Secure Sockets Layer (SSL), Transport Layer Security (TLS), Virtual Private Networks (VPNs), and Internet Protocol Security (IPsec) can be used to encrypt all or part of the links.
  • custom and/or dedicated data communication technologies can be used to replace or supplement the aforementioned data communication technologies.
  • the hardware structure of the host 100 in the above computing system is described below.
  • FIG3 is a schematic diagram of the structure of a computing device provided in an embodiment of this application.
  • the computing device 300 includes a memory 301, a processor 302, a communication interface 303, and a bus 304.
  • the memory 301, the processor 302, and the communication interface 303 are interconnected via the bus 304.
  • Memory 301 may be read-only memory (ROM) or other types of static storage devices capable of storing static information and instructions, random access memory (RAM) or other types of dynamic storage devices capable of storing information and instructions, or electrically erasable programmable read-only memory (EEPROM), compact disc read-only memory (CD-ROM) or other optical disc storage, optical disc storage (including compressed optical discs, laser discs, optical discs, digital universal optical discs, Blu-ray discs, etc.), magnetic disk storage media or other magnetic storage devices, or any other medium capable of carrying or storing desired program code having an instruction or data structure form and accessible by a computer, but not limited thereto.
  • memory 301 is used to store at least a piece of program code. When the program code stored in memory 301 is executed by processor 302, processor 302 performs the steps performed by the host in the following method embodiments.
  • the processor 302 may be a network processor (NP), a central processing unit (CPU), a graphics processing unit (GPU), an application-specific integrated circuit (ASIC), or an integrated circuit used to control the execution of the program in this application.
  • the processor 302 may be a single-core (single-CPU) processor or a multi-core (multi-CPU) processor. The number of processors 302 may be one or more.
  • the communication interface 303 uses a transceiver module, such as a transceiver, to enable communication between the computing device 300 and other devices or communication networks. For example, data can be acquired through the communication interface 303.
  • a transceiver module such as a transceiver
  • the memory 301 and the processor 302 can be set separately or integrated together.
  • Bus 304 may include a pathway for transmitting information between various components of computing device 300 (e.g., memory 301, processor 302, communication interface 303).
  • Figure 4 is a flowchart of an AI model training method provided in an embodiment of this application. As shown in Figure 4, taking the application of this method to a host in a computing system as an example, the method schematically includes the following steps 401 to 403.
  • the multiple modules include a first module and a second module.
  • the first module is the module that needs to be recomputed during the backward computation process
  • the second module is the module that needs to use the output data of the first module during the forward computation process.
  • the computing system includes a host and accelerator cards.
  • the host can access a wired or wireless network to provide AI services to users.
  • Performing forward computation on multiple modules of the AI model means that the host sends forward computation tasks for each module to the accelerator cards related to the training requests, and controls the accelerator cards to execute the forward computation of the corresponding modules according to the received tasks.
  • the AI model includes multiple modules. Any module in the AI model can be one or more network layers, or a component or basic element within a network layer. Furthermore, any module can be a non-linear or linear module. This application does not limit the granularity or type of module division in the AI model.
  • the multiple modules include a first module and a second module.
  • the first module is the module that needs to be recalculated during the backward computation process. That is, during the forward computation of the first module, the intermediate activation values of the first module can be discarded and re-executed during the subsequent backward computation to generate new intermediate activation values. This reduces the memory usage of the accelerator card.
  • the second module is the module that needs to use the output data of the first module during the forward computation process. In other words, during the forward computation process, the forward computation of the second module depends on the output data of the first module, which is the output activation value of the first module. The same principle applies below. For example, the second module may be the next module after the first module among multiple modules.
  • the AI model includes multiple network layers, and at least one network layer is executed on each accelerator card related to the training request.
  • This at least one network layer includes the aforementioned multiple modules executed sequentially, and the second module is the last module among multiple modules.
  • an AI model might be a Transformer model, and its network layers might be Transformer layers, also known as Transformer blocks.
  • a Transformer block typically refers to the basic building block of a Transformer model, which may include multi-head self-attention, a feedforward network, and components such as normalization and residual connections. Multiple Transformer blocks are stacked sequentially to form a complete Transformer model, enabling effective feature extraction and representation learning from the input sequence.
  • this network layer includes sequentially executed modules A, B, C, and D
  • the first module can be module A, B, or C
  • the second module is module D.
  • the host computer compiles operators based on the AI model's framework, i.e., generates operators for the AI model.
  • operator generation involves three processes: input tensor description, weight data transformation, and output tensor description.
  • input tensor description information such as the input dimension and memory size of each operator is calculated, and the format of the operator's input data is defined.
  • weight data transformation the weight parameters used by the operator undergo data format, shape transformation, and data compression.
  • the output tensor description information such as the output dimension and memory size of the operator is calculated.
  • the host computer After compiling operators based on the AI model's framework, the host computer adapts the operators to backend device interfaces (i.e., adapts the upper-layer operators to the lower-layer hardware devices), identifies the operator attributes, generates corresponding computation tasks based on the identified operators, and sends the computation tasks to the accelerator cards related to the training request.
  • backend device interfaces i.e., adapts the upper-layer operators to the lower-layer hardware devices
  • performing forward computation on the second module means that the host sends a forward computation task for the second module to the accelerator card, and controls the accelerator card to perform forward computation on the second module based on the output data of the first module according to the received forward computation task.
  • discarding the output data of the first module means that the host sends a discard task to the accelerator card after the forward computation of the second module is completed, and controls the accelerator card to discard the output data of the first module based on the discard task.
  • the discard task indicates that the actual storage of the output data of the first module in the accelerator card should be discarded, but the logical view of the output data should be retained. This ensures the accuracy of the computational logic while freeing up the accelerator card's memory space. It should be understood that since the output data of the first module is not used temporarily after the forward computation of the second module, the discard task can be executed to discard the output data and free up the accelerator card's memory space.
  • performing reverse computation on multiple modules of the AI model means that the host sends the reverse computation task for each module to the accelerator card related to the training request, and controls the accelerator card to execute the reverse computation of the corresponding module according to the received reverse computation task.
  • the first module is recomputed, and the second module is reverse computed based on the data obtained from the recomputation of the first module. That is, before sending the reverse computation task for the second module to the accelerator card, the recomputation task for the first module is sent to the accelerator card, so that the accelerator card executes the reverse computation task for the second module based on the data obtained from the recomputation of the first module.
  • the recomputation of the first module is scheduled to be performed before the reverse computation of the second module.
  • the reverse computation of the second module can depend on the data obtained from the recomputation of the first module. That is, the output data of the first module does not need to be stored from the forward computation process all the way to the reverse computation process of the second module.
  • the output data of the first module is discarded after the forward computation of the second module is completed.
  • the recomputation of the first module is scheduled to be performed before the backward computation of the second module. In this way, when performing the backward computation of the second module, it can rely on the data obtained by recomputing the first module.
  • the computing system can have one or multiple accelerator cards; that is, the implementation method of using recomputation in AI model training provided in this application is applicable to both single and multiple accelerator cards. Based on this, the following describes several other AI model training methods provided in this application with reference to the embodiments shown in Figures 5 to 11.
  • Figure 5 is a flowchart of another AI model training method provided in an embodiment of this application. As shown in Figure 5, the method is described using an application to a computing system as an example. Schematic, taking the interaction between the host and the accelerator card as an example, the method includes the following steps 501 to 503.
  • the host In response to a training request for the AI model, the host identifies multiple modules of the AI model, including a first module to be recomputed.
  • the host can access a wired or wireless network to provide AI services to the user.
  • the AI model includes multiple modules. Any module in the AI model can be one or more network layers, or a component or basic element within a network layer. Furthermore, any module can be a non-linear or linear module. This application does not limit the granularity or type of module division in the AI model.
  • the multiple modules include a first module, which is the module to be recalculated. That is, during the forward computation of the first module, its intermediate activation values are not retained. Instead, in the subsequent backward computation, the forward computation of the first module is re-executed to generate new intermediate activation values.
  • This method reduces the memory usage of the accelerator card.
  • this application does not limit the position of the first module in the AI model.
  • the first module can be determined according to business needs.
  • the first module may be the first module among multiple modules, or each module in the AI model may be the first module, and so on. It should be understood that in this embodiment, the first module refers to the module that needs to be recalculated, and its number can be one or more; this application does not limit this.
  • the host generates forward computation tasks for each module, backward computation tasks for each module, and recomputation tasks for the first module.
  • the host for any one of the multiple modules, the host generates a forward computation task and a backward computation task for that module. Furthermore, for the first module, the host generates a recomputation task for that module. Specifically, the forward computation task for the first module instructs that forward computation be performed on the first module. After the forward computation, the intermediate activation values of the first module are discarded, and the output data of the first module is stored. This process essentially separates the recomputation and backward computation of the first module in the AI model into two functions, extracting the recomputation as a separate function. Thus, the forward computation, backward computation, and recomputation involved in the AI model training process can be independently scheduled.
  • this application does not limit the timing of when the host generates the computing tasks corresponding to each module.
  • the host can generate each computing task sequentially according to the training framework of the AI model, and when any computing task is generated, it can send the computing task to the acceleration card related to the training request for execution.
  • the host can generate the computing tasks corresponding to each module according to the training framework of the AI model, and then send the computing tasks to the acceleration card related to the training request sequentially according to the execution order of each computing task, and so on.
  • This process can be configured according to business needs.
  • the host Before sending the first reverse computation task to the accelerator card related to the training request, the host sends the recomputation task of the first module to the accelerator card. There is a dependency between the first reverse computation task and the recomputation task.
  • the first reverse computation task is a reverse computation task among multiple modules' reverse computation tasks.
  • the dependency between the first reverse computation task and the recomputation task means that the execution of the first reverse computation task depends on the execution result of the recomputation task.
  • the first reverse computation task can be determined based on the AI model's training framework.
  • the host Before sending the first reverse computation task to the accelerator card, the host sends the recomputation task of the first module to the accelerator card, allowing the accelerator card to depend on the execution result of the recomputation task when executing the first reverse computation task.
  • This method of separately scheduling recomputation tasks improves the flexibility of recomputation, allowing users to schedule recomputation tasks to appropriate locations according to business needs, thereby reducing the memory usage of the accelerator card during AI model training or improving model training performance.
  • the AI model training method after determining multiple modules of the AI model according to the training request of the AI model, forward computation tasks and backward computation tasks corresponding to each module are generated. Furthermore, for the first module among the multiple modules to be recomputed, a recomputation task for the first module is generated separately. When distributing computation tasks to the acceleration card related to the training request, the recomputation task is scheduled before the first backward computation task that has a dependency on it. By scheduling the recomputation task separately, the flexibility of recomputation is effectively improved.
  • this application provides an AI model training method that schedules recomputation tasks separately.
  • the first reverse computation task can be determined according to the training framework of the AI model.
  • the training method of the above AI model will be further explained below with different scenarios as examples.
  • Figure 6 is a flowchart of another AI model training method provided in an embodiment of this application. As shown in Figure 6, the method is applied to a computing system, which includes a host and an accelerator card. Schematic, the method includes the following steps 601 to 614.
  • the host In response to a training request for an AI model, the host identifies multiple modules of the AI model, including a first module to be recomputed.
  • multiple modules also include a second module, which is a module that needs to use the output data of the first module during the forward computation process.
  • the second module is the next module after the first module among multiple modules.
  • the host generates forward computation tasks for each module, backward computation tasks for each module, and recomputation tasks for the first module.
  • the host sends the forward computation task of the first module to the accelerator card related to the training request.
  • the accelerator card performs the forward computation task of the first module.
  • the accelerator card performs forward calculations on the first module to obtain the intermediate activation value and output data of the first module, discards the intermediate activation value of the first module, and stores the output data of the first module in the memory of the accelerator card.
  • the host sends the forward computation task of the second module to the accelerator card.
  • the forward computation task of the second module indicates that forward computation should be performed on the second module of the AI model.
  • the accelerator card performs the forward computation task of the second module.
  • the accelerator card performs forward computation on the second module based on the output data of the first module to obtain the intermediate activation value and output data of the second module, and stores the intermediate activation value and output data of the second module in the memory of the accelerator card.
  • the second module can also be a module to be recalculated. For example, if the second module and the first module are both modules to be recalculated, then in this step, after obtaining the intermediate activation value and output data of the second module, the accelerator card discards the intermediate activation value of the second module and stores the output data of the second module in the memory of the accelerator card. This application does not limit this.
  • the host sends a drop task to the accelerator card, which instructs the output data of the first module to be discarded.
  • the discard task instruction to discard the output data of the first module means discarding the actual storage of the output data of the first module in the accelerator card, while retaining the logical view of the output data. In this way, the accuracy of the calculation logic can be ensured while freeing up the memory space of the accelerator card.
  • Accelerator card performs discard task.
  • a discard task can be executed to discard the output data and free up the memory space of the accelerator card.
  • the host sends a recomputation task for the first module to the accelerator card.
  • the recomputation task instructs the first module to be recomputed.
  • Accelerator cards perform recalculation tasks.
  • the accelerator card recalculates the first module to obtain the intermediate activation value and output data of the first module.
  • the execution result of the recalculation task includes the output data of the first module.
  • the host sends a first reverse computation task to the accelerator card.
  • This first reverse computation task is the reverse computation task of the second module.
  • the accelerator card performs the first reverse computation task.
  • the accelerator card performs reverse computation on the second module based on the output data of the first module stored in memory and the intermediate activation value of the second module.
  • the output data of the first module which the accelerator card relies on for reverse computation of the second module, is obtained through the recomputation task of the first module, rather than being stored during the forward computation process.
  • the output data of the first module does not need to be stored from the forward computation process all the way to the reverse computation process of the second module.
  • the host sends the reverse computation task of the first module to the accelerator card.
  • the accelerator card performs the reverse computation task of the first module.
  • the accelerator card performs reverse computation on the first module based on the intermediate activation values of the first module stored in memory. It should be understood that if the first module is the first module of the AI model, the accelerator card also needs to combine the samples input to the AI model when performing the first reverse computation task; if there are other modules before the first module, the accelerator card also needs to combine the output data of the other modules when performing the first reverse computation task, which will not be elaborated here.
  • the first and second modules in the AI model are used as examples for illustration. It should be understood that in some embodiments, if the AI model includes more modules, between steps 608 and 609 above, there is also a step in which the host sequentially sends forward computing tasks of other modules and reverse computing tasks of other modules to the accelerator card.
  • the second module is the next module after the first module.
  • the second module is a subsequent module of the first module.
  • the AI model includes multiple network layers, and at least one network layer is executed on each accelerator card associated with the training request. This at least one network layer includes the aforementioned multiple modules executed sequentially, with the second module being the last module among the multiple modules.
  • the AI model is a Transformer model, and the network layers are Transformer layers, also known as Transformer blocks. It should be understood that a Transformer block typically refers to the basic building block in a Transformer model, which may include multi-head self-attention, feedforward networks, and components such as normalization and residual connections.
  • the network layer executed on the accelerator card is as follows: if the network layer includes modules A, B, C, and D executed sequentially, then the first module can be module A, B, or C, and the second module is module D. Based on this, by using the above method to schedule the recomputation of at least one of modules A, B, and C before the reverse computation task of module D, the memory usage on the accelerator card can be minimized, thereby improving model training performance.
  • FIG 7 is a flowchart of another AI model training method provided in an embodiment of this application. As shown in Figure 7, the method is applied to a computing system, which includes a host and N accelerator cards. Schematic, the method includes the following steps 701 to 703.
  • the host In response to a training request for an AI model, the host identifies multiple modules of the AI model, including a first module to be recomputed.
  • the host computer is communicatively connected to N accelerator cards.
  • the N accelerator cards execute the training requests of the AI model in a pipelined parallel manner, where N is a positive integer.
  • Pipeline parallelism includes, but is not limited to, naive pipelined parallelism, virtual pipelined parallelism, micro-batch pipelined parallelism, etc.
  • the various modules of the AI model are each executed with corresponding computational tasks by the N accelerator cards.
  • the host generates forward computation tasks for each module, backward computation tasks for each module, and recomputation tasks for the first module.
  • FIG 8 is a schematic diagram of a pipelined parallel AI model training method provided in this embodiment.
  • N accelerator devices based on micro-batch and virtual pipelined parallelism as an example
  • F represents forward computation
  • B represents backward computation
  • R represents recomputation (assuming recomputation is performed on every layer).
  • the eight network layers are divided into two groups: layers 1 to 4 form one group, and layers 5 to 8 form the other.
  • Card 0 executes layers 1 and 5, card 1 executes layers 2 and 6, card 2 executes layers 3 and 7, and card 3 executes layers 4 and 8.
  • Unfilled boxes represent the earlier network layers in the model (i.e., layers 1 to 4), and filled boxes represent the later network layers in the model (i.e., layers 5 to 8).
  • the numbers 1-8 in the boxes represent eight micro-batches b1 to b8, respectively. It should be noted that Figure 8 is shown to illustrate the principle of the scheme.
  • the dimensions of the boxes shown in the figure are only illustrative and do not constitute a limitation on the execution time of each computation task. For example, in practical applications, the execution time of the backward computation task is usually longer than that of the forward computation task, such as the execution time of the backward computation task being twice that of the forward computation task.
  • the unfilled box F1 indicates forward computation of layer 1 based on micro-batch b1
  • the filled box F1 indicates forward computation of layer 5 based on micro-batch b1
  • the unfilled box B1 indicates backward computation of layer 1 based on micro-batch b1
  • the filled box B1 indicates backward computation of layer 5 based on micro-batch b1
  • the unfilled box R1 indicates recomputation of layer 1 based on micro-batch b1
  • the filled box R1 indicates recomputation of layer 5 based on micro-batch b1.
  • the forward computation task of the first module is the unfilled box F1 in the time sequence corresponding to card 0, and the recomputation task of the first module is the unfilled box R1 in the time sequence corresponding to card 0;
  • the first module is the fifth layer of the AI model, then the forward computation task of the first module is the filled box F1 in the time sequence corresponding to card 0, and the recomputation task of the first module is the filled box R1 in the time sequence corresponding to card 0;
  • the first module is the second layer of the AI model, then the forward computation task of the first module is the unfilled box F1 in the time sequence corresponding to card 1, and the recomputation task of the first module is the unfilled box R1 in the time sequence corresponding to card 1;
  • the first module is the sixth layer of the AI model, then the forward computation task of the first module is the filled box F1 in the time sequence corresponding to card 1, and the reco
  • the host Before sending the first reverse computation task to the first accelerator card among N accelerator cards, the host sends the recomputation task of the first module to the first accelerator card.
  • the first reverse computation task is the reverse computation task of the first module.
  • the first accelerator card is used to perform computational tasks related to the first module.
  • the first reverse computation task is the unfilled box B1 in the time sequence corresponding to card 0
  • the recomputation task of the first module is the unfilled box R1 in the time sequence corresponding to card 0.
  • the host sends a second reverse computation task to the second accelerator card among the N accelerator cards.
  • the second accelerator card is used to execute computation tasks related to the third module, which is the next module after the first module.
  • the second reverse computation task is the reverse computation task of the third module.
  • the host also sends a first reverse computation task to the first accelerator card, which instructs to perform reverse computation on the first module based on the execution result of the second reverse computation task.
  • the pipeline's "bubbles" can be utilized for recomputation, thereby reducing the waste of computing power.
  • card 0 executes the filled box B6 (reverse calculation of layer 5 based on micro-batch b6) corresponding to card 0.
  • this solution fills in the bubbles indicated by the arrows in the diagram.
  • the filled boxes R6 and B6 corresponding to card 0 are bound together; these two tasks need to wait for card 1 to send the execution result of its corresponding filled box B6 before they can execute.
  • the technical solution provided in this application can effectively reduce the bubbles that occur during pipeline parallelism, thereby improving the utilization rate of the accelerator card.
  • Figure 9 is a flowchart of another AI model training method provided in an embodiment of this application. As shown in Figure 9, the method is applied to a computing system, which includes a host and M accelerator cards. Schematic, the method includes the following steps 901 to 903.
  • the host determines multiple modules of the AI model, including a first module to be recomputed.
  • M accelerator cards execute AI model training requests in parallel via virtual pipelines, where M is a positive integer.
  • the host determines the first module based on the number of virtual pipelines, the number of pipelines, and the number of M accelerator cards. This process essentially means that the host treats a portion of the AI model as modules to be recomputed, thereby freeing up memory space on the accelerator cards.
  • FIG 10 is a schematic diagram of a training method for an AI model based on virtual pipelined parallelism provided in this application embodiment.
  • F represents forward computation
  • B represents backward computation
  • R represents recomputation.
  • the eight network layers are divided into two groups: layers 1 to 4 form one group, and layers 5 to 8 form the other.
  • Card 0 executes layers 1 and 5, card 1 executes layers 2 and 6, card 2 executes layers 3 and 7, and card 3 executes layers 4 and 8.
  • Unfilled boxes represent the earlier network layers in the model (i.e., layers 1 to 4), and filled boxes represent the later network layers in the model (i.e., layers 5 to 8).
  • the numbers 1-8 in the boxes represent eight micro-batches b1 to b8, respectively.
  • Figure 10 is shown to illustrate the principle of the scheme. The dimensions of the boxes shown in the figure are only illustrative and do not constitute a limitation on the execution time of each computation unit. For example, in practical applications, the execution time of the backward computation task is usually longer than that of the forward computation task, such as the execution time of the backward computation task being twice that of the forward computation task.
  • the modules corresponding to unfilled boxes F5, F6, and F7 can be identified as the first module, and the first module can be recomputed.
  • the host generates forward computation tasks for each module, backward computation tasks for each module, and recomputation tasks for the first module.
  • the host Before sending the first reverse computation task to the third accelerator card among the M accelerator cards, the host sends a recomputation task to the third accelerator card.
  • the first reverse computation task is the reverse computation task of the first module
  • the third accelerator card is used to execute the computation tasks related to the first module.
  • the first reverse computation task is the unfilled boxes B5, B6, and B7 in the timing sequence corresponding to card 0.
  • the accelerator card executes the forward computation task of the first module, it can discard the intermediate activation values of that first module, thus freeing up the accelerator card's memory space.
  • the host sends the forward computation task of the first module to the third accelerator card.
  • the forward computation task of the first module instructs that after performing forward computation on the first module, the intermediate activation values of the first module should be discarded.
  • the unfilled box R5 (based on the recalculation of the first layer based on micro-batch b5) can be scheduled to be executed at the bubble before the unfilled box B3
  • the unfilled box R6 (based on the recalculation of the first layer based on micro-batch b6) can be scheduled to be executed at the bubble before the unfilled box B4
  • the unfilled box R7 (based on the recalculation of the first layer based on micro-batch b7) can be scheduled to be executed at the bubble before the filled box B5.
  • the intermediate activation values of the first layer can be discarded, freeing up the accelerator card's memory space while maintaining performance.
  • the recalculation scheduling strategy can be determined according to business needs, and this application does not limit this.
  • the host uses the following scheduling method to issue tasks to that accelerator card to achieve the pipelined parallel execution effect shown in Figure 10 above.
  • this scheduling method includes the following stages:
  • Phase 3 Based on the number of pipelines, micro-batch counts, and virtual pipeline counts, the host sends Z first computation task groups to the accelerator cards.
  • Each first computation task group includes one forward computation task and one backward computation task.
  • Z b ⁇ vpp - pp ⁇ vpp - pp + 2 + 2i. Where pp - 2 - 2i ⁇ 0, the forward computation task of the first first computation task group in the Z first computation task groups needs to be recomputed.
  • the i-th accelerator card as card 0 as an example, there are a total of 6 Z first computation task groups, as indicated by the arrows in the figure.
  • Phase 4 Based on the number of pipelines, the host sends W second computation task groups to the accelerator cards.
  • Each second computation task group includes one recomputation task and one reverse computation task.
  • Schematic, W max(pp-1-2i, 0). Referring to Figure 11, taking the i-th accelerator card as card 0 as an example, there are W second computation task groups as indicated by the arrows in the figure, a total of 3 groups.
  • Phase 5 The host sends K reverse computation tasks to the accelerator cards based on the number of pipelines and the number of virtual pipelines.
  • K min(pp ⁇ vpp - 1, pp ⁇ vpp + pp - 2 - 2i).
  • K min(pp ⁇ vpp - 1, pp ⁇ vpp + pp - 2 - 2i).
  • the above scheduling method is only an example.
  • the modules that need to be recalculated can be set according to business needs, and the execution time of the recalculation task can be set, etc. This application does not limit this.
  • FIG. 12 is a schematic diagram of the structure of an AI model training device provided in an embodiment of this application. As shown in Figure 12, the device includes a forward computation execution unit 1201, a discard unit 1202, and a backward computation execution unit 1203.
  • the forward computation execution unit 1201 is used to perform forward computation on multiple modules of the AI model in response to the training request of the AI model.
  • the multiple modules include a first module and a second module.
  • the first module is the module that needs to be recomputed during the backward computation process
  • the second module is the module that needs to use the output data of the first module during the forward computation process.
  • the discard unit 1202 is used to discard the output data of the first module after the forward computation of the second module is completed;
  • the reverse computation execution unit 1203 is used to perform reverse computation on multiple modules of the AI model. Before performing reverse computation on the second module, it performs recomputation on the first module and performs reverse computation on the second module based on the data obtained from recomputing the first module.
  • the second module is the module following the first module among a plurality of modules.
  • the AI model includes multiple network layers, with at least one network layer executed on each accelerator card.
  • the at least one network layer includes multiple modules executed sequentially, with the second module being the last of the multiple modules.
  • the AI model is a Transformer model
  • the network layer is a Transformer layer
  • the AI model training device provided in the above embodiments is only illustrated by the division of the above functional modules when training the AI model.
  • the above functions can be assigned to different functional modules as needed, that is, the internal structure of the device can be divided into different functional modules to complete all or part of the functions described above.
  • the AI model training device and the AI model training method embodiments provided in the above embodiments belong to the same concept, and their specific implementation process can be found in the method embodiments, which will not be repeated here.
  • FIG. 13 is a schematic diagram of another AI model training device provided in an embodiment of this application. As shown in Figure 13, the device includes a determining unit 1301, a generating unit 1302, and a sending unit 1303.
  • the determining unit 1301 is used to determine multiple modules of the AI model in response to a training request of the AI model, the multiple modules including a first module to be recomputed;
  • the generation unit 1302 is used to generate the forward computation task of each module, the backward computation task of each module, and the recomputation task of the first module in the multiple modules.
  • the sending unit 1303 is used to send a recalculation task to the accelerator card before sending the first reverse computation task to the accelerator card related to the training request. There is a dependency between the first reverse computation task and the recalculation task.
  • training requests are executed by N accelerator cards in a pipelined parallel manner, where N is a positive integer;
  • the sending unit 1303 is used to send a recalculation task to the first accelerator card before sending the first reverse calculation task to the first accelerator card among N accelerator cards.
  • the first reverse calculation task is the reverse calculation task of the first module, and the first accelerator card is used to execute the calculation tasks related to the first module.
  • the transmitting unit 1303 is further configured to:
  • the second accelerator card is used to execute computation tasks related to the third module.
  • the third module is the next module after the first module.
  • the second reverse computation task is the reverse computation task of the third module.
  • a first reverse computation task is sent to the first accelerator card.
  • the first reverse computation task instructs the first module to be reverse computed based on the execution result of the second reverse computation task.
  • training requests are executed by M accelerator cards in parallel via a virtual pipeline, where M is a positive integer.
  • Determining unit 1301 is used to determine the first module in response to a training request, based on the number of virtual pipelines, the number of pipelines, and the number of M accelerator cards;
  • the sending unit 1303 is used to send a recalculation task to the third accelerator card before sending the first reverse calculation task to the third accelerator card among the M accelerator cards.
  • the first reverse calculation task is the reverse calculation task of the first module
  • the third accelerator card is used to execute the calculation tasks related to the first module.
  • the sending unit 1303 is also used to send the forward computation task of the first module to the third accelerator card.
  • the forward computation task of the first module indicates that after performing forward computation on the first module, the intermediate activation value of the first module is discarded.
  • the AI model training device provided in the above embodiments is only illustrated by the division of the above functional modules when training the AI model.
  • the above functions can be assigned to different functional modules as needed, that is, the internal structure of the device can be divided into different functional modules to complete all or part of the functions described above.
  • the AI model training device and the AI model training method embodiments provided in the above embodiments belong to the same concept, and their specific implementation process can be found in the method embodiments, which will not be repeated here.
  • first and second accelerator cards are used to distinguish identical or similar items with substantially the same function and purpose. It should be understood that there is no logical or temporal dependency between “first,” “second,” and “nth,” nor does it limit the quantity or execution order. It should also be understood that although the following description uses the terms “first,” “second,” etc., to describe various elements, these elements should not be limited by the terms. These terms are merely used to distinguish one element from another.
  • a first accelerator card can be referred to as a second accelerator card, and similarly, a second accelerator card can be referred to as a first accelerator card. Both the first and second accelerator cards can be accelerator cards, and in some cases, they can be separate and different accelerator cards.
  • the term "at least one" means one or more, and the term “multiple” means two or more.
  • multiple accelerator cards means two or more accelerator cards.
  • implementation can be achieved, in whole or in part, through software, hardware, firmware, or any combination thereof.
  • software When implemented in software, it can be implemented, in whole or in part, in the form of program structure information.
  • This program structure information includes one or more program instructions. When these program instructions are loaded and executed on a computing device, the processes or functions according to the embodiments of this application are generated, in whole or in part.

Landscapes

  • Engineering & Computer Science (AREA)
  • Theoretical Computer Science (AREA)
  • Physics & Mathematics (AREA)
  • Data Mining & Analysis (AREA)
  • General Health & Medical Sciences (AREA)
  • Biomedical Technology (AREA)
  • Biophysics (AREA)
  • Computational Linguistics (AREA)
  • Life Sciences & Earth Sciences (AREA)
  • Evolutionary Computation (AREA)
  • Artificial Intelligence (AREA)
  • Molecular Biology (AREA)
  • Computing Systems (AREA)
  • General Engineering & Computer Science (AREA)
  • General Physics & Mathematics (AREA)
  • Mathematical Physics (AREA)
  • Software Systems (AREA)
  • Health & Medical Sciences (AREA)
  • Management, Administration, Business Operations System, And Electronic Commerce (AREA)

Abstract

本申请公开了一种人工智能AI模型的训练方法、装置及计算设备,涉及人工智能AI技术领域,方法包括:在对AI模型的多个模块执行前向计算时,对于待执行重计算的第一模块和依赖于第一模块输出数据的第二模块,在对第二模块的前向计算执行完成后,丢弃第一模块的输出数据,在对多个模块执行反向计算时,将第一模块的重计算调度至第二模块的反向计算之前执行,如此,在执行第二模块的反向计算时,能够依赖对第一模块进行重计算得到的数据,通过这种灵活的重计算调度方式,无需将第一模块的输出数据从前向计算过程存储至反向计算过程,节约了加速卡的内存资源。

Description

人工智能AI模型的训练方法、装置及计算设备
本申请要求于2024年06月29日提交的申请号为202410869302.6、发明名称为“人工智能AI模型的训练方法、装置及计算设备”的中国专利申请的优先权,其全部内容通过引用结合在本申请中。
技术领域
本申请涉及人工智能(artificial intelligence,AI)技术领域,特别涉及一种人工智能AI模型的训练方法、装置及计算设备。
背景技术
随着AI技术的快速发展,一系列加速卡涌现出来,用于提供对矩阵、向量等进行计算的计算算力,以加速AI模型的计算过程。目前,在AI模型的训练过程中采用重计算(Recomputation),能够有效减少模型训练过程中对加速卡的内存占用。
相关技术中,以AI模型包括模块A和模块B为例,在模型训练过程中采用重计算的实施方式例如是:将模块A作为需要执行重计算的模块,在前向计算过程中,加速卡顺序执行模块A和模块B的前向计算,其中,加速卡在执行模块A的前向计算时,丢弃模块A的中间激活值,存储模块A的输出激活值;在反向计算过程中,加速卡基于前向计算过程中存储的模块A的输出激活值,执行模块B的反向计算,之后,重新执行一遍模块A的前向计算,得到模块A的中间激活值和输出激活值(也即执行模块A的重计算),进而基于模块A的中间激活值执行模块A的反向计算,在执行模块A的反向计算之后,丢弃模块A的中间激活值,减少对加速卡的内存占用。
上述方法中,由于模块A的重计算和反向计算需要绑定在一起执行,且模块A的反向计算在模块B的反向计算之后执行,因此,为确保加速卡在执行模块B的反向计算时能够获取到模块A的输出激活值,加速卡在执行模块A的前向计算时,必须保留模块A的输出激活值,并将其保留至模块B的反向计算之后,造成加速卡内存资源浪费。
发明内容
本申请实施例提供了一种人工智能AI模型的训练方法、装置及计算设备,能够在AI模型的训练过程中,节约加速卡的内存资源。
第一方面,本申请提供了一种AI模型的训练方法,该方法包括:
响应于AI模型的训练请求,对AI模型的多个模块执行前向计算,多个模块包括第一模块及第二模块,第一模块为后向计算过程中需要执行重计算的模块,第二模块为前向计算过程中需要使用第一模块的输出数据的模块;
在第二模块的前向计算执行完成后,丢弃第一模块的输出数据;
对AI模型的多个模块执行反向计算,在对第二模块执行反向计算之前,对第一模块执行重计算,并基于对第一模块进行重计算得到的数据,对第二模块执行反向计算。
通过上述方式,在对AI模型的多个模块执行前向计算时,对于待执行重计算的第一模块和依赖于第一模块输出数据的第二模块,在对第二模块的前向计算执行完成后,丢弃第一模块的输出数据,在对多个模块执行反向计算时,将第一模块的重计算调度至第二模块的反向计算之前执行,如此,在执行第二模块的反向计算时,能够依赖对第一模块进行重计算得到的数据,通过这种灵活的重计算调度方式,无需将第一模块的输出数据从前向计算过程存储至反向计算过程,节约了加速卡的内存资源。
在一些实施例中,第二模块为多个模块中第一模块的下一个模块。
在一些实施例中,AI模型包括多个网络层,每个加速卡上执行至少一个网络层,至少一个网络层包括顺序执行的多个模块,第二模块为多个模块中的最后一个模块。
在一些实施例中,AI模型为Transformer模型,网络层为Transformer层。
第二方面,本申请提供了一种人工智能AI模型的训练方法,方法包括:
响应于人工智能AI模型的训练请求,确定AI模型的多个模块,多个模块包括待执行重计算的第一模块;
生成多个模块中每个模块的前向计算任务、每个模块的反向计算任务以及第一模块的重计算任务;
在向训练请求相关的加速卡发送第一反向计算任务之前,向加速卡发送重计算任务,第一反向计算任务与重计算任务之间存在依赖关系。
在一些实施例中,训练请求由N个加速卡以流水线并行的方式执行,N为正整数;
在向训练请求相关的加速卡发送第一反向计算任务之前,向加速卡发送重计算任务,包括:在向N个加速卡中第一加速卡发送第一反向计算任务之前,向第一加速卡发送重计算任务,第一反向计算任务为第一模块的反向计算任务,第一加速卡用于执行第一模块相关的计算任务。
在一些实施例中,方法还包括:
向N个加速卡中第二加速卡发送第二反向计算任务,第二加速卡用于执行第三模块相关的计算任务,第三模块为第一模块的下一个模块,第二反向计算任务为第三模块的反向计算任务;
向第一加速卡发送第一反向计算任务,第一反向计算任务指示基于第二反向计算任务的执行结果,对第一模块进行反向计算。
在一些实施例中,训练请求由M个加速卡以虚拟流水线并行的方式执行,M为正整数;
响应于人工智能AI模型的训练请求,确定AI模型的多个模块,包括:响应于训练请求,基于虚拟流水线的数量、流水线的数量以及M个加速卡的数量,确定第一模块;
在向训练请求相关的加速卡发送第一反向计算任务之前,向加速卡发送重计算任务,包括:在向M个加速卡中第三加速卡发送第一反向计算任务之前,向第三加速卡发送重计算任务,第一反向计算任务为第一模块的反向计算任务,第三加速卡用于执行第一模块相关的计算任务;
方法还包括:向第三加速卡发送第一模块的前向计算任务,第一模块的前向计算任务指示对第一模块进行前向计算后,丢弃第一模块的中间激活值。
第三方面,本申请实施例提供了一种AI模型的训练装置,该装置包括至少一个功能单元,该至少一个功能单元用于执行前述第一方面或第一方面的任意一种可能的实现方式所提供的AI模型的训练方法。
第四方面,本申请实施例提供了一种AI模型的训练装置,该装置包括至少一个功能单元,该至少一个功能单元用于执行前述第二方面或第二方面的任意一种可能的实现方式所提供的AI模型的训练方法。
第五方面,本申请提供了一种计算设备,该计算设备包括处理器和存储器,处理器用于执行存储器中存储的至少一段程序代码,以使计算设备实现如前述第一方面或第一方面的任意一种可能的实现方式所提供的AI模型的训练方法,或者,实现如前述第二方面或第二方面的任意一种可能的实现方式所提供的AI模型的训练方法。
第六方面,本申请实施例提供了一种计算系统,该计算系统包括主机和至少一个加速卡,该计算系统用于实现如前述第一方面或第一方面的任意一种可能的实现方式所提供的AI模型的训练方法,或者,实现如前述第二方面或第二方面的任意一种可能的实现方式所提供的AI模型的训练方法。
第七方面,本申请提供了一种计算机可读存储介质,计算机可读存储介质用于存储至少一段程序代码,当该至少一段程序代码由计算设备执行时,使得计算设备实现如前述第一方面或第一方面的任意一种可能的实现方式所提供的AI模型的训练方法,或者,实现如前述第二方面或第二方面的任意一种可能的实现方式所提供的AI模型的训练方法。该存储介质包括但不限于易失性存储器,例如随机访问存储器,非易失性存储器,例如快闪存储器、硬盘(hard disk drive,HDD)、固态硬盘(solid state drive,SSD)。
第八方面,本申请提供了一种计算机程序产品,当计算机程序产品在计算设备上运行时,使得计算设备实现如前述第一方面或第一方面的任意一种可能的实现方式所提供的AI模型的训练方法,或者,实现如前述第二方面或第二方面的任意一种可能的实现方式所提供的AI模型的训练方法。该计算机程序产品可以为一个软件安装包,在需要实现前述AI模型的训练方法的情况下,可以下载该计算机程序产品并在计算设备上执行该计算机程序产品。
附图说明
图1是一种相关技术和本申请的应用场景示意图;
图2是本申请实施例提供的一种实施环境的示意图;
图3是本申请实施例提供的一种计算设备的结构示意图;
图4是本申请实施例提供的一种AI模型的训练方法的流程图;
图5是本申请实施例提供的另一种AI模型的训练方法的流程图;
图6是本申请实施例提供的另一种AI模型的训练方法的流程图;
图7是本申请实施例提供的另一种AI模型的训练方法的流程图;
图8是本申请实施例提供的一种基于流水线并行的AI模型的训练方法示意图;
图9是本申请实施例提供的又一种AI模型的训练方法的流程图;
图10是本申请实施例提供的一种基于虚拟流水线并行的AI模型的训练方法示意图;
图11是本申请实施例提供的另一种基于虚拟流水线并行的AI模型的训练方法示意图;
图12是本申请实施例提供的一种AI模型的训练装置的结构示意图;
图13是本申请实施例提供的另一种AI模型的训练装置的结构示意图。
具体实施方式
为使本申请的目的、技术方案和优点更加清楚,下面将结合附图对本申请实施方式作进一步地详细描述。需要说明的是,本申请所涉及的信息(包括但不限于用户设备信息、用户个人信息等)、数据(包括但不限于用于分析的数据、存储的数据、展示的数据等)以及信号,均为经用户授权或者经过各方充分授权的,且相关数据的收集、使用和处理需要遵守相关国家和地区的相关法律法规和标准。例如,本申请中涉及到的AI模型、计算任务的相关数据等都是在充分授权的情况下获取的。
为了方便理解,下面先对本申请涉及的关键术语和关键概念进行说明。
人工智能(artificial intelligence,AI)模型,是一类用机器学习思想解决实际问题的数学算法模型,通常,AI模型包括大量的参数和计算公式(或计算规则)。
算子(operator,OP),是指运行在计算设备上的计算单元或计算函数,在深度学习领域,神经网络层乃至整个模型均由算子组成,这些算子对应神经网络层中的计算逻辑。例如,卷积层(convolution layer)是一个算子;全连接层(fully-connected layer,FC layer)中的权值求和过程,是一个算子。
加速卡,或称加速设备(acceleration device)、加速器(accelerator)、加速芯片等,是一类专门的硬件加速器或计算机系统,旨在加速AI应用,尤其是神经网络、机器视觉和机器学习等。例如,用于提供对矩阵、向量等进行计算的计算算力,以加速AI模型的计算。示意性地,加速设备例如是图形处理器(graphics processing unit,GPU)、神经网络处理器(neural network processing unit,XPU)、智能处理器(intelligent processing unit,IPU)、张量处理器(tensor processing unit,TPU)、特定域架构(domain specific architecture,DSA)芯片,等等,本申请不限于此。
内存峰值,在本申请中是指AI模型训练过程中,加速卡的内存占用量达到的最大值。通常,内存峰值的大小取决于多个因素,如模型规模(层数、参数数量)、训练数据的大小、批量大小、中间计算结果(如前向计算产生的激活值)的存储。比如说,由于AI模型训练过程中,需要将前向计算产生的激活值留给反向计算使用,因此,在反向计算开始之前,前向计算产生的激活值会越来越多,导致内存占用过高,最终可能达到内存峰值,导致内存不足。换言之,内存峰值一般在前向计算时达到,在反向计算开始之后,内存占用随着反向计算的进行逐渐下降。
重计算(Recomputation),是指在AI模型的训练过程中,执行前向计算时不保留每个模块的中间激活值(activation),而是在执行每个模块的反向计算之前,重新执行一遍对应模块的前向计算,重新产生中间激活值,在执行完该模块的反向计算之后,丢弃这部分中间激活值,从而缓解加速卡内存峰值过高的问题。
流水线并行(pipeline parallelism,PP),是一种模型并行训练方式,通过将模型的不同层切分到不同的设备(如加速卡)上,每个设备负责其中某几层的训练,降低单个设备的显存消耗,从而实现大规模模型训练。比如说,模型包括12个网络层,加速卡0执行0-3层,加速卡1执行4-7层,加速卡2执行8-11层,加速卡之间通过通信链路进行数据交换。在前向计算过程中,卡0执行0-3层的前向计算,将执行结果送给卡1,卡1执行4-7层的前向计算,将执行结果发送给卡2,卡2执行8-11层的前向计算。反向计算过程同理,从卡2开始进行反向计算,得到梯度信息,依次传递回卡1和卡0。在一些场景下,这种流水线并行方式也称为朴素流水线并行。
虚拟流水线并行(virtual pipeline parallelism,VPP),是流水线并行的一种扩展训练方式,在流水线并行过程中,在设备数量不变的情况下,划分出更多的流水线阶段(pipeline stage)。示意性地,以虚拟流水线并行的规模为vpp为例,将模型的所有层分为vpp份,每一份做完流水线并行后,将最后一个设备的结果送给第一个设备,继续做下一份的流水线并行,直到所有份全部做完。比如说,以前述流水线并行的示例为基础,若虚拟流水线规模vpp=2,则将12个网络层分为两份,0-5层为一份,6-11层为另一份,其中,卡0执行0-1层,卡1执行2-3层,卡2执行4-5层,卡0执行6-7层,卡1执行8-9层,卡2执行10-11层。
微批次(micro batch)流水线并行,是流水线并行的一种扩展训练方式,通常,在模型训练过程中,一次迭代采用的样本组成一个批次(global batch),基于这些样本训练产生的梯度累加在一起,用于权重更新,将一个批次的样本划分为多个微批次,每个微批次包括至少一个样本,将微批次作为并行调度的基本单元,如此,能够提升流水线并行时设备的利用率,减小设备处于空闲状态的时间。例如,在流水线并行中,加速卡每次执行的任务为基于一个微批次的某几层的前向计算或反向计算。
气泡(bubble),是指在流水线并行时,某个设备可能需要等待另一个设备的执行结果,收到后才能继续计算,因此处于空闲状态。这种空闲状态被称为气泡,是对设备算力的一种浪费。
下面对本申请的应用场景及实施环境进行介绍。
本申请应用于在AI模型训练过程中采用重计算以减少对加速卡内存占用的场景中,其中,本申请对于AI模型的类型、加速卡的类型均不作限定。下面参考图1,对本申请的应用场景进行介绍。图1是一种相关技术和本申请的应用场景示意图。
如图1中(a)图所示,相关技术中,以AI模型包括模块A和模块B为例,在模型训练过程中采用重计算的实施方式例如是:将模块A作为需要执行重计算的模块,在前向计算过程中,加速卡顺序执行模块A和模块B的前向计算,其中,加速卡在执行模块A的前向计算时,丢弃模块A的中间激活值,存储模块A的输出激活值OutputA;在反向计算过程中,加速卡基于前向计算过程中存储的模块A的输出激活值,执行模块B的反向计算,之后,重新执行一遍模块A的前向计算,得到模块A的中间激活值和输出激活值OutputA(也即执行模块A的重计算),进而基于模块A的中间激活值执行模块A的反向计算,在执行模块A的反向计算之后,丢弃模块A的中间激活值,减少对加速卡的内存占用。可见,由于模块A的重计算和反向计算需要绑定在一起执行,且模块A的反向计算在模块B的反向计算之后执行,因此,为确保加速卡在执行模块B的反向计算时能够获取到模块A的输出激活值OutputA,加速卡在执行模块A的前向计算时,必须保留模块A的输出激活值OutputA,并将其保留至模块B的反向计算之后,造成加速卡内存资源浪费。
如图1中(b)图所示,本申请提供了一种在AI模型训练过程中采用重计算的实施方式,能够有效降低加速卡内存峰值,减少内存资源浪费。在本申请中,将模块A的重计算与反向计算拆分为两个函数,也即,将重计算单独抽离出来,实现为一个单独的函数,如此,AI模型训练过程中涉及的前向计算、反向计算以及重计算可以独立调度。例如,在前向计算过程中,加速卡顺序执行模块A和模块B的前向计算,其中,加速卡在执行模块A的前向计算时,丢弃模块A的中间激活值,存储模块A的输出激活值OutputA,并且,在执行模块B的前向计算后,丢弃OutputA的真实存储,在逻辑视图上保留OutputA,从而在释放加速卡内存空间的基础上,确保计算逻辑的准确性;在反向计算过程中,加速卡先执行模块A的重计算,得到模块A的中间激活值和输出激活值OutputA,进而基于模块A的中间激活值和输出激活值OutputA顺序执行模块B的反向计算和模块A的反向计算,也即,加速卡在执行模块B的反向计算时,能够依赖通过重计算得到的模块A的输出激活值OutputA,另外,在执行模块A的反向计算之后,加速卡丢弃模块A的中间激活值,减少对加速卡的内存占用。可见,本申请提供的技术方案增加了重计算的灵活性,能够降低加速卡的内存峰值,提升模型训练性能。
应理解,上述图1是以模型中两个模块为例进行举例说明的,在一些场景中,模型包括更多数量个模块时,同理适用于上述内容。比如说,模型包括模块A、模块B以及模块C,若模块A为待执行重计算的模块,在采用本申请提供的实施方式时,可以将模块A的重计算调度至模块C的反向计算与模块B的反向计算之间执行,换言之,由于模块B的反向计算依赖于模块A的输出激活值,因此,将模块A的重计算调度至模块B的反向计算之前,或者说,将模块A的重计算调度至加速卡达到内存峰值之后,如此,可以尽可能减少模块A的输出激活值对加速卡的内存占用,进而提升模型训练性能。
下面参考图2,对本申请的实施环境进行介绍。图2是本申请实施例提供的一种实施环境的示意图。如图2所示,该实施环境包括计算系统,计算系统包括主机(host)100和至少一个加速卡200,主机100和至少一个加速卡200之间通信连接。
主机100用于提供AI业务,在本申请实施例中,主机100能够控制加速卡200来执行在AI模型训练过程中所涉及的计算任务。例如,主机100响应于AI模型的训练请求,向加速卡200发送AI模型中各个模块的前向计算任务、后向计算任务以及重计算任务等等,实现对各个模块执行前向计算、反向计算以及重计算。这一过程也可以理解为,将AI模型训练过程中涉及的计算任务加载到加速卡200中执行。另外,主机100的数量可以是一个或多个,本申请对此不作限定。
加速卡200用于为AI模型的训练过程提供计算算力,以加速AI模型的训练过程。例如,加速卡200为GPU、XPU、IPU、TPU、DSA芯片等等,本申请不限于此。示意性地,加速卡200接收主机100发送的AI模型的计算任务(如前向计算任务、后向计算任务、重计算任务等),调用相应算子,对该算子的输入数据进行处理,得到输出数据,也即得到该计算任务的执行结果。另外,加速卡200的数量可以是一个或多个,本申请对此不作限定。在加速卡200的数量为N个的情况下,计算系统具备分布式计算能力,能够以流水线并行的方式执行AI模型的训练请求,N为正整数。其中,流水线并行包括但不限于朴素流水线并行、虚拟流水线并行、微批次流水线并行,等等。
上述主机100和加速卡200可以集成在一个服务器中,也可以分离设置,本申请对此不作限定。示意性地,以主机100和加速卡200集成在一个服务器中为例,主机100和加速卡200通过外围组件互连总线(peripheral component interconnect express,PCIe)链路通信连接,主机100和加速卡200之间通过PCIe链路进行数据交互。上述服务器可以是独立的物理服务器,或者是多个物理服务器构成的服务器集群或者分布式文件系统,又或者是提供云服务、云数据库、云计算、云函数、云存储、网络服务、云通信、中间件服务、域名服务、安全服务、内容分发网络(content delivery network,CDN)以及大数据和人工智能平台等基础云计算服务的云服务器。以计算设备为云服务器为例,计算设备也可以称为是一种云平台(即云计算平台的简称),是指基于硬件资源和软件资源的服务,提供计算、网络和存储能力。通过网络“云”将庞大的数据计算处理在远端进行处理和分析后返回给用户,具有大规模、分布式、虚拟化、高可用性、扩展性、按需服务以及安全性等特点。云平台可以以较小的管理代价,或者用户与业务提供者较低的交互复杂度,实现可配置计算资源的快速发放与发布。
另外,上述涉及的网络包括但不限于数据中心网络(data center network)、存储区域网(storage area network,SAN)、局域网(local area network,LAN)、城域网(metropolitan area network,MAN)、广域网(wide area network,WAN)、移动、有线或者无线网络、专用网络或者虚拟专用网络的任何组合。在一些实现方式中,使用包括超级文本标记语言(hyper text markup language,HTML)、可扩展标记语言(extensible markup language,XML)等技术和/或格式来代表通过网络交换的数据。此外还能够使用诸如安全套接字层(secure sockets layer,SSL)、传输层安全(transport layer security,TLS)、虚拟专用网络(virtual private network,VPN)、网际协议安全(internet protocol security,IPsec)等常规加密技术来加密所有或者部分链路。在另一些实施例中,还能够使用定制和/或专用数据通信技术取代或者补充上述数据通信技术。
下面对上述计算系统中主机100的硬件结构进行介绍。
本申请提供了一种计算设备,能够配置为上述主机100,参考图3,图3是本申请实施例提供的一种计算设备的结构示意图。如图3所示,该计算设备300包括存储器301、处理器302、通信接口303以及总线304。其中,存储器301、处理器302、通信接口303通过总线304实现彼此之间的通信连接。
存储器301可以是只读存储器(read-only memory,ROM)或可存储静态信息和指令的其它类型的静态存储设备,随机存取存储器(random access memory,RAM)或者可存储信息和指令的其它类型的动态存储设备,也可以是电可擦可编程只读存储器(electrically erasable programmable read-only memory,EEPROM)、只读光盘(compact disc read-only memory,CD-ROM)或其它光盘存储、光碟存储(包括压缩光碟、激光碟、光碟、数字通用光碟、蓝光光碟等)、磁盘存储介质或者其它磁存储设备、或者能够用于携带或存储具有指令或数据结构形式的期望的程序代码并能够由计算机存取的任何其它介质,但不限于此。示意性地,存储器301用于存储至少一段程序代码,当存储器301中存储的程序代码被处理器302执行时,处理器302用于执行下述方法实施例中主机所执行的步骤。
处理器302可以是网络处理器(network processor,NP)、中央处理器(central processing unit,CPU)、图像处理器(graphics processing unit,GPU)、特定应用集成电路(application-specific integrated circuit,ASIC)或用于控制本申请方案程序执行的集成电路。该处理器302可以是一个单核(single-CPU)处理器,也可以是一个多核(multi-CPU)处理器。该处理器302的数量可以是一个,也可以是多个。
通信接口303使用例如收发器一类的收发模块,来实现计算设备300与其他设备或通信网络之间的通信。例如,可以通过通信接口303获取数据。
其中,存储器301和处理器302可以分离设置,也可以集成在一起。
总线304可包括在计算设备300各个部件(例如,存储器301、处理器302、通信接口303)之间传送信息的通路。
基于上述图1至图3所示内容,对本申请提供的应用场景及实施环境进行了介绍,下面对本申请提供的AI模型的训练方法进行介绍。
图4是本申请实施例提供的一种AI模型的训练方法的流程图。如图4所示,以该方法应用于计算系统中的主机为例进行介绍,示意性地,该方法包括如下步骤401至步骤403。
401、响应于AI模型的训练请求,对AI模型的多个模块执行前向计算,多个模块包括第一模块及第二模块,第一模块为后向计算过程中需要执行重计算的模块,第二模块为前向计算过程中需要使用第一模块的输出数据的模块。
在本申请实施例中,计算系统包括主机和加速卡,主机能够接入有线网络或无线网络,为用户AI业务。对AI模型的多个模块执行前向计算也即是,主机向训练请求相关的加速卡发送每个模块的前向计算任务,控制加速卡根据接收到的前向计算任务,执行相应模块的前向计算。其中,AI模型包括多个模块,对于AI模型的任意一个模块,该模块可以是AI模型的一个或多个网络层,也可以是一个网络层中的组件、基本元素等,另外,对于任意一个模块,该模块可以是非线性模块,也可以是线性模块,本申请对AI模型中模块的划分粒度及类型不作限定。
多个模块包括第一模块和第二模块,其中,第一模块为后向计算过程中需要执行重计算的模块,也即是,在对第一模块进行前向计算的过程中,可以不保留该第一模块的中间激活值,而是在后续反向计算过程中,重新执行一遍第一模块的前向计算,重新产生中间激活值,通过这种方式,来降低对加速卡的内存占用。第二模块为前向计算过程中需要使用第一模块的输出数据的模块,也即是,在前向计算过程中,第二模块的前向计算依赖于第一模块的输出数据,第一模块的输出数据也即是第一模块的输出激活值,后文同理,不再赘述。比如说,第二模块为多个模块中第一模块的下一个模块,又比如,在一些实施例中,AI模型包括多个网络层,训练请求相关的每个加速卡上执行至少一个网络层,该至少一个网络层包括顺序执行的前述多个模块,第二模块为多个模块中的最后一个模块。例如,AI模型为Transformer模型,网络层为Transformer层,也称为Transformer block,应理解,Transformer block通常指的是Transformer模型中的基本构建层,可以包括多头自注意力机制(Multi-head Self-Attention)、前馈神经网络(Feed Forward Network)以及归一化和残差连接等组件。多个Transformer block依次堆叠,共同构成了完整的Transformer模型,以实现对输入序列的有效特征提取和表示学习。示意性地,以加速卡上执行的网络层为Transforme模型中任意一个Transformer block为例,若该网络层包括顺序执行的模块A、模块B、模块C以及模块D,则第一模块可以是模块A或模块B或模块C,第二模块为模块D。
示意性地,主机响应于AI模型的训练请求,基于AI模型的模型框架进行算子编译,也即是生成AI模型的算子,例如,算子生成分为输入张量描述、权重数据转换和输出张量描述三个流程。其中,在输入张量描述中,计算每个算子的输入维度、内存大小等信息,并定义算子输入数据的形式;在权重数据转换中,对算子使用的权重参数进行数据格式、形状转换、数据压缩等处理;在输出张量描述中,计算算子的输出维度、内存大小等信息。主机基于AI模型的模型框架进行算子编译后,对算子进行后端设备接口适配(也即将上层算子与底层硬件设备进行适配),并识别算子的属性,根据识别到的算子,生成相应的计算任务,将计算任务发送给训练请求相关的加速卡。
402、在第二模块的前向计算执行完成后,丢弃第一模块的输出数据。
在本申请实施例中,对第二模块执行前向计算也即是,主机向加速卡发送第二模块的前向计算任务,控制加速卡根据接收到的前向计算任务,基于第一模块的输出数据,对第二模块进行前向计算。在第二模块的前向计算执行完成后,丢弃第一模块的输出数据也即是,主机在第二模块的前向计算执行完成后,向加速卡发送丢弃任务,控制加速卡基于该丢弃任务,丢弃第一模块的输出数据。其中,丢弃任务指示丢弃第一模块的输出数据在加速卡中的真实存储,保留该输出数据的逻辑视图,如此,能够在释放加速卡内存空间的基础上,确保计算逻辑的准确性。应理解,由于在执行第二模块的前向计算之后,暂时不会用到第一模块的输出数据,因此,可以执行丢弃任务来丢弃该输出数据,释放加速卡的内存空间。
403、对AI模型的多个模块执行反向计算,在对第二模块执行反向计算之前,对第一模块执行重计算,并基于对第一模块进行重计算得到的数据,对第二模块执行反向计算。
在本申请实施例中,对AI模型的多个模块执行反向计算也即是,主机向训练请求相关的加速卡发送每个模块的反向计算任务,控制加速卡根据接收到的反向计算任务,执行相应模块的反向计算。在这一过程中,在对第二模块执行反向计算之前,对第一模块执行重计算,并基于对第一模块进行重计算得到的数据,对第二模块执行反向计算,也即是,在向加速卡发送第二模块的反向计算任务之前,向加速卡发送第一模块的重计算任务,使得加速卡基于对第一模块进行重计算得到的数据,执行第二模块的反向计算任务。换言之,将对第一模块的重计算调度至第二模块的反向计算之前执行,如此,在对第一模块进行重计算之后,对第二模块的反向计算能够依赖于对第一模块进行重计算得到的数据。也即是,第一模块的输出数据无需从前向计算过程一直存储至第二模块的反向计算过程。
通过上述方式,在对AI模型的多个模块执行前向计算时,对于待执行重计算的第一模块和依赖于第一模块输出数据的第二模块,在对第二模块的前向计算执行完成后,丢弃第一模块的输出数据,在对多个模块执行反向计算时,将第一模块的重计算调度至第二模块的反向计算之前执行,如此,在执行第二模块的反向计算时,能够依赖对第一模块进行重计算得到的数据,通过这种灵活的重计算调度方式,无需将第一模块的输出数据从前向计算过程存储至反向计算过程,节约了加速卡的内存资源。
另外,基于前述内容可知,对于AI模型中的第一模块,通过将第一模块的重计算与反向计算拆分为两个函数,将重计算单独抽离出来,实现为一个单独的函数,如此,AI模型训练过程中涉及的前向计算、反向计算以及重计算可以独立调度,通过这种将重计算单独调度的方式,提升了重计算的灵活性,使得用户能够根据业务需求将重计算调度至合适的位置,进而降低AI模型训练过程中加速卡的内存占用,或者,提升模型训练性能。而且,计算系统中加速卡的数量可以是一个,也可以是多个,也即是,本申请提供的在AI模型训练过程中采用重计算的实施方式既适用于单张加速卡,也适用于多张加速卡。基于此,下面参考图5至图11所示实施例,对本申请提供的另外几种AI模型的训练方法进行介绍。
图5是本申请实施例提供的另一种AI模型的训练方法的流程图。如图5所示,以该方法应用于计算系统为例进行介绍,示意性地,以主机和加速卡之间的交互为例,该方法包括如下步骤501至步骤503。
501、主机响应于AI模型的训练请求,确定AI模型的多个模块,其中,多个模块包括待执行重计算的第一模块。
在本申请实施例中,主机能够接入有线网络或无线网络,为用户AI业务。其中,AI模型包括多个模块,对于AI模型的任意一个模块,该模块可以是AI模型的一个或多个网络层,也可以是一个网络层中的组件、基本元素等,另外,对于任意一个模块,该模块可以是非线性模块,也可以是线性模块,本申请对AI模型中模块的划分粒度及类型不作限定。
多个模块包括第一模块,第一模块为待执行重计算的模块,也即是,在对第一模块进行前向计算的过程中,不保留该第一模块的中间激活值,而是在后续反向计算过程中,重新执行一遍第一模块的前向计算,重新产生中间激活值,通过这种方式,来降低对加速卡的内存占用。另外,本申请对于第一模块在AI模型中的位置不作限定,第一模块能够根据业务需求来确定,例如,第一模块为多个模块中的第一个模块,又例如,AI模型中每个模块均为第一模块,等等。应理解,第一模块在本实施例中用于指代需执行重计算的模块,其数量可以是一个,也可以是多个,本申请对此不作限定。
502、主机生成多个模块中每个模块的前向计算任务、每个模块的反向计算任务以及第一模块的重计算任务。
在本申请实施例中,对于多个模块中任意一个模块,主机生成该模块的前向计算任务和反向计算任务,并且,对于第一模块,主机生成该第一模块的重计算任务。其中,对于第一模块的前向计算任务,该前向计算任务指示对第一模块进行前向计算,并且,在前向计算后,丢弃第一模块的中间激活值,存储第一模块的输出数据。这一过程也即是,对于AI模型中的第一模块,将第一模块的重计算与反向计算拆分为两个函数,将重计算单独抽离出来,实现为一个单独的函数,如此,AI模型训练过程中涉及的前向计算、反向计算以及重计算可以独立调度。
另外,本申请对于主机生成各个模块所对应计算任务的时机不作限定,主机可以按照AI模型的训练框架,依次生成各个计算任务,并在生成任意一个计算任务的情况下,将该计算任务下发给训练请求相关的加速卡执行,主机也可以按照AI模型的训练框架,生成各个模块所对应的计算任务之后,按照各个计算任务的执行时序,将计算任务依次下发给训练请求相关的加速卡,等等,这一过程能够根据业务需求进行设置。
503、主机在向训练请求相关的加速卡发送第一反向计算任务之前,向加速卡发送第一模块的重计算任务,第一反向计算任务与重计算任务之间存在依赖关系。
在本申请实施例中,主机与加速卡之间通信连接,主机能够控制加速卡执行相应计算任务。第一反向计算任务是多个模块的反向计算任务中的反向计算任务,第一反向计算任务与重计算任务之间存在依赖关系是指,第一反向计算任务的执行依赖于重计算任务的执行结果。示意性地,第一反向计算任务能够根据AI模型的训练框架来确定,主机向加速卡发送第一反向计算任务之前,向加速卡发送第一模块的重计算任务,使得加速卡在执行第一反向计算任务时,能够依赖重计算任务的执行结果。通过这种将重计算任务单独调度的方式,提升了重计算的灵活性,使得用户能够根据业务需求将重计算任务调度至合适的位置,进而降低AI模型训练过程中加速卡的内存占用,或者,提升模型训练性能。
综上,在本申请实施例提供的AI模型的训练方法中,在根据AI模型的训练请求确定AI模型的多个模块之后,生成各个模块对应的前向计算任务和反向计算任务,并且,对于多个模块中待执行重计算的第一模块,单独生成该第一模块的重计算任务,在向训练请求相关的加速卡下发计算任务时,将该重计算任务调度至与其存在依赖关系的第一反向计算任务之前。通过这种将重计算任务单独调度的方式,有效提升了重计算的灵活性。
基于前述图5所示实施例可知,本申请提供了一种将重计算任务单独调度的AI模型训练方法,其中,第一反向计算任务能够根据AI模型的训练框架来确定,下面以不同场景为例,对上述AI模型的训练方法作进一步说明。
图6是本申请实施例提供的另一种AI模型的训练方法的流程图。如图6所示,该方法应用于计算系统,该计算系统包括主机和加速卡,示意性地,该方法包括如下步骤601至步骤614。
601、主机响应于AI模型的训练请求,确定AI模型的多个模块,其中,多个模块包括待执行重计算的第一模块。
其中,多个模块还包括第二模块,第二模块为前向计算过程中需要使用第一模块的输出数据的模块,例如,第二模块为多个模块中第一模块的下一个模块。
602、主机生成多个模块中每个模块的前向计算任务、每个模块的反向计算任务以及第一模块的重计算任务。
603、主机向训练请求相关的加速卡发送第一模块的前向计算任务。
604、加速卡执行第一模块的前向计算任务。
其中,加速卡对第一模块进行前向计算,得到第一模块的中间激活值和输出数据,丢弃第一模块的中间激活值,将第一模块的输出数据存储至加速卡的内存中。
605、主机向加速卡发送第二模块的前向计算任务。
其中,第二模块的前向计算任务指示对AI模型的第二模块进行前向计算。
606、加速卡执行第二模块的前向计算任务。
其中,加速卡基于第一模块的输出数据,对第二模块进行前向计算,得到第二模块的中间激活值和输出数据,将第二模块的中间激活值和输出数据存储至加速卡内存中。需要说明的是,第二模块也可以是待执行重计算的模块,例如,若第二模块和第一模块同为待执行重计算的模块,则在本步骤中,加速卡在得到第二模块的中间激活值和输出数据后,丢弃第二模块的中间激活值,将第二模块的输出数据存储至加速卡的内存中,本申请对此不作限定。
607、主机向加速卡发送丢弃任务,丢弃任务指示丢弃第一模块的输出数据。
在本申请实施例中,丢弃任务指示丢弃第一模块的输出数据是指,丢弃第一模块的输出数据在加速卡中的真实存储,保留该输出数据的逻辑视图,如此,能够在释放加速卡内存空间的基础上,确保计算逻辑的准确性。
608、加速卡执行丢弃任务。
应理解,由于加速卡在执行第二模块的前向计算任务之后,暂时不会用到第一模块的输出数据,因此,可以执行丢弃任务来丢弃该输出数据,释放加速卡的内存空间。
609、主机向加速卡发送第一模块的重计算任务,重计算任务指示对第一模块重新进行前向计算。
610、加速卡执行重计算任务。
其中,加速卡对第一模块重新进行前向计算,重新得到第一模块的中间激活值和输出数据,也即,重计算任务的执行结果包括第一模块的输出数据。
611、主机向加速卡发送第一反向计算任务,该第一反向计算任务为第二模块的反向计算任务。
612、加速卡执行第一反向计算任务。
其中,加速卡基于内存中存储的第一模块的输出数据以及第二模块的中间激活值,对第二模块进行反向计算。这一过程也即是,加速卡在对第二模块进行反向计算时所依赖的第一模块的输出数据是通过第一模块的重计算任务得到的,而不是在前向计算过程中所存储的,也即是,第一模块的输出数据无需从前向计算过程一直存储至第二模块的反向计算过程。
613、主机向加速卡发送第一模块的反向计算任务。
614、加速卡执行第一模块的反向计算任务。
其中,加速卡基于内存中存储的第一模块的中间激活值,对第一模块进行反向计算。应理解,若第一模块为AI模型的第一个模块,则加速卡在执行第一反向计算任务时,还需结合输入AI模型的样本;若第一模块之前还存在其他模块,则加速卡在执行第一反向计算任务时,还需结合其他模块的输出数据,在此不再赘述。
在上述步骤603至步骤614中,是以AI模型中的第一模块和第二模块为例进行说明的,应理解,在一些实施例中,若AI模型还包括更多个模块,在前述步骤608和步骤609之间,还包括主机向加速卡依次发送其他模块的前向计算任务以及其他模块的反向计算任务的步骤。
通过上述方式,在根据AI模型的训练请求确定AI模型的多个模块之后,生成各个模块对应的前向计算任务和反向计算任务,并且,对于多个模块中待执行重计算的第一模块,单独生成该第一模块的重计算任务,在向训练请求相关的加速卡下发计算任务时,将该重计算任务调度至与其存在依赖关系的第二模块的反向计算任务之前。通过这种将重计算任务单独调度的方式,有效提升了重计算的灵活性,而且,加速卡在执行第二模块的反向计算任务时,能够依赖基于重计算任务得到的第一模块的输出数据,如此,加速卡在执行第二模块的前向计算任务之后,可以及时丢弃第一模块的输出数据,释放加速卡内存空间,减少对加速卡的内存占用。
另外,在上述图6所示实施例中,第二模块是第一模块的下一个模块,在另一些实施例中,第二模块为第一模块的后续模块,在一些实施例中,AI模型包括多个网络层,训练请求相关的每个加速卡上执行至少一个网络层,该至少一个网络层包括顺序执行的前述多个模块,第二模块为多个模块中的最后一个模块。例如,AI模型为Transformer模型,网络层为Transformer层,也称为Transformer block,应理解,Transformer block通常指的是Transformer模型中的基本构建层,可以包括多头自注意力机制(Multi-head Self-Attention)、前馈神经网络(Feed Forward Network)以及归一化和残差连接等组件。多个Transformer block依次堆叠,共同构成了完整的Transformer模型,以实现对输入序列的有效特征提取和表示学习。示意性地,以加速卡上执行的网络层为Transforme模型中任意一个Transformer block为例,若该网络层包括顺序执行的模块A、模块B、模块C以及模块D,则第一模块可以是模块A或模块B或模块C,第二模块为模块D。基于此,采用上述方法,将模块A、模块B、模块C中至少一个模块的重计算调度至模块D的反向计算任务之前,可以尽可能降低对加速卡的内存占用,进而提升模型训练性能。
图7是本申请实施例提供的另一种AI模型的训练方法的流程图。如图7所示,该方法应用于计算系统,该计算系统包括主机和N个加速卡,示意性地,该方法包括如下步骤701至步骤703。
701、主机响应于AI模型的训练请求,确定AI模型的多个模块,其中,多个模块包括待执行重计算的第一模块。
在本申请实施例中,主机与N个加速卡通信连接,N个加速卡以流水线并行的方式执行AI模型的训练请求,N为正整数。其中,流水线并行包括但不限于朴素流水线并行、虚拟流水线并行、微批次流水线并行,等等。其中,AI模型的多个模块分别由N个加速卡执行相应的计算任务。
702、主机生成多个模块中每个模块的前向计算任务、每个模块的反向计算任务以及第一模块的重计算任务。
为便于更清楚地解释N个加速卡以流水线并行的方式执行训练请求的过程,在本实施例中,参考图8来介绍AI模型的训练方法。图8是本申请实施例提供的一种基于流水线并行的AI模型的训练方法示意图。如图8所示,计算系统包括N个加速设备,N=4,N个加速设备分别表示为卡0、卡1、卡2以及卡3,以N个加速卡基于微批次和虚拟流水线并行为例进行介绍,其中,微批次b=8,虚拟流水线vpp=2,AI模型包括8层。图中,F表示前向计算,B表示反向计算,R表示重计算(以每层均执行重计算为例),将8个网络层分为2份,第1层至第4层为一份,第5层至第8层为另一份,卡0执行第1层和第5层,卡1执行第2层和第6层,卡2执行第3层和第7层,卡3执行第4层和第8层,无填充的方框表示模型中前面的网络层(即第1层至第4层),有填充的方框表示模型中后面的网络层(即第5层至第8层),方框中的数字1-8分别表示8个微批次b1至b8。需要说明的是,图8所示是为了介绍方案原理,图中所示各个方框的尺寸仅为举例说明,并不构成对每个计算任务执行时长的限定,比如说,在实际应用中,反向计算任务的执行时长通常大于前向计算任务的执行时长,如反向计算任务的执行时长为前向计算任务的执行时长的2倍。
如图8所示,卡0对应的时序中,无填充方框F1表示基于微批次b1对第1层进行前向计算,有填充方框F1表示基于微批次b1对第5层进行前向计算,无填充方框B1表示基于微批次b1对第1层进行反向计算,有填充方框B1表示基于微批次b1对第5层进行反向计算,无填充方框R1表示基于微批次b1对第1层进行重计算,有填充方框R1表示基于微批次b1对第5层进行重计算。
举例而言,以微批次b1为例,若第一模块为AI模型的第1层,则第一模块的前向计算任务为卡0对应时序中的无填充方框F1,第一模块的重计算任务为卡0对应时序中的无填充方框R1;若第一模块为AI模型的第5层,则第一模块的前向计算任务为卡0对应时序中的有填充方框F1,第一模块的重计算任务为卡0对应时序中的有填充方框R1;若第一模块为AI模型的第2层,则第一模块的前向计算任务为卡1对应时序中的无填充方框F1,第一模块的重计算任务为卡1对应时序中的无填充方框R1;若第一模块为AI模型的第6层,则第一模块的前向计算任务为卡1对应时序中的有填充方框F1,第一模块的重计算任务为卡1对应时序中的有填充方框R1,等等,不再赘述。
703、主机在向N个加速卡中第一加速卡发送第一反向计算任务之前,向第一加速卡发送第一模块的重计算任务,第一反向计算任务为第一模块的反向计算任务。
其中,第一加速卡用于执行第一模块相关的计算任务。
以图8所示内容为例,若第一模块为AI模型的第1层,则第一反向计算任务为卡0对应时序中的无填充方框B1,第一模块的重计算任务为卡0对应时序中的无填充方框R1。应理解,由于在N个加速卡以流水线并行的方式执行训练请求的过程中,将重计算任务单独调度,使得重计算任务无需等待上一个加速卡的执行结果,只有第一模块的反向计算任务需等待上一个加速卡的执行结果,示意性地,主机向N个加速卡中第二加速卡发送第二反向计算任务,第二加速卡用于执行第三模块相关的计算任务,第三模块为第一模块的下一个模块,第二反向计算任务为第三模块的反向计算任务;向第一加速卡发送第一反向计算任务,第一反向计算任务指示基于第二反向计算任务的执行结果,对第一模块进行反向计算。如此,可以将流水线的气泡利用起来进行重计算,从而减少算力的浪费。比如说,继续参考图8,以第一加速卡为卡0,第二加速卡为卡1为例,由于将有填充方框R6(基于微批次b6对第5层的重计算)和有填充方框B6(基于微批次b6对第5层的反向计算)分离调度,因此,卡0在执行有填充方框B5(基于微批次b5对第5层的反向计算)之后,可以立即执行有填充方框R6(基于微批次b6对第5层的重计算),然后卡0在接收到卡1执行后发送的有填充方框B6(卡1对应时序中基于微批次b6对第6层的反向计算)的执行结果之后,执行卡0对应的有填充方框B6(基于微批次b6对第5层的反向计算)。相比于相关技术,填充了如图中箭头所指的气泡,即,在相关技术中,卡0对应的有填充方框R6和有填充方框B6是绑定在一起的,这两个任务需要一起等待卡1将卡1对应的有填充方框B6的执行结果发送过来才能执行。显然,相比于相关技术,采用本申请提供的技术方案能够有效减少流水线并行过程中出现的气泡,进而提升加速卡的利用率。
图9是本申请实施例提供的又一种AI模型的训练方法的流程图。如图9所示,该方法应用于计算系统,该计算系统包括主机和M个加速卡,示意性地,该方法包括如下步骤901至步骤903。
901、主机响应于AI模型的训练请求,确定AI模型的多个模块,其中,多个模块包括待执行重计算的第一模块。
其中,M个加速卡以虚拟流水线并行的方式执行AI模型的训练请求,M为正整数。示意性地,主机基于虚拟流水线的数量、流水线的数量以及M个加速卡的数量,确定第一模块。这一过程也即是,主机将AI模型中的部分模块作为待执行重计算的模块,通过这种方式,来释放加速卡内存空间。
为便于更清楚地解释M个加速卡以虚拟流水线并行的方式执行训练请求的过程,在本实施例中,参考图10来介绍AI模型的训练方法。图10是本申请实施例提供的一种基于虚拟流水线并行的AI模型的训练方法示意图。如图10所示,计算系统包括M个加速设备,M=4,M个加速设备分别表示为卡0、卡1、卡2以及卡3,以M个加速卡基于微批次和虚拟流水线并行为例进行介绍,其中,微批次b=8,虚拟流水线vpp=2,AI模型包括8层。图中,F表示前向计算,B表示反向计算,R表示重计算,将8个网络层分为2份,第1层至第4层为一份,第5层至第8层为另一份,卡0执行第1层和第5层,卡1执行第2层和第6层,卡2执行第3层和第7层,卡3执行第4层和第8层,无填充的方框表示模型中前面的网络层(即第1层至第4层),有填充的方框表示模型种后面的网络层(即第5层至第8层),方框中的数字1-8分别表示8个微批次b1至b8。需要说明的是,图10所示是为了介绍方案原理,图中所示各个方框的尺寸仅为举例说明,并不构成对每个计算单元执行时长的限定,比如说,在实际应用中,反向计算任务的执行时长通常大于前向计算任务的执行时长,如反向计算任务的执行时长为前向计算任务的执行时长的2倍。
如图10所示,卡0对应的时序中,由于卡0在基于微批次b7执行第1层的前向计算时(对应无填充方框F7),内存可用容量达到峰值,在基于微批次b5执行第1层的前向计算时(对应无填充方框F5),内存可用容量达到某一阈值,因此,可以将无填充方框F5、无填充方框F6以及无填充方框F7对应的模块确定为第一模块,对该第一模块进行重计算。
902、主机生成多个模块中每个模块的前向计算任务、每个模块的反向计算任务以及第一模块的重计算任务。
举例而言,参考图10,图10中卡0对应时序中,无填充方框R5、无填充方框R6以及无填充方框R7均为第一模块的重计算任务。
903、主机在向M个加速卡中第三加速卡发送第一反向计算任务之前,向第三加速卡发送重计算任务,第一反向计算任务为第一模块的反向计算任务,第三加速卡用于执行第一模块相关的计算任务。
其中,以图10所示内容为例,若第三加速卡为卡0,第一反向计算任务为卡0对应时序中的无填充方框B5、无填充方框B6以及无填充方框B7。应理解,由于在M个加速卡以虚拟流水线并行执行训练请求的过程中,将部分模块确定为待执行重计算的模块,因此,加速卡在执行第一模块的前向计算任务时,能够丢弃该第一模块的中间激活值,如此,可以释放加速卡内存空间。示意性地,主机向第三加速卡发送第一模块的前向计算任务,第一模块的前向计算任务指示对第一模块进行前向计算后,丢弃第一模块的中间激活值。
另外,通过上述方法,在M个加速卡以虚拟流水线并行的方式执行训练请求时,对部分模块进行重计算,还能够将流水线的气泡利用起来进行重计算,从而减少算力的浪费。比如说,继续参考图10,以第三加速卡为卡0为例,对于卡0而言,可以将无填充方框R5(基于微批次b5对第1层的重计算)调度至无填充方框B3之前的气泡处执行,将无填充方框R6(基于微批次b6对第1层的重计算)调度至无填充方框B4之前的气泡处执行,将无填充方框R7(基于微批次b7对第1层的重计算)调度至有填充方框B5之前的气泡处执行,如此,卡0在执行无填充方框F5(基于微批次b5对第1层的前向计算)、无填充方框F6(基于微批次b6对第1层的前向计算)以及无填充方框F7(基于微批次b7对第1层的前向计算)时,可以丢弃第1层的中间激活值,在维持性能不变的情况下释放加速卡内存空间。当然,在实际应用中,能够根据业务需求确定重计算的调度策略,本申请对此不作限定。
在一些实施例中,对于M个加速卡中任意一个加速卡,主机采用如下调度方式,向该加速卡下发任务,以实现如上述图10所示的流水线并行执行效果。示意性地,以流水线数量表示为pp,虚拟流水线数量表示为vpp,微批次数量表示为b,任意一个加速卡表示为第i个加速卡为例(i的取值范围为0~M-1),该调度方式包括如下几个阶段:
阶段1、主机基于流水线数量和虚拟流水线数量,向加速卡发送X个前向计算任务,X个前向计算任务后续不进行重计算。示意性地,X=min(pp×vpp+pp-2-2i,pp×vpp)。参考图11,以第i个加速卡为卡0为例,X个前向计算任务为前8个方框对应的前向计算。
阶段2、主机基于流水线数量,向加速卡发送Y个前向计算任务,Y个前向计算任务后续需要进行重计算。示意性地,Y=max(pp-2-2i,0)。参考图11,以第i个加速卡为卡0为例,Y个前向计算任务为无填充方框F5(基于微批次b5对第1层的前向计算)、无填充方框F6(基于微批次b6对第1层的前向计算)。
阶段3、主机基于流水线数量、微批次数量以及虚拟流水线数量,向加速卡发送Z个第一计算任务组,一个第一计算任务组包括一个前向计算任务和一个反向计算任务。示意性地,Z=b×vpp-pp×vpp-pp+2+2i。其中,若pp-2-2i≥0,则Z个第一计算任务组中首个第一计算任务组的前向计算任务需要执行重计算。参考图11,以第i个加速卡为卡0为例,Z个第一计算任务组如图中箭头所指,一共6组。
阶段4、主机基于流水线数量,向加速卡发送W个第二计算任务组,一个第二计算任务组包括一个重计算任务和一个反向计算任务。示意性地,W=max(pp-1-2i,0)。参考图11,以第i个加速卡为卡0为例,W个第二计算任务组如图中箭头所指,一共3组。
阶段5、主机基于流水线数量和虚拟流水线数量,向加速卡发送K个反向计算任务。示意性地,K=min(pp×vpp-1,pp×vpp+pp-2-2i)。参考图11,以第i个加速卡为卡0为例,K个反向计算任务如图中箭头所指,一共7个。
需要说明的是,上述调度方式仅为举例说明,在一些实施例中,能够根据业务需求设置需要进行重计算的模块,以及,设置重计算任务的执行时机,等等,本申请对此不作限定。
另外,参考图12,本申请还提供了一种AI模型的训练装置,该装置用于实现前述AI模型的训练方法。图12是本申请实施例提供的一种AI模型的训练装置的结构示意图。如图12所示,该装置包括前向计算执行单元1201、丢弃单元1202、反向计算执行单元1203。
前向计算执行单元1201,用于响应于AI模型的训练请求,对AI模型的多个模块执行前向计算,多个模块包括第一模块及第二模块,第一模块为后向计算过程中需要执行重计算的模块,第二模块为前向计算过程中需要使用第一模块的输出数据的模块;
丢弃单元1202,用于在第二模块的前向计算执行完成后,丢弃第一模块的输出数据;
反向计算执行单元1203,用于对AI模型的多个模块执行反向计算,在对第二模块执行反向计算之前,对第一模块执行重计算,并基于对第一模块进行重计算得到的数据,对第二模块执行反向计算。
在一些实施例中,第二模块为多个模块中第一模块的下一个模块。
在一些实施例中,AI模型包括多个网络层,每个加速卡上执行至少一个网络层,至少一个网络层包括顺序执行的多个模块,第二模块为多个模块中的最后一个模块。
在一些实施例中,AI模型为Transformer模型,网络层为Transformer层。
需要说明的是:上述实施例提供的AI模型的训练装置在训练AI模型时,仅以上述各功能模块的划分进行举例说明,实际应用中,可以根据需要而将上述功能分配由不同的功能模块完成,即将装置的内部结构划分成不同的功能模块,以完成以上描述的全部或者部分功能。另外,上述实施例提供的AI模型的训练装置与AI模型的训练方法实施例属于同一构思,其具体实现过程详见方法实施例,这里不再赘述。
参考图13,本申请还提供了一种AI模型的训练装置,该装置用于实现前述AI模型的训练方法。图13是本申请实施例提供的另一种AI模型的训练装置的结构示意图。如图13所示,该装置包括确定单元1301、生成单元1302、发送单元1303。
确定单元1301,用于响应于人工智能AI模型的训练请求,确定AI模型的多个模块,多个模块包括待执行重计算的第一模块;
生成单元1302,用于生成多个模块中每个模块的前向计算任务、每个模块的反向计算任务以及第一模块的重计算任务;
发送单元1303,用于在向训练请求相关的加速卡发送第一反向计算任务之前,向加速卡发送重计算任务,第一反向计算任务与重计算任务之间存在依赖关系。
在一些实施例中,训练请求由N个加速卡以流水线并行的方式执行,N为正整数;
发送单元1303,用于在向N个加速卡中第一加速卡发送第一反向计算任务之前,向第一加速卡发送重计算任务,第一反向计算任务为第一模块的反向计算任务,第一加速卡用于执行第一模块相关的计算任务。
在一些实施例中,发送单元1303,还用于:
向N个加速卡中第二加速卡发送第二反向计算任务,第二加速卡用于执行第三模块相关的计算任务,第三模块为第一模块的下一个模块,第二反向计算任务为第三模块的反向计算任务;
向第一加速卡发送第一反向计算任务,第一反向计算任务指示基于第二反向计算任务的执行结果,对第一模块进行反向计算。
在一些实施例中,训练请求由M个加速卡以虚拟流水线并行的方式执行,M为正整数;
确定单元1301,用于响应于训练请求,基于虚拟流水线的数量、流水线的数量以及M个加速卡的数量,确定第一模块;
发送单元1303,用于在向M个加速卡中第三加速卡发送第一反向计算任务之前,向第三加速卡发送重计算任务,第一反向计算任务为第一模块的反向计算任务,第三加速卡用于执行第一模块相关的计算任务;
发送单元1303,还用于向第三加速卡发送第一模块的前向计算任务,第一模块的前向计算任务指示对第一模块进行前向计算后,丢弃第一模块的中间激活值。
需要说明的是:上述实施例提供的AI模型的训练装置在训练AI模型时,仅以上述各功能模块的划分进行举例说明,实际应用中,可以根据需要而将上述功能分配由不同的功能模块完成,即将装置的内部结构划分成不同的功能模块,以完成以上描述的全部或者部分功能。另外,上述实施例提供的AI模型的训练装置与AI模型的训练方法实施例属于同一构思,其具体实现过程详见方法实施例,这里不再赘述。
本申请中术语“第一”“第二”等字样用于对作用和功能基本相同的相同项或相似项进行区分,应理解,“第一”、“第二”、“第n”之间不具有逻辑或时序上的依赖关系,也不对数量和执行顺序进行限定。还应理解,尽管以下描述使用术语第一、第二等来描述各种元素,但这些元素不应受术语的限制。这些术语只是用于将一元素与另一元素区别分开。例如,在不脱离各种所述示例的范围的情况下,第一加速卡可以被称为第二加速卡,并且类似地,第二加速卡可以被称为第一加速卡。第一加速卡和第二加速卡都可以是加速卡,并且在某些情况下,可以是单独且不同的加速卡。
本申请中术语“至少一个”的含义是指一个或多个,本申请中术语“多个”的含义是指两个或两个以上,例如,多个加速卡是指两个或两个以上的加速卡。
以上描述,仅为本申请的具体实施方式,但本申请的保护范围并不局限于此,任何熟悉本技术领域的技术人员在本申请揭露的技术范围内,可轻易想到各种等效的修改或替换,这些修改或替换都应涵盖在本申请的保护范围之内。因此,本申请的保护范围应以权利要求的保护范围为准。
在上述实施例中,可以全部或部分地通过软件、硬件、固件或者其任意组合来实现。当使用软件实现时,可以全部或部分地以程序结构信息的形式实现。该程序结构信息包括一个或多个程序指令。在计算设备上加载和执行该程序指令时,全部或部分地产生按照本申请实施例中的流程或功能。
本领域普通技术人员可以理解实现上述实施例的全部或部分步骤可以通过硬件来完成,也可以通过程序来指令相关的硬件完成,该程序可以存储于一种计算机可读存储介质中,上述提到的存储介质可以是只读存储器,磁盘或光盘等。
以上所述,以上实施例仅用以说明本申请的技术方案,而非对其限制;尽管参照前述实施例对本申请进行了详细的说明,本领域的普通技术人员应当理解:其依然可以对前述各实施例所记载的技术方案进行修改,或者对其中部分技术特征进行等同替换;而这些修改或者替换,并不使相应技术方案的本质脱离本申请各实施例技术方案的范围。

Claims (16)

  1. 一种人工智能AI模型的训练方法,其特征在于,所述方法包括:
    响应于AI模型的训练请求,对所述AI模型的多个模块执行前向计算,所述多个模块包括第一模块及第二模块,所述第一模块为后向计算过程中需要执行重计算的模块,所述第二模块为前向计算过程中需要使用所述第一模块的输出数据的模块;
    在所述第二模块的前向计算执行完成后,丢弃所述第一模块的输出数据;
    对所述AI模型的多个模块执行反向计算,在对所述第二模块执行反向计算之前,对所述第一模块执行重计算,并基于对所述第一模块进行重计算得到的数据,对所述第二模块执行反向计算。
  2. 根据权利要求1所述的方法,其特征在于,所述第二模块为所述多个模块中所述第一模块的下一个模块。
  3. 根据权利要求1所述的方法,其特征在于,所述AI模型包括多个网络层,每个加速卡上执行至少一个网络层,所述至少一个网络层包括顺序执行的所述多个模块,所述第二模块为所述多个模块中的最后一个模块。
  4. 根据权利要求3所述的方法,其特征在于,所述AI模型为Transformer模型,所述网络层为Transformer层。
  5. 一种人工智能AI模型的训练方法,其特征在于,所述方法包括:
    响应于人工智能AI模型的训练请求,确定所述AI模型的多个模块,所述多个模块包括待执行重计算的第一模块;
    生成所述多个模块中每个模块的前向计算任务、每个模块的反向计算任务以及所述第一模块的重计算任务;
    在向所述训练请求相关的加速卡发送第一反向计算任务之前,向所述加速卡发送所述重计算任务,所述第一反向计算任务与所述重计算任务之间存在依赖关系。
  6. 根据权利要求5所述的方法,其特征在于,所述训练请求由N个加速卡以流水线并行的方式执行,N为正整数;
    所述在向所述训练请求相关的加速卡发送第一反向计算任务之前,向所述加速卡发送所述重计算任务,包括:在向所述N个加速卡中第一加速卡发送所述第一反向计算任务之前,向所述第一加速卡发送所述重计算任务,所述第一反向计算任务为所述第一模块的反向计算任务,所述第一加速卡用于执行所述第一模块相关的计算任务。
  7. 根据权利要求6所述的方法,其特征在于,所述方法还包括:
    向所述N个加速卡中第二加速卡发送第二反向计算任务,所述第二加速卡用于执行第三模块相关的计算任务,所述第三模块为所述第一模块的下一个模块,所述第二反向计算任务为所述第三模块的反向计算任务;
    向所述第一加速卡发送所述第一反向计算任务,所述第一反向计算任务指示基于所述第二反向计算任务的执行结果,对所述第一模块进行反向计算。
  8. 根据权利要求5所述的方法,其特征在于,所述训练请求由M个加速卡以虚拟流水线并行的方式执行,M为正整数;
    所述响应于人工智能AI模型的训练请求,确定所述AI模型的多个模块,包括:响应于所述训练请求,基于虚拟流水线的数量、流水线的数量以及所述M个加速卡的数量,确定所述第一模块;
    所述在向所述训练请求相关的加速卡发送第一反向计算任务之前,向所述加速卡发送所述重计算任务,包括:在向所述M个加速卡中第三加速卡发送所述第一反向计算任务之前,向所述第三加速卡发送所述重计算任务,所述第一反向计算任务为所述第一模块的反向计算任务,所述第三加速卡用于执行所述第一模块相关的计算任务;
    所述方法还包括:向所述第三加速卡发送所述第一模块的前向计算任务,所述第一模块的前向计算任务指示对所述第一模块进行前向计算后,丢弃所述第一模块的中间激活值。
  9. 一种人工智能AI模型的训练装置,其特征在于,所述装置包括:
    前向计算执行单元,用于响应于AI模型的训练请求,对所述AI模型的多个模块执行前向计算,所述多个模块包括第一模块及第二模块,所述第一模块为后向计算过程中需要执行重计算的模块,所述第二模块为前向计算过程中需要使用所述第一模块的输出数据的模块;
    丢弃单元,用于在所述第二模块的前向计算执行完成后,丢弃所述第一模块的输出数据;
    反向计算执行单元,用于对所述AI模型的多个模块执行反向计算,在对所述第二模块执行反向计算之前,对所述第一模块执行重计算,并基于对所述第一模块进行重计算得到的数据,对所述第二模块执行反向计算。
  10. 根据权利要求9所述的装置,其特征在于,所述第二模块为所述多个模块中所述第一模块的下一个模块。
  11. 根据权利要求9所述的装置,其特征在于,所述AI模型包括多个网络层,每个加速卡上执行至少一个网络层,所述至少一个网络层包括顺序执行的所述多个模块,所述第二模块为所述多个模块中的最后一个模块。
  12. 根据权利要求11所述的装置,其特征在于,所述AI模型为Transformer模型,所述网络层为Transformer层。
  13. 一种计算设备,其特征在于,所述计算设备包括处理器和存储器,所述处理器用于执行所述存储器中存储的至少一段程序代码,以使所述计算设备实现如前述权利要求1至8中任一项所述的AI模型的训练方法。
  14. 一种计算系统,其特征在于,所述计算系统包括主机和至少一个加速卡,所述系统用于实现如前述权利要求1至8中任一项所述的AI模型的训练方法。
  15. 一种计算机程序产品,其特征在于,当所述计算机程序产品在计算设备上运行时,使得所述计算设备实现如前述权利要求1至8中任一项所述的AI模型的训练方法。
  16. 一种计算机可读存储介质,其特征在于,所述计算机可读存储介质用于存储至少一段程序代码,所述至少一段程序代码用于实现如前述权利要求1至8中任一项所述的AI模型的训练方法。
PCT/CN2025/079244 2024-06-29 2025-02-26 人工智能ai模型的训练方法、装置及计算设备 Pending WO2026001038A1 (zh)

Applications Claiming Priority (2)

Application Number Priority Date Filing Date Title
CN202410869302.6 2024-06-29
CN202410869302.6A CN121235014A (zh) 2024-06-29 2024-06-29 人工智能ai模型的训练方法、装置及计算设备

Publications (1)

Publication Number Publication Date
WO2026001038A1 true WO2026001038A1 (zh) 2026-01-02

Family

ID=98151945

Family Applications (1)

Application Number Title Priority Date Filing Date
PCT/CN2025/079244 Pending WO2026001038A1 (zh) 2024-06-29 2025-02-26 人工智能ai模型的训练方法、装置及计算设备

Country Status (2)

Country Link
CN (1) CN121235014A (zh)
WO (1) WO2026001038A1 (zh)

Citations (5)

* Cited by examiner, † Cited by third party
Publication number Priority date Publication date Assignee Title
US20180321980A1 (en) * 2015-12-04 2018-11-08 Cornell University Execution time prediction for energy-efficient computer systems
CN115688893A (zh) * 2022-10-19 2023-02-03 北京百度网讯科技有限公司 内存调度方法及装置、电子设备和存储介质
CN115965074A (zh) * 2022-11-28 2023-04-14 北京百度网讯科技有限公司 深度学习模型的训练方法、数据处理方法、装置和设备
CN116107754A (zh) * 2023-02-24 2023-05-12 华中科技大学 一种面向深度神经网络的内存管理方法及系统
CN118195033A (zh) * 2024-05-20 2024-06-14 科大讯飞股份有限公司 一种语言模型训练方法及相关装置

Patent Citations (5)

* Cited by examiner, † Cited by third party
Publication number Priority date Publication date Assignee Title
US20180321980A1 (en) * 2015-12-04 2018-11-08 Cornell University Execution time prediction for energy-efficient computer systems
CN115688893A (zh) * 2022-10-19 2023-02-03 北京百度网讯科技有限公司 内存调度方法及装置、电子设备和存储介质
CN115965074A (zh) * 2022-11-28 2023-04-14 北京百度网讯科技有限公司 深度学习模型的训练方法、数据处理方法、装置和设备
CN116107754A (zh) * 2023-02-24 2023-05-12 华中科技大学 一种面向深度神经网络的内存管理方法及系统
CN118195033A (zh) * 2024-05-20 2024-06-14 科大讯飞股份有限公司 一种语言模型训练方法及相关装置

Also Published As

Publication number Publication date
CN121235014A (zh) 2025-12-30

Similar Documents

Publication Publication Date Title
CN113435682B (zh) 分布式训练的梯度压缩
CN110998570B (zh) 具有带有块浮点处理的矩阵矢量单元的硬件节点
US11176449B1 (en) Neural network accelerator hardware-specific division of inference into groups of layers
US9053067B2 (en) Distributed data scalable adaptive map-reduce framework
CN114008586A (zh) 使用处理元件阵列的转置运算
US9159028B2 (en) Computing regression models
US10628740B2 (en) Asynchronous stochastic gradient descent
US11354573B2 (en) Dynamically resizing minibatch in neural network execution
CN117112145B (zh) 训练模型分配方法、装置、计算机设备和存储介质
US20240086719A1 (en) Sparse encoding and decoding at mixture-of-experts layer
CN111047025B (zh) 一种卷积计算方法及装置
US11734214B2 (en) Semi-programmable and reconfigurable co-accelerator for a deep neural network with normalization or non-linearity
CN118297113A (zh) 自动显存调优的模型并行训练方法、装置及设备
WO2026001038A1 (zh) 人工智能ai模型的训练方法、装置及计算设备
WO2025085808A1 (en) Method and system for efficient encryption of fine tuned large language models
US20230342320A1 (en) Semi-programmable and reconfigurable co-accelerator for a deep neural network with normalization or non-linearity
CN114357371A (zh) 一种矩阵数据的处理方法、装置、电子设备及存储介质
EP4202774A1 (en) Runtime predictors for neural network computation reduction
US11520781B2 (en) Efficient bulk loading multiple rows or partitions for a single target table
CN119578473A (zh) 数据处理方法、装置、加速器及计算设备
CN120234516B (zh) 用于包括多个计算单元的处理器执行矩阵乘法运算的方法
CN119624748B (zh) 计算图的处理方法、装置及计算机可读存储介质
US12587360B1 (en) Multi-chip accelerator architecture for fully homomorphic encryption inference
CN119005271B (zh) 基于算子分块的神经网络模型并行优化方法及装置
WO2026007433A1 (zh) Ai模型分布式计算过程中的任务执行方法及装置

Legal Events

Date Code Title Description
121 Ep: the epo has been informed by wipo that ep was designated in this application

Ref document number: 25824433

Country of ref document: EP

Kind code of ref document: A1