Video memory scheduling system based on dynamic sequence length assembly line parallel training

Abstract

The invention belongs to the technical field of large-scale deep learning model training, and particularly relates to a video memory scheduling system based on dynamic sequence length assembly line parallel training. The system comprises a dynamic recalculation module, a video memory arrangement and prefetching module and a mapping reconstruction module, the system collects the length and load information of each micro-batch before each round of iteration, establishes an optimization model with the goal of minimizing the iteration time, dynamically generates a re-calculation plan, and realizes inter-stage video memory load balancing; by introducing a collating stream and an asynchronous scheduler outside a main computing stream, overlapping execution of video memory collating and computing prefetching is realized, and the influence of fragment collating on training performance is reduced; a dynamic physical block mapping strategy is designed based on a CUDA virtual memory management mechanism, the physical block granularity is adaptively adjusted according to the idle state of the video memory, and the API calling and context switching overhead is reduced. Experimental results show that the video memory utilization rate and the system stability are remarkably improved, and an efficient video memory scheduling scheme is provided for large model training.

Description

Technical Field

[0001] This invention belongs to the field of large-scale deep learning model training technology, specifically relating to a memory scheduling system based on dynamic sequence length pipeline parallel training. Background Technology

[0002] Variable-length sequence training refers to the process in which the length of the input sequence of the model varies significantly among samples during deep learning training. This is common in tasks such as natural language processing and speech recognition. Compared with traditional deep neural network (DNN) training with fixed-length input, variable-length sequence training brings stronger task adaptability, but also raises many challenges in terms of computational efficiency and system resource utilization. In recent years, Transformer models (such as T5, GPT, etc.) have widely adopted mixed task datasets for multi-task training. This training method is inherently accompanied by the diversity of input length. For example, in the text summarization task, the average length of the input sequence is as high as 977.93 tokens[1], while in the text implication task, it is only 51.59 tokens[2]. This huge length difference makes the traditional padding-based training method inefficient, because a large number of padding operations introduce redundant computation, which not only wastes valuable computational resources, but also reduces the overall training throughput. In order to alleviate the resource waste problem caused by variable-length sequences, techniques such as packing reduce the computational overhead by reducing the invalid padding of samples. These methods have achieved remarkable results in a single iteration, but most of them only focus on optimization within each training round, while ignoring the uneven use of system-level resources caused by fluctuations in input length between rounds.

[0003] In distributed training scenarios, this problem is further amplified. Current mainstream distributed training frameworks, such as Megatron-LM and DeepSpeed, generally adopt a static parallel strategy, that is, setting the batch size and other settings before training begins [3]. This approach often configures system resources based on the maximum possible sequence length, thereby ensuring that the model can run even in extreme cases. However, for a large number of short and medium sequence samples in actual iterations, this conservative strategy not only leads to a large amount of idle memory and computing power, but also causes a continuous decrease in system throughput during long-term training, significantly reducing the overall training efficiency. However, if only the normal input length is considered for resource allocation, performance degradation or even training interruption is likely to occur due to insufficient memory when encountering occasional ultra-long sequences.

[0004] The memory scheduling problem in variable-length sequence training faces two main challenges:

[0005] (1) The dynamic change in the length of the input sequence will cause the memory load to fluctuate greatly during the training process. The memory utilization of the GPU is inconsistent in different iterations, which makes some devices risk being overloaded or even crashing, while other devices are in a low utilization state for a long time, resulting in serious waste of resources. [4]

[0006] (2) Frequent changes in video memory allocation can cause video memory fragmentation, especially during long training periods, when video memory space is repeatedly requested and released, resulting in discontinuous memory space. This can lead to video memory allocation failure when requesting large blocks of video memory, further threatening the stability of training. Summary of the Invention

[0007] The purpose of this invention is to propose a high-efficiency and stable memory scheduling system based on dynamic sequence length pipelined parallel training, so as to solve the problems of memory load fluctuation and severe fragmentation caused by variable-length sequence input in the current variable-length sequence pipelined parallel training process.

[0008] This invention proposes a memory scheduling system based on dynamic sequence length pipelined parallel training. First, before each training iteration, the length and load information of each micro-batch are collected to establish an optimization model aimed at minimizing iteration time, dynamically generating a recomputation plan to achieve memory load balancing between stages. Second, by introducing an independent defragmentation stream and an asynchronous scheduler outside the main computation stream, overlapping execution of memory defragmentation and computation prefetching is achieved, reducing the impact of defragmentation on training performance. Finally, a dynamic physical block mapping strategy is designed based on the CUDA virtual memory management mechanism, adaptively adjusting the physical block granularity according to the memory idle state to reduce API call and context switching overhead. Specifically, it includes three parts: a dynamic recomputation module, a memory defragmentation and computation prefetching overlapping module, and an efficient memory mapping reconstruction module, as detailed below:

[0009] (i) The dynamic recomputation module aims to reduce peak memory usage while maintaining training throughput, thereby improving memory resource utilization and scheduling flexibility. Specifically, before each training iteration, the system first collects the length, position, and load information of the current micro-batch, assessing the memory usage and potential idle time of each micro-batch in this training iteration. Subsequently, combining the memory release benefits and computational overhead of recomputation, the system dynamically generates a recomputation plan for each micro-batch, selecting whether to release intermediate activations and postpone recomputation at certain stages. This module emphasizes global awareness of resource usage at each stage of the pipeline, prioritizing recomputation in areas with lighter GPU loads or idle time in the pipeline stages, thereby achieving a significant reduction in peak memory usage with minimal performance sacrifice. This strategy is more flexible than the traditional static recomputation mechanism and can significantly improve system stability and resource utilization efficiency under variable input lengths.

[0010] (II) The memory defragmentation and computation prefetching overlap module addresses the increasingly severe memory fragmentation problem under frequent scheduling of variable-length sequence training and dynamic recomputation, making memory defragmentation overhead a key factor hindering training performance. To address this, this module introduces an overlap mechanism for memory defragmentation and computation prefetching. Specifically, during training, the system predicts the memory size required for subsequent computations based on existing scheduling information and opens another stream in the current computation phase to perform memory defragmentation. Through pre-defragmentation, the system can interleave tasks during computation gaps, thereby minimizing the blocking impact of defragmentation on the main computation flow. Furthermore, the prefetching mechanism can pre-map the defragmented memory space to virtual addresses and register available regions for upper-layer programs, improving the response speed of the next round of memory allocation. This "computation-defragmentation" overlap strategy significantly reduces the impact of defragmentation latency on training throughput.

[0011] (III) Regarding the memory mapping reconstruction module, in traditional memory management mechanisms, to improve the consistency and simplicity of virtual address and physical address mapping, the system typically uses fixed-granularity physical blocks as the basic unit of mapping. This type of mechanism generally relies on the API provided by CUDA Virtual Memory Management (VMM) to request several fixed-size physical blocks (e.g., 2MB), and strictly bases all mapping and remapping operations on this fixed granularity. While this design ensures simple mapping logic and easy debugging in static scenarios, it faces significant performance bottlenecks in large-scale deep learning model training scenarios. Especially during periods of frequent dynamic task scheduling and active memory remapping mechanisms, such as when processing dynamic sequence length inputs, performing dynamic recalculation switching, or reusing gradients across layers, the problems caused by the fixed-granularity mapping method become increasingly prominent. Specifically, the main bottlenecks include two aspects: First, the number of API calls required for each remapping operation increases linearly with the number of physical pages. When the number of physical blocks is large, frequent calls to system interfaces will lead to a sharp increase in call overhead. Second, each API call is accompanied by GPU driver communication and context state synchronization, which can easily trigger a large number of context switches, thereby further increasing latency.

[0012] To address the aforementioned issues, this invention proposes a memory mapping reconstruction mechanism based on dynamic physical block combination and constructs a memory mapping reconstruction module. It breaks through the constraints of traditional fixed-length page blocks, enabling flexible selection of physical blocks of different sizes for mapping based on actual memory needs, thereby reducing API call counts and significantly improving mapping efficiency. The basic approach is as follows: when a memory allocation request is triggered, the system allocator first accurately estimates the target tensor size and selects the most suitable granularity for allocation from the candidate physical block pool. For large tensors (typically exceeding several hundred MB), the system prioritizes coarse-grained physical page blocks (e.g., 64MB, 256MB) to reduce the number of mapping and allocation calls; while for small tensors (several MB or less), small-grained physical page blocks (e.g., 2MB, 4MB) are joined together to ensure allocation flexibility and fine-grained memory management capabilities. Through this differentiated granularity strategy, the system can maintain mapping accuracy and flexibility while effectively reducing the number of physical blocks, thus alleviating the performance overhead caused by frequent remapping.

[0013] Furthermore:

[0014] (i) The dynamic recomputation module specifically includes: firstly, defining which part of the model each GPU should be responsible for executing, thus achieving memory load balancing on the GPUs during pipelined training. The L-layer Transformer model will be divided into S stages. Furthermore, each mini-batch is also divided into M micro-batches when executed in parallel on the pipeline. .use Description of micro-batch In the stage Forward propagation execution time, Description of micro-batch In the stage The backpropagation execution time. Furthermore, using... To describe micro-batch In the stage The recalculation scheme, Description of micro-batch In the stage The total activation value size. The recalculation scheme determines the backpropagation time. and actual activation value usage Before training begins, a series of simulations are performed. Performance analysis is used to obtain the forward propagation time for micro-batches of different sequence lengths, the backpropagation time under different recomputation strategies, and memory usage. A related cost model is also established for subsequent simulations of pipelined computation. Specifically:

[0015] Restrictions:

[0016] (1) The phase partitioning consists of disjoint subsets, i.e.:

[0017] ;

[0018] ;

[0019] (2) The memory usage for model training cannot exceed the upper limit.

[0020] Stages can be obtained After calculating the micro-batch Forward propagation and Memory usage of activation values ​​during backpropagation:

[0021] ;

[0022] The memory constraint can then be expressed as:

[0023] ;

[0024] Where M is the GPU's video memory capacity, Weight is the model weight size, and Buffer is the size of various buffers used for auxiliary calculations and communication during training. These are usually determined before training begins and can be considered fixed values. Therefore, we only need to control the video memory usage of the activation values.

[0025] (3) Data dependency of pipeline parallelism.

[0026] Micro-batch during pipeline training In the stage Only after calculating the forward propagation can we proceed to the next stage. To carry out subsequent forward propagation; in the stage Only after calculating backpropagation can we proceed to the next stage. Then proceed with the subsequent forward propagation. Furthermore, both the forward and backward propagation of the micro-batch are calculated sequentially within the same phase. Forward propagation must have occurred earlier than Forward propagation; its backward propagation must also precede that. Backpropagation.

[0027] The critical path in the pipeline training process consists of the stepwise computation of the longest-running microbatches in the initial and final stages, and the computation of all microbatches in the final stage of the steady phase (see [link to documentation]). Figure 1 The upper bound of the time used is:

[0028] ;

[0029] We need to minimize the upper bound of the critical path, so the optimization problem becomes:

[0030] ;

[0031] ;

[0032] ;

[0033] After applying a batch-level recomputation strategy, the recomputation strategy for some micro-batches in certain stages can be adjusted to make fuller use of video memory. For example... Figure 2 As shown, in stages 1 and 2, micro-batches 0-2 undergo recalculation, while micro-batches 3 do not. In stage 3, micro-batches 0-1 undergo recalculation, while micro-batches 2-3 do not. Stage 4, however, does not undergo recalculation at all. This method is relatively... Figure 1 Stage 1, which has high memory pressure, will not significantly increase memory pressure, while Stage 4, which has low memory pressure, will reduce computation time, thereby improving overall training efficiency.

[0034] (ii) The aforementioned memory reclaiming and computation prefetching overlap module, by introducing an independent reclaiming stream and an asynchronous scheduler into the computation flow, executes memory reclamation and pre-allocation operations in parallel with the main computation flow. This ensures that computation is not blocked while waiting for memory reclaiming, thereby reducing the impact on overall throughput. Specifically:

[0035] During the scheduling phase, the system collects the input dimensions, pipeline scheduling strategies, and recomputation strategies from all micro-batches in the current iteration. The collected information undergoes internal prediction processing to generate estimates of the memory block size and lifetime required for subsequent computation stages. These prediction results serve as input to the sorting and prefetching tasks, driving the scheduler to arrange subsequent operations.

[0036] While the main computation stream executes, the memory defragmentation stream continuously acquires prefetching tasks based on prediction information and asynchronously defragments fragmented memory. The prefetch defragmentation operation includes the following four stages:

[0037] (1) Obtain the prefetch task to be executed from the "video memory prefetch queue" and attempt to acquire the memory pool lock. At this time, check the main thread priority flag variable need_lock. If the flag is true, release the lock immediately to yield the resource; otherwise, continue to execute the prefetch task and enter the subsequent steps;

[0038] (2) Collect free video memory block information from the ordered free pool and calculate the total amount. If it is insufficient to meet the prefetching requirements, request new physical blocks from the driver layer. Then search for suitable virtual blocks in the free virtual video memory pool. If no matching size is found, it is also necessary to request new physical blocks from the driver layer.

[0039] (3) Start the timer and enter the video memory mapping loop: map the fragmented physical blocks to the virtual video memory area one by one, and set the access permissions. If the mapping loop is completed successfully, proceed directly to the next step; if the timer reaches the threshold ∆t, actively release the lock, and decide whether to enter sleep mode and wait for the main thread to wake up, or reacquire the lock and continue the mapping task according to the need_lock status;

[0040] (4) Record the number of sleep cycles during the mapping process and calculate the sleep ratio after actively releasing the lock. If the ratio exceeds the threshold p, shorten the timer threshold ∆t; otherwise, appropriately extend the timer threshold. After completing the tuning, put the mapped virtual memory block into the prefetch pool.

[0041] The scheduler plays a crucial coordinating role throughout the process. It monitors the execution status of the computation flow through hook mechanisms and attempts to insert prefetch tasks during gaps in the main computation flow (such as data-parallel gradient updates or pipelined bubble operations). The scheduler maintains a task priority strategy: when the main computation flow contains large-scale matrix computations or communication tasks, the scheduler lowers the priority of the scrambling task, allowing only low-overhead release operations; when memory is scarce and large tensors are about to be needed, the scheduler raises the priority of the scrambling task, accelerating merging and migration.

[0042] To support the scheduler's real-time decision-making, the module introduces a background thread pool. This thread pool manages the task queue and triggers memory compaction operations via an asynchronous API. Upon completion of the compaction operation, it returns a pointer to available memory and its corresponding size to the scheduler. These pointers are written to a shared metadata table and can be directly queried upon the next allocation request. In this way, the main computation stream can obtain available memory blocks instantly without interacting with the memory compaction process.

[0043] In addition, the module defines three levels for the compaction operation: fast release (releasing only invalid tensors), lightweight merge (merging free blocks), and deep reordering (releasing memory to the operating system and re-requesting compaction). The scheduler automatically selects the appropriate level based on the predicted time window. For example, when a communication delay exceeding several hundred microseconds is detected, deep reordering is allowed to generate a larger contiguous block of memory for use in subsequent stages.

[0044] The prefetch overlap module achieves overlap between computation and memory management through multi-stream parallelism, asynchronous scheduling, and collaboration with a background thread pool. Its implementation details ensure that the memory defragmentation process is almost transparent to the main computation stream, and can continuously provide efficient and stable defragmentation support in training scenarios with high memory pressure and severe fragmentation.

[0045] (III) In the memory mapping reconstruction module, the specific operation of memory mapping reconstruction is to divide the physical block pool into two categories: a large pool and a small pool, to support allocation requests for large and small tensors respectively. Physical blocks in the large pool are kept intact as much as possible when not necessary to avoid frequent splitting leading to fragmentation accumulation. When there are no large blocks to satisfy the request, the system selects suitable blocks from the small pool and completes the allocation through splitting or splicing. Furthermore, to avoid excessively long memory allocation paths and allocation request blocking, block merging and large-scale defragmentation operations do not directly interfere with foreground allocation requests, but are handled by a background daemon thread that performs defragmentation gradually in a lazy-merging manner. The system also continuously collects memory allocation and release data during the training iteration phase and dynamically optimizes the physical block pool structure based on statistical results. For example, when some large blocks in the large pool are idle for a long time, they can be split and transferred to the small pool; while blocks frequently used for splicing in the small pool will be directly merged into large blocks for efficient reuse in subsequent iterations. By combining hierarchical pool management with dynamic optimization, the system achieves higher memory utilization while ensuring low-latency allocation, effectively reducing fragmentation rate, and significantly mitigating the performance overhead caused by frequent mapping and remapping.

[0046] The main technical features and functional advantages of this invention are as follows:

[0047] (1) Dynamic recomputation and elastic memory scheduling: Before each iteration, the length and memory usage information of each micro-batch are collected to dynamically generate a recomputation plan. The system selectively releases intermediate activations and postpones recomputation based on the stage load, which significantly reduces the peak memory usage. Under the premise of ensuring throughput, the system achieves global awareness and optimal allocation of resources at each stage of the pipeline, which is more flexible and adaptable than the static recomputation mechanism. (2) Memory defragmentation and prefetching overlap mechanism: To address the memory fragmentation problem caused by frequent scheduling, a scheme for parallel execution of memory defragmentation and computation prefetching is proposed. During the computation stage, the system uses idle streams to perform memory defragmentation and address mapping prefetching in advance, effectively reducing the blockage of the main computation process by defragmentation delay. Through the "computation-defragmentation" overlap execution strategy, the efficiency of memory defragmentation and allocation response speed are greatly improved. (3) Efficient mapping reconstruction mechanism: To reduce the overhead of memory remapping, a dynamic physical block mapping strategy based on CUDA VMM is designed. The system intelligently combines physical page blocks of different sizes based on the current idle state and allocation needs, completing the mapping operation with minimal API calls and reducing context switching and defragmentation overhead. This mechanism maintains a unified virtual interface while achieving adaptive optimization of the underlying mapping granularity, significantly improving the memory management efficiency and system stability of large-scale models under high-fluctuation loads.

[0048] Results Display

[0049] To verify the effectiveness of the memory scheduling system based on dynamic sequence length pipelined parallel training proposed in this invention, a small-scale verification experiment was designed and conducted. The experiment was conducted on one A100-PCIe server with a total of eight NVIDIA A100 GPUs interconnected via PCIe. The experiment used the Megatron-LM dataset to train a GPT-13B model on the Flan-V2 dataset. The parallel strategy was a hybrid approach combining 4-way pipelined parallelism and 2-way data parallelism. The baseline scheme employed the most aggressive recomputation method while ensuring no OutOfMemoryError (OOM) and PyTorch's native memory scheduling scheme. The optimized scheme was the memory scheduling system proposed in this invention.

[0050] Figure 4 The comparison of training throughput between the present invention and the baseline system under different batch sizes and sequence lengths is shown in the figures. As can be seen, the present invention significantly improves throughput. For example, with a batch size of 131072, the throughput of the baseline system Megatron is 24280.5 tokens / s, while the throughput of the present invention reaches 35106.4 tokens / s, an improvement of 44.6%. Furthermore, for a sequence length of 4096, the present invention can increase the throughput from 15221.2 tokens / s of the baseline to 22581.8 tokens / s (an improvement of 48.4%). Simultaneously, the performance improvement consistently exceeds 10% across all tested batch sizes and sequence lengths, and the improvement increases progressively with increasing batch size and sequence length, indicating that the present invention performs better in large-scale training scenarios. This significant performance improvement is attributed to our optimized communication strategy, computation scheduling mechanism, and memory management and computation pipeline mechanism, effectively solving communication bottlenecks, resource utilization efficiency issues, and memory bottlenecks in large-scale training. Attached Figure Description

[0051] Figure 1 This is a pipelined parallel mechanism for model training.

[0052] Figure 2 This is an example of a hybrid recalculation scheme.

[0053] Figure 3 It is a memory manager architecture based on dynamic sequence length pipelined parallel training.

[0054] Figure 4 This represents the throughput of the present invention and the baseline method during the training process. Detailed Implementation

[0055] Step 1: Dynamic Recomputation Planning. Before each iteration, the system collects information on the length, position, and memory load of the current microbatch, evaluating the memory usage and potential free space of each microbatch at each stage of the pipeline. Based on the analysis results, the module establishes an optimization model with minimizing iteration time as the objective function and memory usage as the constraint. The system linearly models the forward propagation, backward propagation, and recomputation overhead of each pipeline stage, and estimates the computational load of each stage using profiling data. Subsequently, a recomputation plan for each microbatch is generated using a heuristic algorithm, prioritizing recomputation to pipeline gap stages. This strategy dynamically reduces peak memory usage without significantly affecting training performance, achieving a balance of memory load between stages.

[0056] Step 2: Memory Defragmentation. The system introduces an independent CUDA Stream outside the main compute stream. Before execution, the compute stream acquires the compute graph and sends its memory requirements to the defragmentation stream. The defragmentation stream includes the following four steps: (a) retrieving tasks from the memory prefetch queue and attempting to acquire locks; (b) collecting information on free memory blocks and searching for matching blocks in the virtual memory pool; (c) executing a mapping loop between physical and virtual blocks and dynamically adjusting timer thresholds based on the running status to balance defragmentation efficiency with competition from the main thread; (d) registering defragmented memory blocks to the prefetch pool for use in the next round. The scheduler automatically selects the defragmentation level based on the time window and communication latency. When a long communication gap is detected, deep rearrangement is triggered to generate contiguous large blocks of memory, thereby significantly reducing the fragmentation rate. Through the overlapping execution strategy of "compute-defragmentation," the system effectively reduces the impact of memory defragmentation on training throughput and improves memory allocation response speed.

[0057] Step 3: Memory Mapping. The system uses the CUDA Virtual Memory Management (VMM) mechanism to allocate physical blocks of different sizes for mapping virtual memory. During each remapping process, based on the current free memory status and the size to be allocated, the system intelligently selects physical blocks of different granularities for combination mapping to reduce mapping time. Simultaneously, the system dynamically adjusts the physical block pool structure based on memory usage characteristics during training, further improving mapping flexibility and memory utilization.

[0058] Step 4: Dynamic Recalculation Execution Module. A recalculation strategy control mechanism is introduced at the model level. During each forward propagation, instead of relying on predefined fixed parameters, the recalculation strategy dynamically determines whether to perform a recalculation operation. This enables control over the recalculation scheme.

[0059] References:

[0060] [1]HERMANN K M, KOČISKÝ T, GREFENSTETTE E, et al. Teaching machines to read and comprehend[C] / / Advances in Neural Information Processing Systems 28 - Proceedings of the 2015 Conference: volume 2015. 2015: 1693 - 1701.

[0061] [2]WILLIAMS A, NANGIA N, BOWMAN S R. A Broad - Coverage Challenge Corpus for Sentence Understanding through Inference[J / OL]. 2017. DOI:10.48550 / arxiv.1704.05426.

[0062] [3]NARAYANAN D, SHOEYBI M, CASPER J, et al. Efficient large - scale language model training on GPU clusters using megatron - LM[C / OL] / / Proceedings of SC 2021: The International Conference for High Performance Computing, Networking, Storage and Analysis: Science and Beyond. New York, NY, USA: ACM, 2021: 1 - 15. DOI:10.1145 / 3458817.3476209.

[0063] [4]JIANG C, JIA Z, ZHENG S, et al. DynaPipe: Optimizing Multi-task Training through Dynamic Pipelines[C / OL] / / EuroSys 2024 - Proceedings of the 2024 European Conference on Computer Systems. New York, NY, USA: ACM, 2024: 542-559. DOI:10.1145 / 3627703.3629585。

Claims

1. A dynamic sequence length pipeline parallel training based memory scheduling system, characterized in that, First, the length and load information of each mini-batch is collected before each iteration of the training process, an optimization model is established to minimize iteration time, a dynamic recalculation plan is generated, and inter-stage memory load balancing is achieved; second, by introducing an independent arrangement stream and asynchronous scheduler outside the main computing stream, overlapping execution of memory arrangement and calculation prefetching is achieved, and the impact of fragment arrangement on training performance is reduced; finally, based on the CUDA virtual memory management mechanism, a dynamic physical block mapping strategy is designed, the physical block granularity is adaptively adjusted according to the memory idle state, and the API call and context switching overhead is reduced; specifically, the dynamic recalculation module, the memory arrangement and prefetching overlap module, and the efficient mapping reconstruction module are included, which are three parts, as follows: The dynamic recalculation module aims to reduce peak memory occupancy while ensuring training throughput, thereby improving memory resource utilization and scheduling flexibility; specifically, before the start of each iteration of the training, the length, position and load information of the current mini-batch is first collected, and the memory occupancy and potential idle of each mini-batch in the current iteration is evaluated; Subsequently, combined with the memory release benefit and calculation overhead of recalculation, a recalculation plan is dynamically generated for each mini-batch, and the release of intermediate activations and the delay of recalculation in some stages are selected; this module has global awareness of resource usage in each stage of the pipeline, and preferentially arranges recalculation in parts of the pipeline where the GPU load is lighter or there is idle, thereby achieving a substantial reduction in memory peak with as little performance sacrifice as possible; The memory arrangement and calculation prefetching overlap module, during the training process, the system predicts the memory size required for subsequent calculation based on the existing scheduling information in advance, and opens another stream to perform memory fragment arrangement operation in the current calculation stage; Through pre-arrangement, the system inserts tasks in the calculation gap, thereby minimizing the blocking effect of the arrangement operation on the main calculation process; In addition, the prefetching mechanism maps the arranged memory space to a virtual address in advance, and registers the available area for the upper program, thereby improving the response speed of the next round of memory allocation; The memory mapping reconstruction module, when the memory allocation request is triggered, the system allocator first accurately estimates the size of the target tensor, and selects the most suitable granularity from the candidate physical block pool for allocation; for large-scale tensors, the system preferentially selects coarse-grained physical page blocks to reduce the number of mapping and allocation calls; For small-scale tensors, small-grained physical page blocks are used for splicing to ensure allocation flexibility and memory fine management capability; through this differentiated granularity strategy, the system can maintain the accuracy and flexibility of mapping, and effectively reduce the number of physical blocks, thereby alleviating the performance overhead caused by frequent remapping.

2. The video memory scheduling system of claim 1, wherein, The dynamic recalculation module specifically includes: Firstly, define which part of the model each GPU should be responsible for executing, to achieve memory load balancing of pipeline training on GPU; divide the L-layer Transformer model into S stages ; each mini-batch is also divided into M micro-batches when performing pipeline parallelism ; describe the micro-batch forward propagation execution time in stage , describe the micro-batch backward propagation execution time in stage ; use to describe the micro-batch recalculation scheme in stage , describe the micro-batch total activation value size in stage ; the recalculation scheme determines the backward propagation time and the actual activation value occupancy ; Before the start of training, a series of simulation calculations are performed, the forward propagation time of different sequence length mini-batches and the backward propagation time and memory occupancy under different recalculation strategies are obtained through performance analysis, and a related cost model is established for subsequent simulation of pipeline calculation conditions; wherein: The restrictions are: (1) The phase partitioning consists of disjoint subsets, i.e.: ； ； (2) The memory usage for model training cannot exceed the upper limit; Stage of obtaining After calculating the micro-batch Forward propagation and Memory usage of activation values ​​during backpropagation: ； The memory constraint is then expressed as: ； Where M is the GPU's video memory capacity; Weight is the model weight size; and Buffer is the size of various buffers used for auxiliary calculations and communication during training. These are usually determined before training begins and are considered fixed values. Therefore, it is only necessary to control the video memory usage of the activation values. (3) Data dependencies in pipeline parallelism; That is, in the pipeline training process, the micro-batch After the forward propagation is calculated, the subsequent forward propagation can be performed in the phase After the backward propagation is calculated, the subsequent forward propagation can be performed in the phase ; in the phase , the forward propagation of the micro-batch must be earlier than the forward propagation of the micro-batch ; and the backward propagation of the micro-batch must be earlier than the backward propagation of the micro-batch . The critical path in the pipeline training process consists of the step-by-step computation of the longest-running microbatches in the initial and final stages, and the computation of all microbatches in the last stage of the steady phase. Therefore, the upper bound of the computation time is: ； We need to minimize the upper bound of the critical path, so the optimization problem becomes: ； ； ； After applying a batch-level recalculation strategy, the recalculation strategy for some micro-batches in certain stages is adjusted to make fuller use of video memory.

3. The video memory scheduling system of claim 2, wherein, The memory reclamation and computation prefetching overlap module introduces an independent reclamation stream and an asynchronous scheduler into the computation flow, executing memory reclamation and pre-allocation operations in parallel with the main computation flow. This ensures that computation is not blocked while waiting for memory reclamation, thereby reducing the impact on overall throughput. Specifically: During the scheduling phase, the system collects the input dimensions, pipeline scheduling strategies, and recomputation strategies in the dynamic recomputation module for all micro-batches in the current iteration. The collected information is processed internally to generate estimates of the memory block size and lifecycle required for subsequent computation stages. The prediction results serve as input for the sorting and prefetching tasks, driving the scheduler to arrange subsequent operations. While the main computation stream is executing, the memory defragmentation stream continuously acquires prefetching tasks based on prediction information and asynchronously defragments fragmented memory. The prefetch defragmentation operation includes the following four stages: (1) Obtain the prefetch task to be executed from the "video memory prefetch queue" and attempt to acquire the memory pool lock; at this time, check the main thread priority flag variable need_lock. If the flag is true, release the lock immediately to give up the resource; otherwise, continue to execute the prefetch task and enter the subsequent steps. (2) Collect free video memory block information from the ordered free pool and calculate the total amount. If it is insufficient to meet the prefetching requirements, request a new physical block from the driver layer. Then search for a suitable virtual block in the free virtual video memory pool. If no matching size is found, request a new physical block from the driver layer. (3) Start the timer and enter the video memory mapping loop: map the fragmented physical blocks to the virtual video memory area one by one and set the access permissions; if the mapping loop is completed successfully, proceed directly to the next stage; if the timer reaches the threshold ∆t, actively release the lock and decide whether to enter sleep mode and wait for the main thread to wake up, or reacquire the lock and continue the mapping task according to the need_lock status; (4) Record the number of sleeps during the mapping process and calculate the sleep ratio after actively releasing the lock; if the ratio exceeds the threshold p, shorten the timer threshold ∆t; otherwise, extend the timer threshold; after completing the tuning, put the mapped virtual memory block into the prefetch pool. The scheduler monitors the execution status of the computation flow through a hook mechanism and tries to insert prefetch tasks during the execution gaps of the main computation flow. The scheduler maintains a task priority strategy: when the main computation stream has large-scale matrix computation or communication tasks, the priority of the sorting stream is reduced, and only low-overhead release operations are allowed; when memory is scarce and large tensors are about to be needed, the priority of the sorting task is increased to speed up merging and migration.

4. The video memory scheduling system of claim 3, wherein, In the memory defragmentation and prefetching overlap module, a background thread pool is introduced to support the scheduler's real-time decision-making. This thread pool is responsible for managing the defragmentation task queue and triggering memory defragmentation operations through an asynchronous API. When the defragmentation operation is completed, it returns an available pointer and its corresponding size to the scheduler. These pointers are written to a shared metadata table and can be directly queried when the next allocation request arrives. In this way, the main computing stream can obtain available memory blocks instantly without interacting with the memory defragmentation process. In addition, the module defines three levels for the compaction operation: fast release, which only releases invalid tensors; lightweight merge, which merges free blocks; and deep rearrangement, which releases memory to the operating system and re-requests compaction. The scheduler automatically selects the appropriate level based on the predicted time window, including allowing deep rearrangement to be triggered when a communication delay of more than several hundred microseconds is detected, in order to generate a larger contiguous memory block for use in subsequent stages.

5. The video memory scheduling system of claim 4, wherein, The specific operations of memory mapping reconstruction in the memory mapping reconstruction module are as follows: The physical block pool is divided into two categories: a large pool and a small pool, to support allocation requests for large and small tensors respectively. Physical blocks in the large pool are kept as intact as possible to avoid fragmentation caused by frequent splitting. When there are no large block resources to satisfy the request, the system selects suitable blocks from the small pool and completes the allocation by splitting or splicing. In addition, to avoid excessively long memory allocation paths and blocked allocation requests, block merging and large-scale defragmentation operations do not directly interfere with foreground allocation requests, but are instead handled by background daemon threads that perform defragmentation gradually in a lazy merging manner. The system also continuously collects memory allocation and release data during the training iteration phase, and dynamically optimizes the physical block pool structure based on statistical results; This includes cutting up large sections of material in a large pool and transferring them to smaller pools when those sections are idle for an extended period. Blocks frequently used for splicing in the small pool are directly merged into larger blocks for efficient reuse in subsequent iterations; By combining hierarchical pool management with dynamic optimization, the system achieves higher memory utilization while ensuring low-latency allocation, effectively reducing fragmentation rate, and significantly mitigating the performance overhead caused by frequent mapping and remapping.

Images