Neural network training in a distributed system

Abstract

In one example, a method comprises: performing backward propagation computations for a second layer of a neural network to generate second weight gradients; splitting the second weight gradients into a plurality of subsets each associated with a second exchange operation over a computer network, a number of the second weight gradients included in the each subset being based on at least one of first characteristics of the computer network or second characteristics of the neural network; performing backward propagation computations for a first layer of the neural network to generate first weight gradients in parallel with at least one of the second exchange operations; performing a first exchange operation to exchange the first weight gradients after the at least one of the second exchange operations completes; and after the first exchange operation completes, perform the remaining second exchange operations to exchange the remaining subsets of the second weight gradients.