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 method performed by 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, a first target device of the group of target devices comprising an integrated circuit chip having an on-chip memory of a size less than an entirety of the AI model, the first target device comprising a downloader, a data manager, and an executer, the method comprising, by the first target device:
downloading, by the downloader, 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;
storing, by the data manager, a set of microbatches of a dataset;
executing, by the executer and via 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
downloading, by the downloader, weights for a second subportion of the downloaded portion of the AI from the parameter server.
2 . The method of claim 1 , further comprising, by the first target device:
executing, by the executer, the set of microbatches of the dataset on the second subportion using the downloaded weights for the second subportion; and
downloading, by the downloader, 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, by the executer, the set of microbatches of the dataset on the second subportion and the downloading, by the downloader, the weights for the third subportion are performed contemporaneously.
3 . The method 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;
or
wherein the first target device comprises at least one of an application-specific integrated circuit, a graphics processing unit, or an edge device.
4 . The method of claim 3 , 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 the first target device and the parameter server.
5 . The method of claim 1 , wherein the first target device further comprises an output manager, the method further comprising, by the first target device
saving restoration data suitable for recomputing a first activation of the activations;
discarding the first activation; and
restoring, by the output manager and utilizing the restoration data, the first activation for a backward pass during the training of the AI model.
6 . The method of claim 1 , the first target device further comprises an output manager, the method further comprising, by the first target device:
executing, by the executer, the set of microbatches on a third subportion of the downloaded portion of the AI model to generate gradients;
downloading, by the downloader, weights and output activations for a fourth subportion of the downloaded portion of the AI model; and
contemporaneously by the first target device:
executing, by the executer, the set of microbatches on the fourth subportion using the downloaded weights and output activations,
downloading, by the downloader, weights and output activations for a fifth subportion of the downloaded portion of the AI model from the parameter server, and
sending, by the output manager, the gradients for the third subportion to the parameter server.
7 . At least one computer-readable storage medium having program instructions recorded thereon that, when executed by at least one processing circuit of a group of target devices, configure a first target device of the group of target devices to perform a method, 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, the first target device comprising an integrated circuit chip having an on-chip memory, a downloader, a data manager, and an executer, the method comprising:
downloading, by the downloader, a portion of the AI model from the parameter server, a size of the portion being based at least on a size of the on-chip memory and a size of one or more layers of the AI model;
storing, by the data manager, a set of microbatches of a dataset;
executing, by the executer and via 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
downloading, by the downloader, weights for a second subportion of the downloaded portion of the AI from the parameter server.
8 . A target device comprising an integrated circuit chip having an on-chip memory of a size less than an entirety of an artificial intelligence (AI) model, the target device configured to:
download a portion of the AI model from a 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;
store a set of microbatches of a dataset;
execute 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
download weights for a second subportion of the downloaded portion of the AI model from the parameter server.
9 . The target device of claim 8 , being further configured to:
serially:
execute the set of microbatches on the second subportion using the downloaded weights for the second subportion, and
download weights for a third subportion of the downloaded portion of the AI model from the parameter server.
10 . The target device of claim 8 , 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.
11 . The target device of claim 10 , 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.
12 . The target device of claim 8 , wherein the target device is further configured to:
send a first activation of the activations for the first subportion to the parameter server; and
receive the first activation from the parameter server for a backward pass during the training of the AI model.
13 . The target device of claim 8 , being further configured to:
download weights and output activations for a fourth subportion of the downloaded portion of the AI model; and
contemporaneously
execute the set of microbatches on the fourth subportion using the downloaded weights and output activations,
download weights and output activations for a fifth subportion of the downloaded portion of the AI model from the parameter server, and
send the gradients for the third subportion to the parameter server.
14 . A computer-readable storage medium having program instructions recorded thereon that, when executed by at least one processing circuit of a target device, configure the target device to perform a method that comprises:
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, and a size of the portion being based at least on a size of the on-chip memory 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.
15 . The computer-readable storage medium of claim 14 , wherein the method further comprises:
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 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 the weights for the third subportion are performed contemporaneously;
or
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;
wherein the executing the set of microbatches of the dataset on the second subportion and the downloading the weights for the third subportion are performed serially.
16 . The computer-readable storage medium of claim 14 , 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.
17 . The computer-readable storage medium of claim 16 , 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 the target device and the parameter server.
18 . The computer-readable storage medium of claim 14 , the method further comprising:
saving the activations on the target device for a forward pass during training of the AI model.
19 . The computer-readable storage medium of claim 14 , wherein the method further comprises:
executing the set of microbatches on a third subportion of the downloaded portion of the AI model to generate gradients;
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.
20 . The computer-readable storage medium of claim 14 , wherein the target device comprises at least one of:
an application-specific integrated circuit,
a graphics processing unit, or
an edge device.