Training ultra-large-scale vision transformer neural networks

A system comprising one or more computers and one or more storage devices storing instructions that when executed by the one or more computers cause the one more computers to perform operations for processing an input through each of a plurality of layers of a neural network to generate an output using a plurality of hardware accelerators, wherein the plurality of layers comprise a fully connected layer having a plurality of parameters arranged in a row dimension and a column dimension, and wherein the operations comprise:

• generating a plurality of parameter blocks by partitioning the plurality of parameters along the row dimension and the column dimension according to a number of the plurality of hardware accelerators;
• determining a ratio of a number of parameters along the row dimension relative to a number of parameters along the column dimension;
• and
• determining, from the ratio, whether to use row sharding or column sharding with the plurality of hardware accelerators to calculate an output for the fully connected layer and then calculating the output for the fully connected layer using either row sharding or column sharding.

One or more computer storage media storing instructions that when executed by one or more computers cause the one more computers to perform operations for processing an input through each of a plurality of layers of a neural network to generate an output using a plurality of hardware accelerators, wherein the plurality of layers comprise a fully connected layer having a plurality of parameters arranged in a row dimension and a column dimension, and wherein the operations comprise:

• generating a plurality of parameter blocks by partitioning the plurality of parameters along the row dimension and the column dimension according to a number of the plurality of hardware accelerators;
• determining a ratio of a number of parameters along the row dimension relative to a number of parameters along the column dimension;
• and
• determining, from the ratio, whether to use row sharding or column sharding with the plurality of hardware accelerators to calculate an output for the fully connected layer and then calculating the output for the fully connected layer using either row sharding or column sharding 8.

A system for performing a machine learning task on a network input to generate a network output, the system comprising one or more computers and one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or more computers to implement:

• a vision Transformer neural network configured to perform the machine learning task, the attention neural network comprising a plurality of attention blocks, each attention block comprising an attention layer and an input multi-layer perceptron (MLP), the attention block configured to:
• receive an input sequence for the block comprising a respective input element at each of a plurality of positions;
• provide the input sequence to the attention layer and to the input MLP,
• the attention layer configured to generate an attended input sequence that includes a respective attended input element for each of the plurality of positions at least in part by applying an attention mechanism to the input sequence for the block, wherein applying the attention mechanism (i) requires applying a query parameter matrix, a key parameter matrix, and a value parameter matrix to the input sequence for the layer but (ii) does not require adding any bias vector, and
• the input MLP configured to generate a transformed input sequence that includes a respective transformed input element for each of the plurality of positions by using one or more feed-forward neural network layers included in the input MLP to process the input sequence for the block, wherein generating the transformed input sequence requires both (i) applying a feed-forward input projection parameter matrix to the input sequence for the layer to generate a projected input sequence and (ii) adding an input projection bias vector to the projected input sequence;
• and
• generate the output sequence for the block from the attended input sequence and the transformed input sequence.