TRAINING ULTRA-LARGE-SCALE VISION TRANSFORMER NEURAL NETWORKS

• Main Authors: DJOLONGA, Josip, HOULSBY, Neil Matthew Tinmouth, DEHGHANI, Mostafa, GILMER, Justin Morgan, MUSTAFA, Basil, PADLEWSKI, Piotr Michal, HEEK, Jonathan
• Format: Patent
• Language: English; French
• Published: 02.08.2024

Methods, systems, and apparatus, including computer programs encoded on computer storage media, 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. The plurality of layers comprise a fully connected layer having a plurality of parameters arranged in a row dimension and a column dimension. One of the methods comprises: generating a plurality of parameter blocks by partitioning the plurality of parameters along the row dimension and the column dimension; determining a ratio of a number of parameters along the row dimension relative to a number of parameters along the column dimension; and determining 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. Procédés, systèmes et appareil, comprenant des programmes informatiques codés sur des supports d'enregistrement informatiques, servant à traiter une entrée par l'intermédiaire de chaque couche d'une pluralité de couches d'un réseau de neurones artificiels pour générer une sortie à l'aide d'une pluralité d'accélérateurs matériels. La pluralité de couches comprend une couche entièrement connectée possédant une pluralité de paramètres agencés dans une dimension de rangée et une dimension de colonne. L'un des procédés consiste à : générer une pluralité de blocs de paramètres par partitionnement de la pluralité de paramètres le long de la dimension de rangée et de la dimension de colonne ; déterminer un rapport d'un nombre de paramètres le long de la dimension de rangée par rapport à un certain nombre de paramètres le long de la dimension de colonne ; et déterminer s'il convient d'utiliser une fragmentation de rangée ou une fragmentation de colonne avec la pluralité d'accélérateurs matériels en vue de calculer une sortie pour la couche entièrement connectée, puis calculer la sortie pour la couche entièrement connectée à l'aide d'une fragmentation de rangée ou d'une fragmentation de colonne.