Memory management method, artificial intelligence processing system, and computer program product for training-stage-based memory allocation in artificial intelligence model training

Abstract

A memory management method, including: acquiring multiple architecture parameters of an artificial intelligence model; obtaining a memory allocation size used in a training process of the artificial intelligence model according to the multiple architecture parameters; configuring a target memory region in random access memory according to the memory allocation size, wherein train processing stage in the training process of the artificial intelligence model shares the target memory region; and temporarily storing intermediate data corresponding to each train processing stage of the training process of the artificial intelligence model in the target memory region, wherein the intermediate data is transferred between accelerator processor module, the random access memory, and a storage device. Thereby, the present invention can automatically calculate the optimized memory allocation size, and through the design of shared memory space, effectively improve memory fragmentation issues.