Layered gradient accumulation and modular pipeline parallelism for improved training of machine learning models

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 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)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 Mth compute node, (L−1)th-layer activation states for each of the micro-batches;
• sequentially applying, by the Mth compute node, each of the (L−1)th-layer activation states to an Lth layer of the L layers to generate respective Lth-layer activation states for each of the micro-batches;
• based on the Lth-layer activation states, generating respective Lth-layer error states for each of the micro-batches;
• sequentially applying, by the Mth compute node, each of the Lth-layer error states to the Lth layer to generate respective (L−1)th-layer error states for each of the micro-batches and respective Lth-layer parameter update information for each of the micro-batches;
• transmitting, from the Mth 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 Mth 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 Mth compute node, (L−M)th-layer error states for each of the micro-batches;
• and
• sequentially applying, by the Mth 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.

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 Lth layer of the L layers to generate respective Lth-layer activation states for each of the micro-batches;
• and
• based on the Lth-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.