Adaptive memory management system for efficient training of large language models

A system for adaptive memory management, comprising:

• a computing system comprising processing circuitry including a central processing unit (CPU), graphics processing units (GPUs), and memory;
• adaptive memory management executable by the computing system, the adaptive memory management configured to, when executed by the computing system, at least:
• determine a chunk size for organizing parameters, the chunk size determined based on execution order of parameters and selected to reduce memory waste;
• determine a swapping interval comprising one activation swapping block followed by an integer number of gradient checkpointing blocks, the swapping interval determined by dividing a time required to swap one transformer block by a computation time of the one transformer block;
• determine a number of persistent chunks and non-persistent chunks to offload model states from the GPUs;
• determine a number of chunk buffers on the GPUs for prefetching and reusing the model states;
• determine a number of activation swapping blocks to offload activations from the GPUs;
• determine a number of transformer blocks to apply gradient checkpointing;
• and
• initiate a forward computation pass based upon the swapping interval, the number of persistent chunks and non-persistent chunks to offload, the number of chunk buffers, the number of activation swapping blocks and the number of transformer blocks.