Executing large artificial intelligence models on memory-constrained devices
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 manually or automatically adjusted to reduce the communication overhead.
1. A system, comprising:
a group of target devices, the group of target devices comprising one or more target devices, each target device of the group of target devices being communicatively connected to a parameter server that stores a master copy of an artificial intelligence (AI) model, the group of target devices being configured to run an instance of the AI model, at least one target device comprising an integrated circuit chip having an on-chip memory of a size less than an entirety of the AI model, each target device comprises:
a downloader configured to download a portion of the AI model from the parameter server, 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;
a data manager configured to store a set of microbatches of a dataset; and
an executer configured to execute by the integrated circuit chip the set of microbatches on a first subportion of the downloaded portion of the AI model stored in the on-chip memory to generate activations; and
the downloader being further configured to download weights for a second subportion of the downloaded portion of the AI from the parameter server.
2. The system of claim 1 , wherein
the executer is further configured to execute the set of microbatches of the dataset on the second subportion using the downloaded weights for the second subportion;
the downloader is further configured to download weights for a third subportion of the downloaded portion of the AI model into the memory of the target device from the parameter server;
wherein the executing the set of microbatches of the dataset on the second subportion and the downloading weights for the third subportion are performed contemporaneously.
3. The system of claim 1 , wherein
the executer is further configured to execute the set of microbatches on the second subportion using the downloaded weights for the second subportion; and
the downloader is further configured to download weights for a third subportion of the downloaded portion of the AI model from the parameter server;
wherein the executing the set of microbatches of the dataset on the second subportion and the downloading weights for the third subportion are performed serially.
4. The system of claim 1 , wherein the set of microbatches comprises a plurality of microbatches that are configured to be executed in sequential order, the set of microbatches forming a minibatch that comprises a number of samples per update for training or a number of samples served in every inference cycle for inference.
5. The system of claim 4 , wherein each of the target device further comprises an output manager configured to:
send the activations for the first subportion to the parameter server or save the activations on the target device for a forward pass during training of the AI model; and
restore the activations for a backward pass during the training of the AI model.
6. The system of claim 4 , wherein each microbatch of the set of microbatches has a microbatch size selected based on a rate of execution of the plurality of microbatches and a rate of communication between each target device and the parameter server.
7. A system, comprising:
a group of target devices, the group of target devices comprising one or more target devices, each target device of the group of target devices being communicatively connected to a parameter server that stores a master copy of an artificial intelligence (AI) model, the group of target devices being configured to run an instance of the AI model, each target device comprises:
a downloader configured to download a portion of the AI model from the parameter server;
a data manager configured to store a set of microbatches of a dataset; and
an executer configured to execute the set of microbatches on a first subportion of the downloaded portion of the AI model to generate activations; and
the downloader being further configured to download weights for a second subportion of the downloaded portion of the AI from the parameter server, wherein
the executer is further configured to execute the set of microbatches on a third subportion of the downloaded portion of the AI model to generate gradients;
the downloader is further configured to download weights and output activations for a fourth subportion of the downloaded portion of the AI model; and
contemporaneously
the executer is further configured to execute the set of microbatches on the fourth subportion using the downloaded weights and output activations,
the downloader is configured to download weights and output activations for a fifth subportion of the downloaded portion of the AI model from the parameter server, and
the output manager is further configured to send the gradients for the third subportion to the parameter server.
8. The system of claim 1 , wherein each target device comprises at least one of an application-specific integrated circuit, a graphics processing unit or an edge device.
9. A method implemented in a target device, comprising:
downloading a portion of an artificial intelligence (AI) model from a parameter server, 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;
storing a set of microbatches of a dataset;
executing by the integrated circuit chip at the target device the set of microbatches on a first subportion of the downloaded portion of the AI model stored in the on-chip memory to generate activations; and
downloading weights for a second subportion of the downloaded portion of the AI model from the parameter server.
10. The method of claim 9 , further comprising:
contemporaneously
executing the set of microbatches of the dataset on the second subportion using the downloaded weights for the second subportion; and
downloading weights for a third subportion of the downloaded portion of the AI model from the parameter server.
11. The method of claim 9 , further comprising:
serially
executing the set of microbatches on the second subportion using the downloaded weights for the second subportion; and
downloading weights for a third subportion of the downloaded portion of the AI model from the parameter server.
12. The method of claim 9 , wherein the set of microbatches comprises a plurality of microbatches that are configured to be executed in sequential order, the set of microbatches forming a minibatch that comprises a number of samples per update for training or a number of samples served in every inference cycle for inference.
13. The method of claim 12 , further comprises:
sending the activations for the first subportion to the parameter server or save the activations on the target device for a forward pass during training of the AI model; and
restoring the activations for a backward pass during the training of the AI model.
14. The method of claim 12 , wherein each microbatch of the set of microbatches has a microbatch size that is configurable based on a rate of executing the plurality of microbatches and a rate of communication between the target device and the parameter server.
15. A method implemented in a target device, comprising:
downloading a portion of an artificial intelligence (AI) model from a parameter server;
storing a set of microbatches of a dataset;
executing the set of microbatches on a first subportion of the downloaded portion of the AI model to generate activations and a third subportion of the downloaded portion of the AI model to generate gradients;
downloading weights for a second subportion of the downloaded portion of the AI model from the parameter server
downloading weights and output activations for a fourth subportion of the downloaded portion of the AI model; and
contemporaneously
executing the set of microbatches on the fourth subportion using the downloaded weights and output activations,
downloading weights and output activations for a fifth subportion of the downloaded portion of the AI model from the parameter server, and
sending the gradients for the third subportion to the parameter server.
16. A system, comprising:
a parameter server communicatively connected to a group of target devices, the group of target devices comprising one or more target devices, the group of target devices being configured to run an instance of an artificial intelligence (AI) model, at least one target device comprising an integrated circuit chip having an on-chip memory of a size less than an entirety of the AI model, the parameter server comprising:
a data manager configured to store a master copy the AI model;
a batch manager configured to determine a microbatch size suitable for each target device of the group of target devices, and
a transmitter configured to transmit a portion of the AI model to the group of target devices, 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 training dataset being executed by the integrated circuit chip at the at least one target device on a first subportion of the transmitted portion of the AI model stored in the on-chip memory to generate activations,
a weight updater is configured to perform 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.
17. The system of claim 16 , wherein the microbatch size is based on a rate of executing a set of microbatches at each target device and a rate of communication between the target device and the parameter server.
18. The system of claim 16 , wherein parameter server further comprises an output data manager configured to:
receive activations from each target device after each minibatch is executed; and
generate output activations for a subportion of the downloaded portion of the AI model based on the received activations.
19. The system of claim 16 , wherein the parameter server further comprises a weight updater configured to:
update weights of the AI model based on gradients received from each target device.
20. The system of claim 16 , wherein the parameter server comprises a central processing unit, a field programmable gate array, or an application-specific integrated circuit.