Training giant neural networks using pipeline parallelism

A method comprising:

• obtaining data specifying a partitioning of a neural network having multiple network layers into N composite layers that form a sequence of composite layers, wherein each composite layer comprises a distinct plurality of layers from the multiple network layers of the neural network;
• obtaining data assigning each of the N composite layers to one or more computing devices from a set of N computing devices;
• partitioning a mini-batch of training examples for training the neural network into a plurality of micro-batches;
• and
• training the neural network on the mini-batch of training examples, comprising:
• performing a forward pass through the neural network over a plurality of forward time steps until output activations have been computed for each micro-batch for a final composite layer in the sequence, comprising, for a particular forward time-step:
• determining whether, for any of the plurality of micro-batches, input activations to the first composite layer assigned to a first computing device are available as of the particular forward time-step,
• in response to determining that the input activations to the composite layer assigned to the first computing device are available as of the particular forward time-step for a first micro-batch, performing, at the first computing device, a forward pass through the first composite layer assigned to the first computing device using the available input activations for the first micro-batch to generate output activations for the composite layer for the first micro-batch,
• determining whether input activations to another composite layer that follows the first composite layer and is assigned to another computing device are available as of the particular forward time-step, wherein the input activations to the other composite layer are output activations from the first composite layer for another micro-batch, and
• in response to determining that the input activations to the other composite layer are available as of the particular forward time-step for the other micro-batch, while performing, at the first computing device, the forward pass through the first composite layer using the available input activations for the first micro-batch, performing, at the other computing device, the forward pass through the other composite layer using the available input activations for the other micro-batch,
• performing a backward pass through the neural network over a plurality of backward time steps until output gradients have been computed for each micro-batch for a first composite layer in the sequence, comprising, at a particular backward time-step:
• determining whether, for any of the plurality of micro-batches, gradient inputs for computing gradients of an objective function with respect to parameters of a second composite layer assigned to a second computing device are available as of the particular backward time-step, and
• in response to determining that input gradients to the second composite layer assigned to the second computing device are available as of the particular backward time-step for a second micro-batch, performing, at the second computing device, a backward pass through the second composite layer assigned to the second computing device using the available input gradients for second micro-batch to generate output gradients for the second composite layer for the second micro-batch,
• determining whether input gradients to another composite layer that precedes the second composite layer and is assigned to another computing device are available as of the particular backward time-step, wherein the input gradients to the other composite layer are output gradients from the second composite layer for another micro-batch, and
• in response to determining that the input gradients to the other composite layer are available as of the particular backward time-step for the other micro-batch, while performing, at the second computing device, the backward pass through the second composite layer using the available input gradients for the second micro-batch, performing, at the other computing device, the backward pass through the other composite layer using the available input gradients for the other micro-batch;
• for each composite layer, accumulating, by the respective computing device assigned to the respective composite layer and after performing the backward pass for all of the plurality of micro-batches of the mini-batch, respective gradients computed by the computing device for the respective composite layer for all of the plurality of micro-batches of the mini-batch to generate respective accumulated gradients for the composite layer, and
• for each composite layer, updating, by the respective computing device assigned to the respective composite layer, parameters of the respective composite layer using the respective accumulated gradients.

One or more non-transitory computer-readable storage media encoded with instructions that, when executed by one or more computers, cause the one or more computers to perform operations, operations comprising:

• obtaining data specifying a partitioning of a neural network having multiple network layers into N composite layers that form a sequence of composite layers, wherein each composite layer comprises a distinct plurality of layers from the multiple network layers of the neural network;
• obtaining data assigning each of the N composite layers to one or more computing devices from a set of N computing devices;
• partitioning a mini-batch of training examples for training the neural network into a plurality of micro-batches;
• and
• training the neural network on the mini-batch of training examples, comprising:
• performing a forward pass through the neural network over a plurality of forward time steps until output activations have been computed for each micro-batch for a final composite layer in the sequence, comprising, for a particular forward time-step:
• determining whether, for any of the plurality of micro-batches, input activations to the first composite layer assigned to a first computing device are available as of the particular forward time-step,
• in response to determining that the input activations to the composite layer assigned to the first computing device are available as of the particular forward time-step for a first micro-batch, performing, at the first computing device, a forward pass through the first composite layer assigned to the first computing device using the available input activations for the first micro-batch to generate output activations for the composite layer for the first micro-batch,
• determining whether input activations to another composite layer that follows the first composite layer and is assigned to another computing device are available as of the particular forward time-step, wherein the input activations to the other composite layer are output activations from the first composite layer for another micro-batch, and
• in response to determining that the input activations to the other composite layer are available as of the particular forward time-step for the other micro-batch, while performing, at the first computing device, the forward pass through the first composite layer using the available input activations for the first micro-batch, performing, at the other computing device, the forward pass through the other composite layer using the available input activations for the other micro-batch,
• performing a backward pass through the neural network over a plurality of backward time steps until output gradients have been computed for each micro-batch for a first composite layer in the sequence, comprising, at a particular backward time-step:
• determining whether, for any of the plurality of micro-batches, gradient inputs for computing gradients of an objective function with respect to parameters of a second composite layer assigned to a second computing device are available as of the particular backward time-step, and
• in response to determining that input gradients to the second composite layer assigned to the second computing device are available as of the particular backward time-step for a second micro-batch, performing, at the second computing device, a backward pass through the second composite layer assigned to the second computing device using the available input gradients for second micro-batch to generate output gradients for the second composite layer for the second micro-batch,
• determining whether input gradients to another composite layer that precedes the second composite layer and is assigned to another computing device are available as of the particular backward time-step, wherein the input gradients to the other composite layer are output gradients from the second composite layer for another micro-batch, and
• in response to determining that the input gradients to the other composite layer are available as of the particular backward time-step for the other micro-batch, while performing, at the second computing device, the backward pass through the second composite layer using the available input gradients for the second micro-batch, performing, at the other computing device, the backward pass through the other composite layer using the available input gradients for the other micro-batch;
• for each composite layer, accumulating, by the respective computing device assigned to the respective composite layer and after performing the backward pass for all of the plurality of micro-batches of the mini-batch, respective gradients computed by the computing device for the respective composite layer for all of the plurality of micro-batches of the mini-batch to generate respective accumulated gradients for the composite layer, and
• for each composite layer, updating, by the respective computing device assigned to the respective composite layer, parameters of the respective composite layer using the respective accumulated gradients.