Runtime extension for neural network training with heterogeneous memory

A method comprising:

• managing buffer transfers between a first memory device and a second memory device during training of a neural network comprising a plurality of layers
• during a forward propagation pass of the neural network:
• during processing of a first layer of the neural network:
• allocating space in the second memory device for each buffer used as either an input or an output during processing of a current layer of a first layer of the plurality of layers;
• storing an identification of a first buffer to be transferred from the second memory device to the first memory device, in response to detecting the first buffer is used as an input during processing of the first layer;
• and
• during processing of a second layer of the plurality of layers of the neural network, reserving space used by the first buffer in the second memory for use in storing an output of the second layer, without allocating new space for the output of the second layer.

An apparatus comprising:

• a first processor comprising circuitry configured to execute a runtime manager;
• and
• a second processor comprising circuitry configured to execute a neural network comprising a plurality of executable layers;
• wherein when executed by the first processor, the runtime manager is configured to:
• manage buffer transfers between a first memory device and a second memory device during training of a neural network comprising a plurality of layers;
• during a forward propagation pass of the neural network:
• during processing of a first layer of the neural network:
• allocate space in the second memory device for each buffer used as either an input or an output during processing of a current layer of the plurality of layers;
• store an identification of a first buffer to be transferred from the second memory device to the first memory device, in response to detecting the first buffer is used as an input during processing of the first layer;
• and
• during processing of a second layer of the plurality of layers of the neural network, reserve space used by the first buffer in the second memory device for use in storing an output of the second layer, without allocating new space for the output of the second layer.