Layered gradient accumulation and modular pipeline parallelism for improved training of machine learning models
A method is provided including: (i) assigning sequentially-ordered layers of a machine learning model to a plurality of compute nodes, each of the layers being assigned to exactly one of the nodes; (ii) dividing training data into micro-batches; (iii) forward-propagating the micro-batches through the model, each node operating in parallel to generate respective activation states for the micro-batches with their assigned layers, and with the activation states being communicated between the nodes according to the layers' sequential ordering; and (iv) backward-propagating the micro-batches through the model, each node operating in parallel to generate respective error states for the micro-batches with their assigned layers, with the error states being communicated between the nodes according to the layers' reverse sequential ordering, wherein each of the nodes completes the backward-propagation of all micro-batches through a given layer prior to performing backward-propagation through any layer that precedes the given layer in the sequential ordering.
1 . A computer-implemented method comprising:
operating M compute nodes to train a machine learning model based on a batch of N training examples, wherein the batch of N training examples is divided into n micro-batches, wherein the machine learning model comprises L layers, each layer defined by a respective plurality of parameters of the machine learning model, wherein operating the M compute nodes to train the machine learning model based on the batch of N training examples comprises updating the parameters of the machine learning model by:
sequentially applying, by a first compute node of the M compute nodes, each of the micro-batches to a first layer of the L layers to generate respective first-layer activation states;
transmitting, from the first compute node to a second compute node of the M compute nodes, the first-layer activation states;
sequentially applying, by the second compute node, each of the first-layer activation states to a second layer of the L layers to generate respective second-layer activation states for each of the micro-batches, wherein the second compute node applying a particular one of the first layer activation states to the second layer at least partially overlaps in time with the first compute node transmitting a subsequent one of the first layer activation states to the second compute node;
transmitting, from an M th compute node of the M compute nodes to the first compute node, M th -layer activation states for each of the micro-batches;
sequentially applying, by the first compute node subsequent to generating the first-layer activation states, each of the M th -layer activation states to an (M+1) th layer of the L layers to generate respective (M+1) th -layer activation states for each of the micro-batches;
transmitting, from an (M−1) th compute node of the M compute nodes to the M th compute node, (L−1) th -layer activation states for each of the micro-batches;
sequentially applying, by the M th compute node, each of the (L−1) th -layer activation states to an L th layer of the L layers to generate respective L th -layer activation states for each of the micro-batches;
based on the L th -layer activation states, generating respective L th -layer error states for each of the micro-batches;
sequentially applying, by the M th compute node, each of the L th -layer error states to the L th layer to generate respective (L−1) th -layer error states for each of the micro-batches and respective L th -layer parameter update information for each of the micro-batches;
transmitting, from the M th compute node to the (M−1) th compute node, the (L−1) th -layer error states;
sequentially applying, by the (M−1) th compute node, each of the (L−1) th -layer error states to the (L−1) th -layer to generate respective (L−2) th -layer error states for each of the micro-batches and respective (L−1) th -layer parameter update information for each of the micro-batches, wherein the (M−1) th compute node applying a particular one of the (L−1) th -layer error states to the (L−1) th -layer at least partially overlaps in time with the M th compute node transmitting a subsequent one of the (L−1) th -layer error states to the (M−1) th compute node;
transmitting, from the first compute node to the M th compute node, (L−M) th -layer error states for each of the micro-batches; and
sequentially applying, by the M th compute node subsequent to generating the (L−1) th -layer error states, each of the (L−M) th -layer error states to an (L−M) th layer of the L layers to generate respective (L−M−1) th -layer activation states for each of the micro-batches and respective (L−M) th -layer parameter update information for each of the micro-batches.
2 . The computer-implemented method of claim 1 , wherein the machine learning model comprises a transformer.
3 . The computer-implemented method of claim 2 , wherein each layer of the L layers includes a discrete number of layers of the transformer.
4 . The computer-implemented method of claim 1 , further comprising:
transmitting, to the first compute node, a portion of the plurality of parameters of the machine learning model that define the (M+1) th -layer prior to the first compute node applying each of the M th -layer activation states to the (M+1) th layer to generate respective (M+1) th -layer activation states for each of the micro-batches,
wherein the first compute node applying at least one of the micro-batches to the first layer of the L layers to generate at least one respective first-layer activation states at least partially overlaps in time with transmitting the portion of the plurality of parameters of the machine learning model that define the (M+1) th -layer to the first compute node.
5 . The computer-implemented method of claim 1 , wherein the number of micro-batches n equals the number of compute nodes M.
6 . The computer-implemented method of claim 1 , wherein each micro-batch includes a single training example from the batch of N training examples.
7 . The computer-implemented method of claim 1 , further comprising:
based on the L th -layer parameter update information for each of the micro-batches, determining parameter updates for the plurality of parameters of the machine learning model that define the L th -layer.
8 . The computer-implemented method of claim 7 , wherein determining parameter updates for the plurality of parameters of the machine learning model that define the L th -layer commences prior to the (M−1) th compute node generating at least one of the (L−2) th -layer error states.
9 . The computer-implemented method of claim 7 , wherein determining parameter updates for the plurality of parameters of the machine learning model that define the L th -layer comprises:
transmitting, from the M th compute node to an additional compute node that is not one of the M compute nodes, the L th -layer parameter update information for each of the micro-batches; and
transmitting, from the additional compute node to the M th compute node, L th -layer parameter update information for one or more additional micro-batches that are not part of the n micro-batches.
10 . The computer-implemented method of claim 7 , further comprising:
updating the plurality of parameters of the machine learning model that define the L th -layer based on the parameter updates.
11 . An article of manufacture including a non-transitory computer-readable medium, having stored thereon program instructions that, upon execution by a computing device, cause the computing device to perform operations comprising:
operating M compute nodes to train a machine learning model based on a batch of N training examples, wherein the batch of N training examples is divided into n micro-batches, wherein the machine learning model comprises L layers, each layer defined by a respective plurality of parameters of the machine learning model, wherein operating the M compute nodes to train the machine learning model based on the batch of N training examples comprises updating the parameters of the machine learning model by:
sequentially applying, by a first compute node of the M compute nodes, each of the micro-batches to a first layer of the L layers to generate respective first-layer activation states;
transmitting, from the first compute node to a second compute node of the M compute nodes, the first-layer activation states;
sequentially applying, by the second compute node, each of the first-layer activation states to a second layer of the L layers to generate respective second-layer activation states for each of the micro-batches, wherein the second compute node applying a particular one of the first layer activation states to the second layer at least partially overlaps in time with the first compute node transmitting a subsequent one of the first layer activation states to the second compute node;
transmitting, from an M th compute node of the M compute nodes to the first compute node, M th -layer activation states for each of the micro-batches;
sequentially applying, by the first compute node subsequent to generating the first-layer activation states, each of the M th -layer activation states to an (M+1) th layer of the L layers to generate respective (M+1) th -layer activation states for each of the micro-batches;
transmitting, from an (M−1) th compute node of the M compute nodes to the M th compute node, (L−1) th -layer activation states for each of the micro-batches;
sequentially applying, by the M th compute node, each of the (L−1) th -layer activation states to an L th layer of the L layers to generate respective L th -layer activation states for each of the micro-batches;
based on the L th -layer activation states, generating respective L th -layer error states for each of the micro-batches;
sequentially applying, by the M th compute node, each of the L th -layer error states to the L th layer to generate respective (L−1) th -layer error states for each of the micro-batches and respective L th -layer parameter update information for each of the micro-batches;
transmitting, from the M th compute node to the (M−1) th compute node, the (L−1) th -layer error states;
sequentially applying, by the (M−1) th compute node, each of the (L−1) th -layer error states to the (L−1) th -layer to generate respective (L−2) th -layer error states for each of the micro-batches and respective (L−1) th -layer parameter update information for each of the micro-batches, wherein the (M−1) th compute node applying a particular one of the (L−1) th -layer error states to the (L−1) th layer at least partially overlaps in time with the M th compute node transmitting a subsequent one of the (L−1) th -layer error states to the (M−1) th compute node;
transmitting, from the first compute node to the M th compute node, (L−M) th -layer error states for each of the micro-batches; and
sequentially applying, by the M th compute node subsequent to generating the (L−1) th -layer error states, each of the (L−M) th -layer error states to an (L−M) th layer of the L layers to generate respective (L−M−1) th -layer activation states for each of the micro-batches and respective (L−M) th -layer parameter update information for each of the micro-batches.
12 . The article of manufacture of claim 11 , wherein the operations further comprise:
transmitting, to the first compute node, a portion of the plurality of parameters of the machine learning model that define the (M+1) th -layer prior to the first compute node applying each of the M th -layer activation states to the (M+1) th layer to generate respective (M+1) th -layer activation states for each of the micro-batches,
wherein the first compute node applying at least one of the micro-batches to the first layer of the L layers to generate at least one respective first-layer activation states at least partially overlaps in time with transmitting the portion of the plurality of parameters of the machine learning model that define the (M+1) th -layer to the first compute node.
13 . The article of manufacture of claim 11 , wherein the operations further comprise:
based on the L th -layer parameter update information for each of the micro-batches, determining parameter updates for the plurality of parameters of the machine learning model that define the L th -layer.
14 . The article of manufacture of claim 13 , wherein determining parameter updates for the plurality of parameters of the machine learning model that define the L th -layer comprises:
transmitting, from the M th compute node to an additional compute node that is not one of the M compute nodes, the L th -layer parameter update information for each of the micro-batches; and
transmitting, from the additional compute node to the M th compute node, L th -layer parameter update information for one or more additional micro-batches that are not part of the n micro-batches.
15 . A computer-implemented method comprising:
operating M compute nodes to train a machine learning model based on a batch of N training examples, wherein the batch of N training examples is divided into n micro-batches, wherein the machine learning model comprises L layers, each layer defined by a respective plurality of parameters of the machine learning model, wherein operating the M compute nodes to train the machine learning model based on the batch of N training examples comprises updating the parameters of the machine learning model by:
sequentially applying, by a first compute node of the M compute nodes, each of the micro-batches to a first layer of the L layers to generate respective first-layer activation states;
transmitting, from the first compute node to a second compute node of the M compute nodes, the first-layer activation states;
sequentially applying, by the second compute node, each of the first-layer activation states to a second layer of the L layers to generate respective second-layer activation states for each of the micro-batches, wherein the second compute node applying a particular one of the first layer activation states to the second layer at least partially overlaps in time with the first compute node transmitting a subsequent one of the first layer activation states to the second compute node;
transmitting, from an Mth compute node of the M compute nodes to the first compute node, Mth-layer activation states for each of the micro-batches;
sequentially applying, by the first compute node subsequent to generating the first-layer activation states, each of the Mth-layer activation states to an (M+1) layer of the L layers to generate respective (M+1)th-layer activation states for each of the micro-batches;
sequentially applying, by the Mth compute node, (L−1)th-layer activation states received from an (M−1)th compute node of the M compute nodes to an L th layer of the L layers to generate respective L th -layer activation states for each of the micro-batches; and
based on the L th -layer activation states, generating parameter update information for each of the micro-batches and for each of the L layers of the machine learning model.
16 . The computer-implemented method of claim 15 , further comprising:
transmitting, to the first compute node, a portion of the plurality of parameters of the machine learning model that define the (M+1) th -layer prior to the first compute node applying each of the M th -layer activation states to the (M+1) th layer to generate respective (M+1) th -layer activation states for each of the micro-batches,
wherein the first compute node applying at least one of the micro-batches to the first layer of the L layers to generate at least one respective first-layer activation states at least partially overlaps in time with transmitting the portion of the plurality of parameters of the machine learning model that define the (M+1) th -layer to the first compute node.
17 . The computer-implemented method of claim 15 , wherein the number of micro-batches n equals the number of compute nodes M.
18 . The computer-implemented method of claim 15 , wherein each micro-batch includes a single training example from the batch of N training examples.
19 . The computer-implemented method of claim 15 , wherein generating parameter update information for each of the micro-batches and for each of the L layers of the machine learning model comprises:
commencing, by the M th compute node, determination of parameter updates for the plurality of parameters of the machine learning model that define the L th -layer prior to the (M−1) th compute node completing the generation of (L−2) th -layer error states by applying (L−1) th -layer error states received from the M th compute node to the (L−1) th -layer.