Deep learning model training method and memory optimization device

Abstract

A deep learning model training method and a GPU optimization device are disclosed. The method comprises: in a first training stage, obtaining GPU lifetime information and GPU size information of each tensor in the stage; classifying the tensors according to the GPU lifetime information, and generating a GPU allocation scheme according to the tensor classification result and the GPU size information; in a second training stage, performing GPU allocation based on the GPU allocation scheme. The disclosure uses the prior knowledge of the predictable allocation mode in model training to plan the GPU allocation in advance, effectively reduces the GPU fragmentation without affecting the training performance, and thus improves the resource use efficiency in the training process. The disclosure can also improve the generation efficiency of the GPU allocation scheme through spatial aggregation of temporary tensors grouped by size and tensor merging before aggregation, and by introducing a static and dynamic GPU allocator, the GPU optimization scheme of the disclosure can achieve efficient GPU allocation on dense and MoE models.

Description

Technical Field

[0001] This disclosure relates to the field of artificial intelligence technology, and in particular to a deep learning model training method and a memory optimization device. Background Technology

[0002] With the rapid development of artificial intelligence technology, deep learning models have been widely used in fields such as image recognition, speech recognition, and natural language processing. To improve the expressive power and inference performance of models, the scale of neural networks is constantly expanding, and the number of model parameters continues to grow. In recent years, Large Language Models (LLMs), as an important development direction of deep learning, have seen their parameter scale gradually expand from hundreds of millions to tens of billions or even trillions, placing higher demands on computing resources and GPU memory capacity during the training process.

[0003] Against this backdrop, LLM training is often difficult to complete on a single GPU. To address the issue of insufficient GPU memory, research and engineering practices mainly employ two types of optimization techniques: one is model parallelism, including tensor parallelism and pipeline parallelism, which distributes the GPU memory load by dividing the model's parameters or computational tasks across multiple GPUs; the other is non-parallel memory optimization techniques, such as activation recomputation, which saves storage space by recompiling intermediate results during the backpropagation phase, and tensor offload, which removes temporarily unused tensors from GPU memory and retrieves them when necessary.

[0004] While the aforementioned technologies effectively alleviate memory pressure, they also introduce new challenges, particularly at the memory allocation level. These technologies disrupt the default memory management logic of deep learning frameworks. Taking the deep learning framework PyTorch as an example, its memory allocation mechanism is based on a "best-fit" strategy, relying on relatively stable lifecycles and size patterns for memory requests. However, by adopting strategies such as activation recomputation, tensor unloading, and tensor reuse, the memory allocation and release behavior of tensors becomes dynamic and uncertain. This makes it difficult for the original memory scheduling strategy to remain efficient, leading to increased memory fragmentation, decreased memory utilization, and impacting resource efficiency and system stability during training. Summary of the Invention

[0005] To address this, this disclosure proposes a deep learning model training method and a memory optimization device for deep learning model training. By utilizing prior knowledge of predictable allocation patterns during model training, memory allocation can be planned in advance, thereby effectively reducing memory fragmentation without affecting training performance and improving resource utilization efficiency during training.

[0006] According to a first aspect of this disclosure, a deep learning model training method is provided, comprising: in a first training phase of the deep learning model, acquiring memory behavior information, the memory behavior information including memory lifetime information and memory size information of each tensor in the first training phase, wherein the first training phase includes at least three iteration cycles; generating a memory allocation scheme based on the memory behavior information, wherein tensors are classified according to the memory lifetime information, and the memory allocation scheme is generated based on the tensor classification results and the memory size information; and in a second training phase of the deep learning model, performing memory allocation based on the memory allocation scheme.

[0007] Optionally, generating a memory allocation scheme based on the memory behavior information includes: classifying tensors with memory lifetimes less than a first lifetime threshold as spatially aggregable tensors; grouping spatially aggregable tensors according to the memory size information; and spatially aggregating multiple spatially aggregable tensors within the same group whose memory lifetimes do not overlap to generate a first memory allocation sub-scheme for memory reuse of spatially aggregable tensors.

[0008] Optionally, generating a video memory allocation scheme based on the video memory behavior information includes: classifying tensors with a video memory lifetime greater than a second lifetime threshold as non-space-aggregateable tensors and generating a second video memory allocation sub-scheme for allocating contiguous video memory to the non-space-aggregateable tensors based on the video memory size information of the non-space-aggregateable tensors.

[0009] Optionally, classifying tensors with a video memory lifetime less than a first lifetime threshold as spatially aggregateable tensors includes: merging multiple tensors of a specific type with a video memory lifetime less than the first lifetime threshold and belonging to the same third lifetime threshold interval into a virtual tensor, wherein other spatially aggregateable tensors that do not belong to the specific type are non-virtual tensors, and grouping spatially aggregateable tensors according to the video memory size information includes: combining multiple virtual tensors and multiple non-virtual tensors to form a tensor set; and grouping the tensors in the tensor set according to the video memory size information of each tensor in the tensor set.

[0010] Optionally, spatially aggregating multiple spatially aggregable tensors with non-overlapping memory lifetimes within the same group to generate a first memory allocation sub-scheme for memory reuse of spatially aggregable tensors includes: sorting the spatially aggregated groups in descending order; sequentially arranging memory for each group according to the sorted order, wherein the memory arrangement operation for the current group includes: searching for an unoccupied space region in the allocated memory space that can accommodate the current group; in response to finding an unoccupied space region that meets the requirements, arranging the current group into the unoccupied space region; in response to not finding an unoccupied space region that meets the requirements, arranging the current group into the unallocated memory space.

[0011] Optionally, memory allocation based on the memory allocation scheme includes: a static allocator allocating memory to the current memory request based on the memory allocation scheme; and a dynamic allocator allocating memory to the current memory request in response to detecting that the current memory request does not conform to the memory allocation scheme.

[0012] Optionally, the video memory lifecycle information of the tensor is determined based on the video memory allocation node and video memory release node of the tensor, and the video memory allocation by the static allocator based on the video memory allocation scheme includes: allocating a complete video memory space based on the video memory allocation scheme; the static allocator reading the pointers and sizes of each static video memory request in the video memory allocation scheme; and in response to receiving a video memory request with corresponding video memory allocation node and video memory size information, returning the pointer corresponding to the video memory request in the video memory space.

[0013] Optionally, the deep learning model is a Mixture of Experts (MoE) model including a dynamic layer, and the memory allocation based on the memory allocation scheme includes: in response to receiving a memory request from the dynamic layer, a static allocator performs memory allocation based on the memory allocation scheme; in response to the static allocator not finding available memory, a dynamic allocator performs memory allocation for the memory request from the dynamic layer.

[0014] Optionally, the first training phase corresponds to the first N training iterations before the model starts training, where N is an integer not less than 3, and the second training phase corresponds to a complete training process where the model restarts.

[0015] Optionally, the memory allocation scheme includes a first iteration cycle memory allocation scheme and a second iteration cycle memory allocation scheme, and in the second training phase of the deep learning model, memory allocation based on the memory allocation scheme includes: in the first iteration cycle of the restarted complete training process, memory allocation is performed based on the first iteration cycle memory allocation scheme; and in each iteration cycle after the first iteration cycle of the restarted complete training process, memory allocation is performed based on the second iteration cycle memory allocation scheme.

[0016] According to a second aspect of this disclosure, a deep learning model training method is provided, comprising: recording memory behavior information of the deep learning model for the first N training iterations, the memory behavior information including memory lifetime information and memory size information of each tensor in the first N training iterations, wherein N is an integer not less than 3; generating a memory allocation scheme based on the recorded memory behavior information; and allocating memory based on the memory allocation scheme in the last M training iterations of the deep learning model, wherein M is an integer greater than N, and the last M training iterations correspond to: M training iterations following the first N training iterations in a complete training process of the deep learning model, or M training iterations in a newly started complete training process. Generating a memory allocation scheme based on the recorded memory behavior information includes: classifying tensors into parameter tensors, optimizer tensors, and activation value tensors based on memory lifecycle information; generating a resident memory allocation sub-scheme based on the memory size information of the parameter tensors and the optimizer tensors; grouping the activation value tensors based on their memory size information, and spatially aggregating multiple activation value tensors with non-overlapping lifecycles within the same group to generate a spatially aggregated memory allocation sub-scheme; and merging the resident memory allocation sub-scheme and the spatially aggregated memory allocation sub-scheme to generate the memory allocation scheme.

[0017] According to a third aspect of this disclosure, a memory optimization apparatus for training deep learning models is provided, comprising a runtime component and an offline analysis component. The runtime component includes a memory allocator for recording memory behavior information in a memory behavior recording mode. The memory behavior information includes memory lifetime information and memory size information for each tensor. The memory behavior recording mode lasts for at least three iterations. The offline analysis component includes a memory behavior analysis module for classifying tensors based on the memory lifetime information; and a memory allocation scheme generation module for generating a memory allocation scheme based on the tensor classification results and the memory size information. Furthermore, the memory allocator is also used to allocate memory based on the memory allocation scheme in a memory allocation mode.

[0018] Optionally, the video memory allocator includes: a dynamic allocator for allocating video memory in the video memory behavior recording mode; and a static allocator for allocating video memory based on the video memory allocation scheme in the video memory scheme allocation mode.

[0019] Optionally, the memory allocator further includes: a semi-dynamic allocator, configured to, when the deep learning model is an expert hybrid model including dynamic layers, in response to receiving a memory request from the dynamic layer, cause a static allocator to allocate memory based on the memory allocation scheme, and in response to the static allocator not finding available memory, cause a dynamic allocator to allocate memory for the memory request from the dynamic layer.

[0020] Optionally, the runtime component further includes an error detection and recovery module, used to detect whether the current video memory request conforms to the video memory allocation scheme in the video memory allocation mode, wherein the video memory allocator is further used to allocate video memory for the current video memory request in response to detecting that the current video memory request does not conform to the video memory allocation scheme.

[0021] Optionally, the offline analysis component further includes: an abnormal memory detection module, used to diagnose and debug abnormal situations in the memory usage process based on the memory analysis report containing tensor classification information generated by the memory behavior analysis module, and the memory scheme generation module generates a memory allocation scheme based on the memory analysis report obtained through abnormal memory detection.

[0022] Therefore, the deep learning model training method and memory optimization device disclosed herein improve the generation efficiency of the memory allocation scheme by spatial aggregation of non-resident tensors grouped by size and tensor merging before aggregation, and by introducing a static and dynamic memory allocator, the memory optimization scheme disclosed herein can achieve efficient memory allocation on both dense models and MoE models. Attached Figure Description

[0023] The above and other objects, features and advantages of this disclosure will become more apparent from the more detailed description of exemplary embodiments thereof taken in conjunction with the accompanying drawings, wherein like reference numerals generally denote like parts.

[0024] Figure 1 A schematic flowchart of a deep learning model training method according to an embodiment of the present disclosure is shown.

[0025] Figure 2 The typical Transformer structure and the differences between the FFN module in the Dense and MoE models are shown.

[0026] Figure 3 A schematic diagram of a memory optimization apparatus according to an embodiment of the present disclosure is shown.

[0027] Figure 4 The present invention illustrates a specific application process of a memory optimization method according to an embodiment of the present disclosure in the training process of a deep learning model. Detailed Implementation

[0028] Preferred embodiments of the present disclosure will now be described in more detail with reference to the accompanying drawings. While preferred embodiments of the present disclosure are shown in the drawings, it should be understood that the present disclosure may be implemented in various forms and should not be limited to the embodiments set forth herein. Rather, these embodiments are provided so that the present disclosure will be thorough and complete, and will fully convey the scope of the disclosure to those skilled in the art.

[0029] As used in the specification and appended claims of this application, the singular expressions “a,” “the,” “the,” and “the” are intended to also include expressions such as “one or more,” unless the context explicitly indicates otherwise. The term “comprising” and its variations, as used herein, indicate an open-ended inclusion, i.e., “including but not limited to.” Unless specifically stated otherwise, the term “or” means “and / or.” The term “according to” means “at least in part according to.” The terms “an example embodiment” and “an embodiment” mean “at least one example embodiment.” The term “another embodiment” means “at least one additional embodiment.” The terms “first,” “second,” etc., may refer to different or the same objects. Other explicit and implicit definitions may also be included below.

[0030] Because the parallel strategies introduced during large model training alter the scheduling and lifecycle of tensors in the original training graph, they become incompatible with the default memory management mechanisms in mainstream deep learning frameworks (such as PyTorch's "best fit" strategy). Some systems choose to directly call the native memory allocation and deallocation interfaces provided by the underlying computing platform for greater flexibility. These interfaces are typically provided by the device driver layer or runtime and are used to dynamically allocate and deallocate memory blocks in the physical memory of the computing device. For example, on GPU platforms, memory management can be accomplished by calling native allocation functions (e.g., cudaMalloc and cudaFree in CUDA). However, native calls frequently result in kernel-mode context switching and lack a reuse mechanism, leading to high allocation and deallocation overhead, which significantly impacts training efficiency during iterative training.

[0031] To address the above issues, existing technical solutions, such as GMLake, propose a low-level Virtual Memory Stitching (VMS) mechanism. By calling the underlying Virtual Memory Management Interface (VMM), multiple non-contiguous physical video memory fragments are combined into a logically contiguous virtual address space, enabling programs to logically use contiguous video memory and improving fragment reuse capabilities.

[0032] In actual operation, GMLake uses the `fraglimit` parameter to control the smallest fragment unit participating in the stitching process. A smaller setting maximizes the utilization of fragmented GPU memory, but frequent stitching operations increase system overhead. A larger setting prevents the utilization of smaller fragments, affecting GPU memory resource utilization. This makes it difficult for GMLake to strike a balance between running speed and GPU memory efficiency, limiting its applicability in high-intensity training tasks.

[0033] As can be seen from the above, the existing technology lacks a memory optimization mechanism that takes into account both efficient operation and memory utilization, making it difficult to meet the actual needs of large-scale deep learning model training.

[0034] In view of this, this disclosure proposes a memory optimization mechanism, which can be specifically implemented as a deep learning model training method and a memory optimization device for training. Observations of most deep learning models show that the tensor allocation involved in each iteration during training has strong repeatability and predictability, and the memory lifecycle pattern remains relatively stable between iterations. Therefore, this disclosure records memory behavior information from earlier iterations, analyzes and extracts prior knowledge of predictable allocation patterns, and plans a memory allocation scheme for memory allocation in subsequent training iterations. This effectively reduces memory fragmentation and significantly improves resource utilization efficiency during training without affecting training performance.

[0035] Figure 1 A schematic flowchart illustrating a deep learning model training method according to an embodiment of this disclosure is shown. The memory optimization mechanism of this disclosure is applicable to various deep learning models where the memory lifecycle pattern remains relatively stable across iterations during training, especially suitable for scenarios where the input data size remains consistent across iterations, including cases where the input size remains consistent after padding operations. Based on this, this solution can fully utilize the predictability of memory allocation behavior to achieve efficient memory planning and reuse. In the following description, the training method of this disclosure will be primarily illustrated using the LLM training process as an example. The memory optimization mechanism of this disclosure can also be applied to other deep learning models, such as CNN convolutional neural networks, as long as their input data size remains consistent across iterations.

[0036] In step S110, during the first training phase of the deep learning model, memory behavior information is acquired. This information includes the memory lifetime information and memory size information of each tensor in the first training phase. Since the memory lifetime pattern remains relatively stable across iterations, the first training phase needs to include at least three iteration cycles. The main reason is that the first iteration often involves operations that only occur in the first round, such as model parameter initialization, optimizer state allocation, cache mechanism activation, and memory allocator warm-up. Furthermore, the memory usage pattern of dense models tends to stabilize from the second iteration cycle onwards, and theoretically, the memory behavior in each subsequent iteration cycle remains consistent. Therefore, the first training phase can correspond to the beginning of model training, such as the first N iterations of model training, where N can be an integer not less than 2. In some embodiments, to verify the consistency of memory behavior across iteration cycles from the second iteration cycle, N can be set to 3 or even larger. However, to maximize overall training efficiency, N is usually not set too large; for example, N is usually not greater than 5 (e.g., equal to 4 or 5). That is, in one embodiment, the first training phase corresponds to the first N iterations of model training, where N is at least equal to 3, and in a preferred embodiment, N is 3, 4 or 5.

[0037] In this first training phase, the memory behavior information of each tensor can be recorded. Specifically, the memory allocation and deallocation behaviors of each tensor can be recorded to obtain the memory lifecycle information of the tensor. The memory allocation behavior of a tensor includes the memory size information when the tensor allocates memory and the memory allocation node. The memory deallocation behavior of a tensor corresponds to the memory deallocation node of that tensor. If a tensor has not been deallocated by the end of the first training phase, it can be assumed that the memory deallocation node of that tensor is located after the end of the first training phase (for example, in the case of parameter tensors and optimizer tensors).

[0038] It should be understood that since the purpose of obtaining memory behavior information is to generate memory allocation schemes for subsequent iterations with similar memory allocation patterns, the memory allocation and release nodes of a tensor can be the model execution sub-stages corresponding to the allocation and release of that tensor in the first training phase, rather than simply the timing of the behavior execution. Specifically, the memory allocation node reflects which computation node or module (e.g., a certain layer of a forward or backward operation) the tensor first allocates memory in the training process, and the memory release node reflects which computation node or module (e.g., a certain layer of a forward or backward operation) the tensor is released in the training process, thus forming a memory lifecycle description closely related to the model structure and training logic.

[0039] In step S120, a memory allocation scheme is generated based on the memory behavior information. Tensors can be classified according to the memory lifecycle information, and the memory allocation scheme is generated based on the tensor classification results and the memory size information. Specifically, tensors can be divided into spatially aggregateable tensors and non-spatially aggregateable tensors based on the memory lifecycle information of each tensor included in the memory behavior information.

[0040] Here, "spatial aggregation" refers to mapping multiple tensors with non-overlapping lifecycles to the same memory region, thereby reusing memory space on the time axis to reduce total memory usage and fragmentation; that is, performing memory reuse. Spatially aggregable tensors have short memory lifecycles (corresponding to tensors with memory lifecycles less than a first lifecycle threshold), thus allowing the generation of memory allocation sub-schemes involving memory reuse. Specifically, spatially aggregable tensors can be grouped according to memory size information, and multiple spatially aggregable tensors with non-overlapping memory lifecycles within the same group can be spatially aggregated to generate a first memory allocation sub-scheme for memory reuse of spatially aggregable tensors.

[0041] Since non-space-aggregateable tensors have relatively long memory lifetimes (corresponding to tensors whose memory lifetimes are greater than the second lifetime threshold, where the second lifetime threshold is not less than the first lifetime threshold), a second memory allocation sub-scheme that does not involve memory reuse is generated for them. Because it does not involve memory reuse, the second memory allocation sub-scheme can allocate contiguous memory to non-space-aggregateable tensors. These two memory allocation sub-schemes can be merged to generate the final memory allocation scheme.

[0042] Spatially aggregateable and non-spatially aggregateable tensors are categorized based on memory lifetime, a classification of tensors used during training, from the perspective of whether they can be reused in memory. If classified according to their purpose during model training, tensors can be classified into parameter tensors, optimizer tensors, and activation tensors. Parameter tensors store the model's learnable parameters; they are allocated at the start of training and remain in memory until training ends, their lifetime covering the entire training process, making them long-lived tensors. Optimizer tensors store optimizer state information, such as momentum and gradient accumulation terms; they are allocated after the first training iteration and persist in subsequent training until completion. Therefore, their lifetime begins at the end of the first iteration and continues throughout the remaining training phases, also making them long-lived tensors. In contrast, activation tensors are intermediate computation results temporarily generated during each forward and backward propagation; they are dynamically allocated and released within each training iteration, their lifetime covering only a single iteration, making them short-lived tensors. Due to their distinct lifecycle characteristics as described above, tensors during training can be categorized into parameter tensors, optimizer tensors, and activation tensors based on their memory lifecycle information. Furthermore, since parameter and optimizer tensors have long lifecycles, while activation tensors have short lifecycles, in some embodiments, parameter and optimizer tensors can be considered as non-space-aggregateable tensors, and activation tensors as space-aggregateable tensors. Further, tensors can be classified into resident tensors and non-resident tensors based on their memory residency time. Based on these lifecycle characteristics, in some embodiments, parameter and optimizer tensors can be considered as resident tensors, and activation tensors as temporary tensors. Regardless of the specific classification method used, tensor classification can be based on memory lifecycle information. It should be understood that the first training phase needs to cover the first iteration cycle of training to distinguish parameter and optimizer tensors based on different memory allocation nodes (the memory allocation node for parameter tensors is at the beginning of the first iteration cycle, and for optimizer tensors, it is at the end of the first iteration cycle).

[0043] After generating the memory allocation information, memory allocation based on the memory allocation scheme can be performed in step S130, during the second training phase of the deep learning model.

[0044] Here, both "first training phase" and "second training phase" refer to training phases, i.e., phases where model parameters are optimized using training data through forward propagation, loss calculation, and backpropagation. Since the prior knowledge of memory allocation extracted in the "first training phase" is needed, the "second training phase" obviously follows the "first training phase." The second training phase may include M complete iteration cycles. To fully utilize the extracted prior knowledge of memory allocation, M is an integer greater than N. In one embodiment, the "first training phase" and "second training phase" are located within the same complete training process. The first training phase, which acquires information about memory allocation behavior, corresponds to the first few iteration cycles after the start of model training (e.g., the first 3 to 5 iteration cycles), while the second training phase corresponds to all iteration cycles from the end of the first training phase until the model training is complete. In another embodiment, the "first training phase" and "second training phase" may not be located within the same complete training process. At this time, the first training phase corresponds to the first N training iterations before the model starts training, where N is an integer not less than 3, and the second training phase corresponds to at least a part of a complete training process that the model restarts, such as the first M iterations of a complete training process that has restarted, and preferably corresponds to a complete training process that has restarted.

[0045] Since a memory allocation scheme based on prior knowledge from previous iterations is already available, memory allocation can be statically performed during this second training phase. Here, "static" allocation is the opposite of "dynamic" allocation. Static allocation refers to pre-determining the memory allocation location and size of all tensors before model training or inference, allocating all memory resources required for the entire execution process in one go. Dynamic allocation refers to dynamically requesting and releasing memory resources based on runtime needs during model execution. Tensor memory is allocated only when actually needed and released after use.

[0046] When performing static memory allocation based on a memory allocation scheme, a complete block of memory space can be pre-allocated at the beginning of the second training phase of the model, and memory allocation can be performed for each memory request according to the method specified by the memory allocation scheme during execution. Specifically, during the execution of the second training phase, the corresponding memory request in the memory allocation scheme can be determined based on the memory request behavior of each memory request (i.e., memory request node and memory size information), and the corresponding memory can be allocated to it.

[0047] Since the memory allocation for model training is repeatable across iteration cycles, the memory allocation scheme generated in step S130 needs to cover the memory allocation of each tensor within at least one complete iteration cycle. However, in the second training phase, which includes the first iteration cycle and multiple subsequent iteration cycles, the memory allocation behavior differs between the first and subsequent iteration cycles (the first iteration cycle involves memory allocation of parameter tensors during training initialization, memory allocation of optimizer tensors at the end of the first iteration cycle, and optionally memory allocation of a small number of temporary tensors during training initialization; subsequent iteration cycles do not involve the above). Therefore, the memory allocation scheme generated in step S130 needs to guide the allocation in the first iteration cycle and multiple subsequent iteration cycles of the second training phase.

[0048] Therefore, the training method disclosed herein analyzes the memory behavior during the iteration cycle and extracts predictable allocation patterns as prior information to plan the memory allocation scheme for subsequent training. This can significantly reduce the memory fragmentation rate in subsequent training and improve the memory utilization efficiency and overall resource utilization during the training process without affecting training performance.

[0049] Since an iteration of LLM training involves a massive number of activation value tensors, to improve the efficiency of spatial aggregation, spatially aggregable tensors can be grouped according to their memory size. Then, multiple spatially aggregable tensors within the same group with non-overlapping memory lifetimes are spatially aggregated to generate a first memory allocation sub-scheme for memory reuse of spatially aggregable tensors. In one embodiment, spatially aggregable tensors with the same memory size can be grouped together. In other embodiments, spatially aggregable tensors with similar memory sizes (e.g., within a predetermined memory size range) can also be grouped together. Memory reuse is then performed on multiple tensors within the same group with non-overlapping memory lifetimes. Different algorithms can be used to determine the spatial aggregation method for tensors within the same group. In one embodiment, a group can be divided into multiple layers based on the tensor's lifetime, where the lifetimes of all tensors in each layer do not overlap, thus allowing them to be mapped to the same memory region. This layering process is implemented using a greedy algorithm, aiming to minimize the number of layers while ensuring that the lifetimes of tensors in the same layer do not overlap, thereby maximizing the efficiency of memory space reuse.

[0050] To further reduce the computational difficulty of spatial aggregation within a group, some spatially aggregateable tensors can be selected for merging. In this case, classifying tensors with a memory lifetime less than a first lifetime threshold as spatially aggregateable tensors includes: merging multiple tensors of a specific type whose memory lifetime is less than the first lifetime threshold and which belong to the same third lifetime threshold interval into a single virtual tensor; and considering other spatially aggregateable tensors that do not belong to the specific type as non-virtual tensors. Grouping spatially aggregateable tensors according to the memory size information includes: combining multiple virtual tensors and multiple non-virtual tensors to form a tensor set; and grouping the tensors in the tensor set according to the memory size information of each tensor in the tensor set.

[0051] The third lifecycle threshold range can correspond to a micro-batch. In the training process of a deep learning model, a training iteration typically corresponds to a complete parameter update process, that is, completing the forward propagation, back propagation, and gradient update of a batch of training samples. However, in training scenarios with large models or limited GPU memory, it is often impossible to process all samples in the entire training batch at once. Therefore, a micro-batch mechanism can be used to further divide a training batch into several smaller sub-batches. Each micro-batch independently performs forward and back propagation, accumulating gradients across multiple micro-batches. After all micro-batches have been processed, a unified parameter update operation is performed. In other words, a training iteration can consist of multiple micro-batches, which are processed sequentially and finally jointly complete an effective model weight update, thereby alleviating GPU memory pressure and ensuring training stability without changing the global batch size.

[0052] Within the same micro-batch, activation value tensors can be categorized into two types based on their lifetime: (1) tensors that allocate memory during the forward propagation of the micro-batch and release memory during the backward propagation of the same micro-batch, i.e., tensors spanning both forward and backward propagation; and (2) tensors that are allocated and released within the forward or backward propagation of the micro-batch, i.e., tensors within the forward or backward propagation. Here, since all tensors spanning both forward and backward propagation within the same micro-batch follow the lifetime characteristic of continuous allocation and release, these tensors can be considered as a large virtual tensor. In such an embodiment, all tensors spanning both forward and backward propagation within the same micro-batch correspond to the aforementioned multiple specific types of tensors. The memory space required for a virtual tensor is equal to the theoretical peak value of all tensors spanning both forward and backward propagation within the micro-batch, i.e., equal to the sum of the memory sizes of these tensors. Other spatially aggregateable tensors that do not belong to the aforementioned specific types may include forward or backward internal tensors and correspond to the aforementioned non-virtual tensors. Therefore, by merging tensors across forward and backward propagation within the same micro-batch, the size of the tensors that need to be arranged during intra-group spatial aggregation is reduced, thereby lowering the computational complexity of intra-group spatial aggregation. For example, if an iteration includes four micro-batches, each micro-batch contains 100 tensors across forward and backward propagation and 500 tensors within forward or backward propagation. Without merging, grouping and intra-group spatial aggregation of 2400 tensors (2400 = (500 + 100) x 4) would be required. However, with merging, each micro-batch containing 100 tensors across forward and backward propagation is merged into one virtual tensor. In this case, grouping and intra-group spatial aggregation of 4 virtual tensors and 2000 non-virtual tensors (2000 = 500 x 4), totaling 2004 tensors, can be performed, thus reducing the computational complexity compared to the case of 2400 tensors. It should be understood that this example is only given to illustrate the effect of merging computations, and in practice, much more tensors across forward and backward propagation and internal tensors in forward or backward propagation are usually involved.

[0053] In one embodiment, the memory regions corresponding to the layered results of each group can be arranged sequentially or randomly. However, in another embodiment, the memory allocation requirements can be further reduced by filling gaps in the allocated memory space. In this case, generating the first memory allocation sub-scheme further includes: sorting the multiple groups after spatial aggregation in descending order; and sequentially arranging the memory for each group according to the sorted order, wherein the memory arrangement operation for the current group includes: searching for an unoccupied space region in the allocated memory space that can accommodate the current group; in response to finding an unoccupied space region that meets the requirements, arranging the current group into the unoccupied space region; and in response to not finding an unoccupied space region that meets the requirements, arranging the current group into the unallocated memory space.

[0054] The memory optimization mechanism disclosed herein can be implemented with the participation of a dynamic allocator and a static allocator. The dynamic allocator can be used for dynamic memory allocation during the first training phase. In this case, the dynamic allocator can have an implementation that is consistent with or inconsistent with the default dynamic allocator of the deep learning framework (e.g., a dynamic allocator employing a "best fit" strategy). If an Out of Memory (OOM) error occurs during the first training phase using the dynamic allocator, the native memory allocation and deallocation interface provided by the underlying computing platform (e.g., cudaMalloc and cudaFree in CUDA) can be used for dynamic memory allocation. The static allocator can then be used for memory allocation based on the memory allocation scheme during the second training phase. In this case, memory allocation by the static allocator based on the memory allocation scheme can include: allocating a complete memory space based on the memory allocation scheme; the static allocator reading the pointers and sizes of each static memory request in the memory allocation scheme; and, in response to receiving a memory request with corresponding memory request node and memory size information, returning the memory pointer corresponding to the memory request in the memory space. In some embodiments, the memory allocator, acting as a top-level module, may allocate a complete memory space to the static allocator, which is its sub-module, based on the memory allocation scheme. In other embodiments, the static allocator may allocate a complete memory space itself based on the memory allocation scheme.

[0055] As mentioned earlier, since the memory allocation between iterations is highly repeatable, in the second training phase, when an actual memory request is received, the system searches for the preceding memory request corresponding to this actual memory request in the memory allocation scheme and provides the actual memory request with the memory pointer allocated for this preceding memory request. The corresponding preceding memory request in the memory allocation scheme can be found through the memory request node and memory size information of the actual memory request.

[0056] In one embodiment, an error detection and recovery mechanism can be introduced for the second training phase. Specifically, when a memory allocation error is detected, dynamic memory allocation is performed by the dynamic allocator. In this case, memory allocation based on the memory allocation scheme includes: memory allocation by the static allocator based on the memory allocation scheme; and memory allocation by the dynamic allocator for the current memory request in response to the detection that the current memory request does not conform to the memory allocation scheme. After processing the current memory request, the static allocator can still continue to allocate memory based on the memory allocation scheme.

[0057] In deep learning model design, common architectures include dense models and MoE models. In dense models, all network modules are activated and participate in computation in each forward propagation, resulting in a fixed computation path. MoE models, on the other hand, employ a sparse activation mechanism. Through a gating network, they dynamically select a subset of expert subnetworks for activation, ensuring that only a few experts are computed in each forward propagation. This significantly reduces computational resource consumption while maintaining a large parameter capacity. Expert subnetworks are typically composed of different feedforward neural network (FFN) modules, exhibiting structural similarities but independent parameters. This architecture is particularly suitable for LLMs containing Transformer structures. For ease of understanding, Figure 2 The typical Transformer structure and the differences between the FFN module in the Dense and MoE models are shown.

[0058] As shown in the figure, a Transformer structure consists of self-attention, layer normalization, a feedforward neural network (FFN), and a layer normalization module. In dense models, the FFN structure is always activated in each forward propagation, and all inputs are processed by the same FFN module. However, in the MoE model, the FFN part is replaced by multiple parallel expert subnetworks, i.e., multiple feedforward sublayers (e.g., FFN1 and FFN2 in the figure), and a gating network (router) dynamically selects a subset of experts for activation based on the input content. For example, in... Figure 2 In the example, the router selected FFN1 for activation. Layers in the MoE model that dynamically select execution paths through routing mechanisms can be called dynamic layers, for example... Figure 2 The feedforward selection structure on the right.

[0059] Due to the introduction of an "expert selection" mechanism based on dynamic activation of input, the memory behavior of the MoE model is nondeterministic during training. A semi-dynamic mechanism can be introduced to make the training method of this disclosure applicable to the training process of the MoE model. In this case, memory allocation based on the memory allocation scheme can include: in response to receiving a memory request from the dynamic layer, a static allocator allocates memory based on the memory allocation scheme; and in response to the static allocator not finding available memory, a dynamic allocator allocates memory for the memory request from the dynamic layer. Therefore, while prioritizing the use of the planned memory layout, it is possible to flexibly respond to the uncertain memory demands of the dynamic layer during operation due to sparse activation, thereby ensuring resource utilization while also considering flexible response and efficient scheduling of the memory demands of the dynamic layer in the MoE model.

[0060] In one embodiment, this disclosure can also be implemented as a deep learning model training method, comprising: recording memory behavior information of the deep learning model for the first N training iterations, the memory behavior information including memory lifetime information and memory size information of tensors in the first N training iterations, wherein N is an integer not less than 2, and preferably an integer not less than 3; generating a memory allocation scheme based on the recorded memory behavior information; and allocating memory for tensors based on the memory allocation scheme in the last M training iterations of the deep learning model, wherein M is an integer greater than N. The last M training iterations correspond to: M training iterations following the first N training iterations in a complete training process of the deep learning model, or M training iterations in a newly started complete training process, and the last M training iterations preferably continue until the end of training. The process of generating a memory allocation scheme based on the recorded memory behavior information includes: classifying tensors into parameter tensors, optimizer tensors, and activation value tensors based on memory lifecycle information; generating a resident memory allocation sub-scheme based on the memory size information of the parameter tensors and the optimizer tensors; grouping the activation value tensors based on their memory size information, and spatially aggregating multiple activation value tensors with non-overlapping lifecycles within the same group to generate a spatially aggregated memory allocation sub-scheme; and merging the resident memory allocation sub-scheme and the spatially aggregated memory allocation sub-scheme to obtain the memory allocation scheme.

[0061] As previously mentioned, the first training phase corresponds to the first N training iterations before the model begins training, where N is an integer not less than 3, and the second training phase corresponds to a complete training process where the model restarts. In this case, the memory allocation scheme includes a first iteration memory allocation scheme and a second iteration memory allocation scheme. The first iteration memory allocation scheme can be a memory allocation scheme generated for the memory allocation behavior of the first iteration of model training included in the first training phase. The second iteration memory allocation scheme can be a memory allocation scheme generated for the memory allocation behavior of the second iteration (or possibly the third, fourth, or fifth iteration) of model training included in the first training phase, since the memory allocation behavior of each iteration after the first iteration is consistent, or at least largely consistent. In the second training phase of the deep learning model, memory allocation based on the memory allocation scheme includes: memory allocation based on the first iteration memory allocation scheme in the first iteration of a restarted complete training process; and memory allocation based on the second iteration memory allocation scheme in each iteration after the first iteration of the restarted complete training process. Furthermore, in the implementations of the first memory allocation sub-scheme for generating spatially aggregateable tensors (equivalent to a spatially aggregated memory allocation sub-scheme in some embodiments) and the second memory allocation sub-scheme for generating non-spatially aggregateable tensors (equivalent to a resident memory allocation sub-scheme in some embodiments), the first iteration cycle memory allocation scheme needs to include both the second and first memory allocation sub-schemes, while the second iteration cycle memory allocation scheme does not include the second memory allocation sub-scheme but must include the first memory allocation sub-scheme. When the activation value tensor allocations in the first and subsequent iteration cycles are the same or sufficiently similar, the first and second iteration cycle memory allocation schemes can reuse the same first memory allocation sub-scheme. However, when there is a sufficient difference in the activation value tensor allocations in the first and subsequent iteration cycles, the first and second iteration cycle memory allocation schemes can use different first memory allocation sub-schemes.

[0062] This disclosure enables the implementation of the above-described deep learning model training method using memory optimization devices. Figure 3 A schematic diagram of a memory optimization apparatus according to an embodiment of the present disclosure is shown. As shown, the memory optimization apparatus includes a runtime component 100 and an offline analysis component 200.

[0063] Runtime component 100 is responsible for memory allocation and management during deep learning model training, and mainly includes memory allocator 110. Memory allocator 110 is responsible for memory allocation during model training, supporting two working modes: memory behavior recording mode and memory allocation scheme mode. In memory behavior recording mode, memory allocator 110 records memory behavior information. As mentioned earlier, the memory behavior information includes memory lifetime information and memory size information for each tensor, and the memory behavior recording mode lasts for at least three iterations. The duration of the memory behavior recording mode corresponds to the first training phase described above. In memory allocation scheme mode, memory allocator 110 allocates memory based on the memory allocation scheme generated by offline analysis component 200. The duration of the memory allocation scheme mode corresponds to the second training phase described above, for example, restarting a complete training process.

[0064] In order to perform memory behavior recording and memory allocation according to memory scheme, memory allocator 110 may include a dynamic allocator 111 capable of allocating memory in memory behavior recording mode and a static allocator 112 capable of allocating memory based on memory allocation scheme in memory allocation mode.

[0065] In the memory behavior recording state, the memory allocator 110 defaults to selecting the dynamic allocator 111 for memory allocation and records memory behavior information. When the model's memory usage approaches the GPU's limit, the memory allocator 110 can use the native memory allocation and deallocation interface provided by the underlying computing platform (e.g., GPU) to allocate memory and record memory behavior information. In one embodiment, the dynamic allocator 111 can adopt an implementation consistent with the memory allocator of the deep learning framework used for model training, and dynamically allocate memory blocks according to memory requests at runtime.

[0066] When allocating memory according to a memory allocation scheme, memory allocator 110 reads the memory allocation scheme and initializes the memory space of static allocator 112, then issues memory requests. Memory allocator 110 can allocate a complete memory space to static allocator 112 according to the memory allocation scheme. Static allocator 112 reads the pointer and size of each static memory request in the memory allocation scheme. When it receives a memory request, it returns the memory pointer in its memory space according to the memory allocation scheme. In one implementation example, static allocator 112 executes cudaMalloc() once during initialization to request a complete memory space to accommodate all the memory required for parameter tensors, optimizer tensors, and activation value tensors. Parameter tensors and optimizer tensors are allocated to contiguous locations in this memory space according to the memory allocation scheme (e.g., the second memory allocation sub-scheme mentioned above).

[0067] To optimize the training of the MoE model, the memory allocator 110 may also include a semi-dynamic allocator 113 to handle dynamic memory allocation during MoE model training. The memory allocator 110 dispatches memory requests from dynamic layers in the MoE model to the semi-dynamic allocator 113. Upon receiving a memory request, the semi-dynamic allocator first searches for available free memory blocks in the static allocator 112. If no available free memory blocks are found, the memory request is forwarded to the dynamic allocator 111. It should be understood that dynamic layers constitute a small proportion of the MoE model; therefore, during the training process, memory allocation is primarily handled by the static allocator 112 according to the memory allocation scheme.

[0068] In the illustrated implementation, runtime component 100 further includes an error detection and recovery module 120. The error detection and recovery module 120 is used to detect whether the current video memory request conforms to the video memory allocation scheme in the video memory allocation mode. The video memory allocator 110 can, in response to detecting that the current video memory request does not conform to the video memory allocation scheme, perform video memory allocation for the current video memory request by the dynamic allocator.

[0069] The error detection and recovery module 120 is responsible for checking the correctness of the static allocator 112 during the allocation process to avoid memory allocation errors. When an error is detected, it can switch to using the dynamic allocator 111 to perform memory allocation behavior that meets the user's expectations. For example, during dense model training, theoretically, the memory behavior of all subsequent iterations from the second iteration onwards should be the same as the second iteration. In this case, if the error detection and recovery module 120 detects that the current memory request does not conform to the memory allocation scheme, it considers an error detected and allocates the memory request to the dynamic allocator 111 for processing, but still uses the static allocator 112 to handle subsequent memory requests.

[0070] The offline analysis component 200 is responsible for analyzing the memory behavior during model training and generating an efficient memory allocation scheme. The offline analysis component 200 includes a memory behavior analysis module 210 for classifying tensors based on the memory behavior information; and a memory scheme generation module 220 for generating a memory allocation scheme based on the classification results.

[0071] The memory behavior analysis module 210 can analyze the lifecycle characteristics of tensors, such as tensor memory lifecycle information, based on the model training memory behavior recorded in the memory behavior recording mode of the memory allocator 110. Based on these lifecycle characteristics, parameter tensors, optimizer tensors, and activation tensors in model training can be accurately identified. Parameter tensors and optimizer tensors are resident tensors in memory and correspond to the aforementioned non-spatially aggregateable tensors. Activation tensors, due to their short lifecycle, can be used for memory reuse and therefore correspond to spatially aggregateable tensors. When virtual vector merging is required on activation tensors propagating across forward and backward passes to reduce the computational complexity of subsequent intra-group spatial aggregation, activation tensors can be further classified into mergeable tensors (e.g., mergeable tensors within the same micro-batch are merged into a single virtual tensor) and non-mergeable tensors (corresponding to the aforementioned non-virtual tensors) based on their lifecycle characteristics. The memory behavior analysis module 210 can output a memory analysis report based on the tensor classification results.

[0072] The memory allocation module 220 generates an efficient memory allocation scheme based on the memory analysis report. This module can first tightly arrange the requests for resident memory such as parameter tensors and optimizer tensors (corresponding to the generation of the aforementioned resident memory allocation sub-scheme), and then obtain the memory arrangement of activation value tensors through the spatial aggregation algorithm described earlier in this disclosure (corresponding to the generation of the aforementioned spatial aggregation memory allocation sub-scheme), and finally generate a complete memory allocation scheme for subsequent model training.

[0073] In the illustrated example, the offline analysis component 200 also includes an abnormal memory detection module 230, which diagnoses and debugs abnormal situations during memory usage based on the memory analysis report containing tensor classification information generated by the memory behavior analysis module 210. The abnormal memory detection module 230 can detect errors in the memory analysis report to determine whether there are any anomalies in the current model training, thereby identifying potential problems in the training framework, such as memory leaks or unexpected recomputation configurations. Subsequently, memory debugging can be performed to eliminate the detected errors.

[0074] The following will combine Figure 4 This describes a specific implementation example of the memory optimization device of this disclosure when performing a training method. Figure 4 The present invention illustrates a specific application process of a memory optimization method according to an embodiment of the present disclosure in the training process of a deep learning model.

[0075] As shown in the figure, when starting the model training task, the GPU memory tool is loaded. At this time, the GPU memory tool that needs to be loaded corresponds to the runtime components of this disclosure, such as those described above. Figure 3The runtime component 100 in the aforementioned memory optimization device can be loaded using the memory management interface of a deep learning framework. For example, when the underlying runtime environment is based on the CUDA architecture and a training task is built based on the PyTorch framework, the torch.cuda.memory.CUDAPluggableAllocator interface can be used to load the memory allocator of this disclosure (e.g., ...). Figure 3 The memory allocator 110 in the runtime component 100 shown is loaded into the system, and the native PyTorch memory allocator is replaced by the memory allocator 110 through the torch.cuda.memory.change_current_allocator interface. This allows the memory tool to be seamlessly embedded into the existing PyTorch training process and fully leverage its performance optimization capabilities.

[0076] During the initial training of the model, the memory allocator 110 enters its "memory behavior recording mode" and performs a short-iteration training task. The short iteration period can be 3 to 5 training iterations, meaning the first training phase corresponds to the first N iterations after model training begins, where N is 3, 4, or 5. After the short-iteration training is completed, the memory allocator 110 saves the memory request records from the training process. These records can include entries for memory allocation and deallocation. Memory allocation corresponds to memory size information and allocation nodes, while deallocation corresponds to deallocation nodes. If a tensor is not deallocated by the end of the memory behavior recording mode, its deallocation node is considered to be after the end of the mode.

[0077] The memory request record is then output to the offline analysis component disclosed herein, for example, in combination with the above. Figure 3 The described memory optimization device includes an offline analysis component 200. First, a memory behavior analysis module analyzes memory request records to output a memory analysis report. An abnormal memory detection module then detects errors in the memory analysis report and generates an anomaly detection report. Based on the memory analysis report and the anomaly detection report, a memory allocation scheme generation module is invoked to generate a memory allocation scheme for subsequent training of the model.

[0078] Subsequently, the memory allocator disclosed herein can be set to "memory allocation mode" and used for model training in conjunction with the memory allocation scheme obtained in the previous step. This model training should occur during the aforementioned second training phase, preferably corresponding to a completely restarted training process. During this training phase, error detection and recovery modules included in the runtime components (e.g., ...) can be utilized. Figure 3The error detection and recovery module 120 in the system will assign the memory request to the dynamic allocator when it detects that the current memory request does not conform to the memory allocation scheme, so as to avoid memory trampling behavior.

[0079] The deep learning model training method and memory optimization device according to this disclosure have been described in detail above with reference to the accompanying drawings. This disclosure utilizes prior knowledge of predictable allocation patterns during model training to pre-plan memory allocation, effectively reducing memory fragmentation without affecting training performance, thereby improving resource utilization efficiency during training. The efficiency of memory allocation scheme generation is improved by spatial aggregation of non-resident tensors grouped by size and tensor merging before aggregation. Furthermore, by introducing a static-dynamic memory allocator, the memory optimization scheme of this disclosure can achieve efficient memory allocation on both dense models and MoE models.

[0080] Furthermore, the method according to the present invention can also be implemented as a computer program or computer program product, which includes computer program code instructions for performing the steps defined in the above-described method of the present invention.

[0081] Alternatively, the present invention can also be implemented as a non-transitory machine-readable storage medium (or computer-readable storage medium, or machine-readable storage medium) storing executable code (or computer program, or computer instruction code) thereon, which, when executed by a processor of an electronic device (or computing device, server, etc.), causes the processor to perform the various steps of the method described above according to the present invention.

[0082] Those skilled in the art will also understand that the various exemplary logic blocks, modules, circuits, and algorithm steps described in connection with the disclosure herein can be implemented as electronic hardware, computer software, or a combination of both.

[0083] The flowcharts and block diagrams in the accompanying drawings illustrate the architecture, functionality, and operation of possible implementations of systems and methods according to various embodiments of the present invention. In this regard, each block in a flowchart or block diagram may represent a module, segment, or portion of code containing one or more executable instructions for implementing a specified logical function. It should also be noted that in some alternative implementations, the functions marked in the blocks may occur in a different order than those marked in the drawings. For example, two consecutive blocks may actually be executed substantially in parallel, and they may sometimes be executed in reverse order, depending on the functions involved. It should also be noted that each block in the block diagrams and / or flowcharts, and combinations of blocks in the block diagrams and / or flowcharts, can be implemented using a dedicated hardware-based system that performs the specified function or operation, or using a combination of dedicated hardware and computer instructions.

[0084] The various embodiments of the present invention have been described above. These descriptions are exemplary and not exhaustive, nor are they limited to the disclosed embodiments. Many modifications and variations will be apparent to those skilled in the art without departing from the scope and spirit of the described embodiments. The terminology used herein is chosen to best explain the principles, practical application, or improvement of the technology in the market, or to enable others skilled in the art to understand the embodiments disclosed herein.

Claims

1. A deep learning model training method, comprising: In the first training phase of the deep learning model, memory behavior information is acquired, which includes memory lifetime information and memory size information of each tensor in the first training phase. The first training phase includes at least three iteration cycles. A memory allocation scheme is generated based on the memory behavior information, wherein tensor classification is performed based on the memory lifecycle information, and the memory allocation scheme is generated based on the tensor classification results and the memory size information; and In the second training phase of the deep learning model, memory allocation is performed based on the aforementioned memory allocation scheme. The generation of a video memory allocation scheme based on the video memory behavior information includes: Tensors whose memory lifetime is less than the first lifetime threshold are classified as spatially aggregateable tensors. Group the spatially aggregateable tensors according to the aforementioned video memory size information; and Spatial aggregation is performed on multiple spatially aggregable tensors with non-overlapping memory lifetimes within the same group to generate a first memory allocation sub-scheme for memory reuse of spatially aggregable tensors. Tensors whose memory lifetime is less than the first lifetime threshold are classified as spatially aggregateable tensors, including: Multiple activation value tensors whose memory lifetime is less than the first lifetime threshold, belong to the same micro-batch, and have allocated memory during the forward propagation of the same micro-batch and released memory during the backward propagation of the same micro-batch are merged into a single virtual tensor. The memory size of the virtual tensor is equal to the sum of the memory sizes of the individual activation value tensors. Tensors whose memory lifetime is less than the first lifetime threshold and whose memory is allocated and freed within the forward propagation of a micro-batch and / or whose activation value tensors are allocated and freed within the backward propagation of a micro-batch are defined as non-virtual tensors. Grouping spatially aggregateable tensors based on the aforementioned video memory size information includes: Combining multiple virtual tensors and multiple non-virtual tensors to form a tensor set; and Based on the video memory size information of each tensor in the tensor set, the tensors in the tensor set are grouped, wherein each virtual tensor participates in the grouping process as a tensor in the tensor set.

2. The method as described in claim 1, wherein, Generating a video memory allocation scheme based on the aforementioned video memory behavior information includes: Tensors whose memory lifetime exceeds the second lifetime threshold are classified as non-space-aggregateable tensors; and Based on the memory size information of the non-space-aggregateable tensor, a second memory allocation sub-scheme is generated to allocate contiguous memory to the non-space-aggregateable tensor.

3. The method as described in claim 1, wherein, The first memory allocation sub-scheme, which involves spatially aggregating multiple spatially aggregable tensors with non-overlapping memory lifetimes within the same group to generate memory reuse for spatially aggregable tensors, includes: Sort the multiple groups that have undergone spatial aggregation in descending order; Following the sorted order, the video memory is arranged sequentially for each group. The video memory arrangement operation for the current group includes: Find an unoccupied memory region in the allocated video memory that can accommodate the current group; In response to finding an unoccupied space area that meets the requirements, the current group is placed into that unoccupied space area; In response to the failure to find an unoccupied memory area that meets the requirements, the current group is placed in the unallocated video memory space.

4. The method of claim 1, wherein, Memory allocation based on the aforementioned memory allocation scheme includes: The static allocator allocates video memory to the current video memory request based on the aforementioned video memory allocation scheme; and In response to the detection that the current video memory request does not conform to the video memory allocation scheme, the dynamic allocator performs video memory allocation for the current video memory request.

5. The method of claim 1, wherein, The memory lifetime information of the tensor is determined based on the memory allocation and deallocation nodes of the tensor, and the memory allocation by the static allocator based on the memory allocation scheme includes: Allocate a complete block of video memory space based on the aforementioned video memory allocation scheme; The static allocator reads the pointers and sizes of each static memory request in the memory allocation scheme; and In response to receiving a memory request with corresponding memory request node and memory size information, a pointer corresponding to the memory request in the memory space is returned.

6. The method of claim 1, wherein, The deep learning model is an expert hybrid MoE model including dynamic layers, and the memory allocation based on the memory allocation scheme includes: In response to receiving the video memory request from the dynamic layer, the static allocator allocates video memory based on the video memory allocation scheme; In response to the static allocator not finding available video memory, the dynamic allocator allocates video memory for the video memory request of the dynamic layer.

7. The method of claim 1, wherein, The first training phase corresponds to the first N training iterations before the model starts training, where N is an integer not less than 3, and the second training phase corresponds to a complete training process where the model restarts.

8. The method of claim 7, wherein the memory allocation scheme includes a first iteration cycle memory allocation scheme and a second iteration cycle memory allocation scheme, and in the second training phase of the deep learning model, memory allocation based on the memory allocation scheme includes: In the first iteration of the restarted complete training process, memory allocation is performed based on the memory allocation scheme of the first iteration. as well as In each iteration following the first iteration of the restarted complete training process, memory allocation is performed based on the memory allocation scheme of the second iteration.

9. A deep learning model training method, comprising: Record the memory behavior information of the deep learning model for the first N training iterations. The memory behavior information includes the memory lifecycle information and memory size information of each tensor in the first N training iterations, where N is an integer not less than 3. A memory allocation scheme is generated based on the recorded memory behavior information; and In the last M training iterations of the deep learning model, memory allocation is performed based on the memory allocation scheme, where M is an integer greater than N, and the last M training iterations correspond to: M training iterations following the first N training iterations in a complete training process of the deep learning model, or M training iterations in a newly started complete training process. The process of generating a video memory allocation scheme based on the recorded video memory behavior information includes: Tensors are classified into parameter tensors, optimizer tensors, and activation value tensors based on memory lifetime information. A resident memory allocation sub-scheme is generated based on the memory size information of the parameter tensor and the optimizer tensor. Based on the memory size information of the activation value tensors, the activation value tensors are grouped, and multiple activation value tensors with non-overlapping lifetimes within the same group are spatially aggregated to generate a spatially aggregated memory allocation sub-scheme; and The resident memory allocation sub-scheme and the space aggregation memory allocation sub-scheme are combined to generate the memory allocation scheme. Specifically, multiple activation value tensors whose video memory lifetime is less than a first lifetime threshold, belong to the same micro-batch, and request video memory during the forward propagation of the same micro-batch and release video memory during the backward propagation of the same micro-batch are merged into a single virtual tensor. The video memory size of this virtual tensor is equal to the sum of the individual video memory sizes of the multiple activation value tensors. Tensors whose memory lifetime is less than the first lifetime threshold and whose memory is allocated and freed within the forward propagation of a micro-batch and / or whose activation value tensors are allocated and freed within the backward propagation of a micro-batch are defined as non-virtual tensors. Multiple virtual tensors and multiple non-virtual tensors are combined to form an activation value tensor set. Based on the video memory size information of each tensor in the activation value tensor set, the activation value tensors in the activation value tensor set are grouped.

10. A memory optimization device for training deep learning models, comprising a runtime component and an offline analysis component, wherein... The runtime component includes: A memory allocator is used to record memory behavior information in a memory behavior recording mode. This memory behavior information includes memory lifetime information and memory size information for each tensor. The memory behavior recording mode lasts for at least three iterations. The offline analysis component includes: The video memory behavior analysis module is used to classify tensors based on the video memory lifecycle information; The video memory allocation module is used to generate a video memory allocation scheme based on the tensor classification results and the video memory size information. Furthermore, the memory allocator is also used to allocate memory based on the memory allocation scheme in memory allocation mode. The generation of a video memory allocation scheme based on the tensor classification results and the video memory size information includes: Tensors whose memory lifetime is less than the first lifetime threshold are classified as spatially aggregateable tensors. Group the spatially aggregateable tensors according to the aforementioned video memory size information; and Spatial aggregation is performed on multiple spatially aggregable tensors with non-overlapping memory lifetimes within the same group to generate a first memory allocation sub-scheme for memory reuse of spatially aggregable tensors. Tensors whose memory lifetime is less than the first lifetime threshold are classified as spatially aggregateable tensors, including: Multiple activation value tensors whose memory lifetime is less than the first lifetime threshold, belong to the same micro-batch, and have allocated memory during the forward propagation of the same micro-batch and released memory during the backward propagation of the same micro-batch are merged into a single virtual tensor. The memory size of the virtual tensor is equal to the sum of the memory sizes of the individual activation value tensors. Tensors whose memory lifetime is less than the first lifetime threshold and whose memory is allocated and freed within the forward propagation of a micro-batch and / or whose activation value tensors are allocated and freed within the backward propagation of a micro-batch are defined as non-virtual tensors. Grouping spatially aggregateable tensors based on the aforementioned video memory size information includes: Combining multiple virtual tensors and multiple non-virtual tensors to form a tensor set; and Based on the memory size information of each tensor in the tensor set, the tensors in the tensor set are grouped.

11. The memory optimization apparatus of claim 10, wherein, The video memory allocator includes: A dynamic allocator is used to allocate video memory in the video memory behavior recording mode; and A static allocator is used to allocate video memory based on the video memory allocation scheme in the video memory allocation mode.

12. The memory optimization device as described in claim 11, wherein, The video memory allocator also includes: A semi-dynamic allocator is configured to, when the deep learning model is an expert hybrid model including dynamic layers, in response to receiving a memory request from the dynamic layer, cause a static allocator to allocate memory based on the memory allocation scheme, and in response to the static allocator not finding available memory, cause a dynamic allocator to allocate memory for the memory request from the dynamic layer.

13. The memory optimization device as described in claim 11, wherein, The runtime component also includes: The error detection and recovery module is used to detect whether the current video memory request conforms to the video memory allocation scheme in the video memory allocation mode. The video memory allocator is also used to allocate video memory for the current video memory request in response to the detection that the current video memory request does not conform to the video memory allocation scheme.

14. The memory optimization device as described in claim 10, wherein, The offline analysis component also includes: The abnormal memory detection module is used to diagnose and debug abnormal situations during memory usage based on the memory analysis report containing tensor classification information generated by the memory behavior analysis module. Furthermore, the video memory scheme generation module generates a video memory allocation scheme based on the video memory analysis report obtained through abnormal video memory detection.

Images