MEMORY-EFFICIENT DISTRIBUTED NEURAL NETWORK TRAINING USING REMATERIALIZATION

摘要

Various embodiments provide apparatuses, systems, and methods related to a first worker of a distributed neural network (NN), The first worker may execute a forward training pass of a first node of a distributed NN, wherein execution of the forward training pass includes generation of a first computational graph (CG) that is based on inputs related to a second node that is processed by a second worker of the distributed NN. The first worker may also delete, subsequent to the forward training pass of the first node, the CG. The first worker may also execute, a backward pass of the first node, wherein execution of the backward pass includes regeneration of at least a portion of the first CG. Other embodiments may be described and claimed.