Memory management method for training transformer-based models using memory usage profiling

A memory management method for training an untrained transformer-based model using memory usage profiling, the memory management method comprising:

• performing at least two deep learning memory usage profiling sessions, each using a different batch size, for training the untrained model;
• generating a first memory usage profile log during a first memory usage profiling session and a second memory usage profile log during a second memory usage profiling session;
• determining an optimal batch size based on the first memory usage profile log and the second memory usage profile log;
• generating a memory schedule in accordance with the optimal batch size;
• performing a deep learning model training session at the optimal batch size to train the untrained model, performing custom memory allocation in accordance with the memory schedule;
• and
• generating and outputting a trained model.

A computing system for optimizing a training of an untrained model, the computing system comprising:

• processing circuitry and memory storing instructions that, when executed, cause the processing circuitry to:
• perform at least two deep learning memory usage profiling sessions, each using a different batch size, for training the untrained model;
• intercept and log memory allocation calls during each memory usage profiling session to generate a first memory usage profile log during a first memory usage profiling session and a second memory usage profile log during a second memory usage profiling session;
• classify memory objects identified in the first memory usage profile log and the second memory usage profile log as a same-size object or a proportional-size object;
• process the classified memory objects to simulate memory usage for a plurality of candidate batch sizes;
• determine an optimal batch size among the plurality of candidate batch sizes which maximizes memory usage without exceeding an available memory capacity;
• perform a deep learning model training session at the optimal batch size to train the untrained model;
• and
• generate and output a trained model, wherein
• the optimal batch size is determined by iterating through the plurality of candidate batch sizes in a loop, assigning offsets to the classified memory objects and predicting a total memory usage for each candidate batch size, terminating the loop when the predicted total memory usage for the candidate batch size exceeds the available memory capacity to select a last feasible batch size among the plurality of candidate batch sizes as the optimal batch size.