Data parallelism in distributed training of artificial intelligence models

Abstract

Methods, systems, apparatuses, and computer program products are described herein that enable execution of a large AI model on a memory-constrained target device that is communicatively connected to a parameter server, which stores a master copy of the AI model. The AI model may be dissected into smaller portions (e.g., layers or sub-layers), and each portion may be executed as efficiently as possible on the target device. After execution of one portion of the AI model is finished, another portion of the AI model may be downloaded and executed at the target device. To improve efficiency, the input samples may be divided into microbatches, and a plurality of microbatches executing in sequential order may form a minibatch. The size of the group of microbatches or minibatch can be adjusted to reduce the communication overhead. Multi-level parallel parameters reduction may be performed at the parameter server and the target device.

What is claimed is: 
     
         1 . A system, comprising:
 a parameter server communicatively connected to a target device, the parameter server comprises:   a transmitter configured to transmit a portion of an artificial intelligence (AI) model to the target device, the target device comprising an integrated circuit chip having an on-chip memory of a size less than an entirety of the AI model, and a size of the portion being based at least on the on-chip memory size and a size of one or more layers of the AI model; and   contemporaneously, with a set of microbatches of a training dataset being executed by the integrated circuit chip at the target device on a first subportion of the transmitted portion of the AI model stored in the on-chip memory to generate gradients,
 a weight updater is configured to perform a reduction of parameters for a second subportion of the transmitted portion of the AI model, and 
 the transmitter is further configured to send weights for a third subportion of the transmitted portion of the AI model to the target device. 
   
     
     
         2 . The system of  claim 1 , wherein wherein the first subportion of the transmitted portion of the AI model comprises a current layer of the AI model and the second subportion of the transmitted portion of the AI model comprises a prior layer of the AI model. 
     
     
         3 . The system of  claim 1 , wherein the first subportion of the transmitted portion of the AI model comprises a current layer of the AI model and the third subportion of the transmitted portion of the AI model comprises a next layer of the AI model. 
     
     
         4 . The system of  claim 1 , wherein the weight updater is configured to:
 perform reduction of parameters by:
 receiving gradients from the target device, the gradients being generated by the target device executing the set of microbatches of the training dataset on the second subportion at the target device; and 
 generating an average of the received gradients; and 
   update the AI model with the average of the received gradients.   
     
     
         5 . The system of  claim 1 , wherein a microbatch size is configurable based on a rate of executing the set of microbatches at the target device and a rate of communication between the target device and the parameter server. 
     
     
         6 . The system of  claim 1 , wherein the parameter server further comprises a precision formatter configured to:
 convert weights for a fourth subportion of the transmitted portion of the AI model to a first precision format prior to sending the weights to the target device;   convert gradients received from the target device to a second precision format; and   update the weights using the converted gradients.   
     
     
         7 . The system of  claim 1 , wherein the transmitter is further configured to transmit another portion of the AI model to another target device; and the weight updater is further configured to receive gradients from the another target device to perform reduction of parameters for the another portion of the AI model. 
     
     
         8 . A method implemented in a parameter server, comprising:
 transmitting a portion of a stored artificial intelligence (AI) model to a target device, the target device comprising an integrated circuit chip having an on-chip memory of a size less than an entirety of the AI model, and a size of the portion being based at least on the on-chip memory size and a size of one or more layers of the AI model; and   contemporaneously, with a set of microbatches, of a determined size, of a training dataset being executed by the integrated circuit chip at the target device on a first subportion of the transmitted portion of the AI model stored in the on-chip memory to generate gradients, performing a reduction of parameters for a second subportion of the transmitted portion of the AI model and sending weights for a third subportion of the transmitted portion of the AI model to the target device.   
     
     
         9 . The method of  claim 8 , wherein wherein the first subportion of the transmitted portion of the AI model comprises a current layer of the AI model and the second subportion of the transmitted portion of the AI model comprises a prior layer of the AI model. 
     
     
         10 . The method of  claim 8 , wherein the first subportion of the transmitted portion of the AI model comprises a current layer of the AI model and the third subportion of the transmitted portion of the AI model comprises a next layer of the AI model. 
     
     
         11 . The method of  claim 8 , wherein the performing reduction of parameters comprises:
 receiving gradients from the target device, the gradients being generated from executing the set of microbatches of the training dataset on the second subportion at the target device; and   generating an average of the received gradients;   
       wherein the method further comprises:
 updating the AI model with the average of the received gradients. 
 
     
     
         12 . The method of  claim 8 , wherein a microbatch size is configurable based on a rate of executing the set of microbatches at the target device and a rate of communication between the target device and the parameter server. 
     
     
         13 . The method of  claim 8 , further comprising:
 converting weights for a fourth subportion of the transmitted portion of the AI model to a first precision format prior to sending the weights to the target device;   converting gradients received from the target device to a second precision format; and   updating the weights using the converted gradients.   
     
     
         14 . The method of  claim 8 , further comprising:
 transmitting another portion of the AI model to another target device; and   receiving gradients from the another target device to perform reduction of parameters for the another portion of the AI model.   
     
     
         15 . A computer program product comprising a computer-readable storage device having computer program logic recorded thereon that when executed by a processor- based computer system causes the processor-based system to perform a method, the method comprising:
 transmitting a portion of a stored artificial intelligence (AI) model to a target device, the target device comprising an integrated circuit chip having an on-chip memory of a size less than an entirety of the AI model, and a size of the portion being based at least on the on-chip memory size and a size of one or more layers of the AI model;   contemporaneously, with a set of microbatches, of a determined size, of a training dataset being executed by the integrated circuit chip at the target device on a first subportion of the transmitted portion of the AI model stored in the on-chip memory to generate gradients, performing at least one of a reduction of parameters for a second subportion of the transmitted portion of the AI model or sending weights for a third subportion of the transmitted portion of the AI model to the target device.   
     
     
         16 . The computer program product of  claim 15 , wherein the first subportion of the transmitted portion of the AI model comprises a current layer of the AI model and the second subportion of the transmitted portion of the AI model comprises a prior layer of the AI model. 
     
     
         17 . The computer program product of  claim 15 , wherein the first subportion of the transmitted portion of the AI model comprises a current layer of the AI model and the third subportion of the transmitted portion of the AI model comprises a next layer of the AI model. 
     
     
         18 . The computer program product of  claim 15 , wherein the method further comprises:
 converting weights for a fourth subportion of the transmitted portion of the AI model to a first precision format prior to sending the weights to the target device;   converting gradients received from the target device to a second precision format; and   updating the weights using the converted gradients.   
     
     
         19 . The computer program product of  claim 15 , wherein a microbatch size is configurable based on a rate of executing the set of microbatches at the target device and a rate of communication between the target device and the parameter server. 
     
     
         20 . The computer program product of  claim 15 , wherein the method further comprises:
 transmitting another portion of the AI model to another target device; and   receiving gradients from the another target device to perform reduction of parameters for the another portion of the AI model.