Model training method, device and program product based on heterogeneous storage offloading

Abstract

This application provides a model training method, device, program product, and storage medium based on heterogeneous storage offloading. During the backpropagation computation of the AI ​​accelerator, a new model parameter update task is added to the task pool after each layer of the model's gradient calculation is completed. While the AI ​​accelerator performs backpropagation computation, the processor consumes parameter update tasks from the task pool in parallel. Specifically, before the backpropagation reaches the first layer, model parameter update tasks are scheduled in descending order of model layer number. After the gradient calculation of the first layer is completed, the remaining model parameter update tasks can be scheduled in ascending order of model layer number. This scheduling logic ensures that the shallowest parameters most needed are updated and ready before the next iteration begins, reducing processor waiting time caused by parameter update lag and maximizing the parallel efficiency of the CPU and AI accelerator.

Description

Technical Field

[0001] This application relates to the field of artificial intelligence technology, and more specifically, to a model training method, device, program product, and storage medium based on heterogeneous storage offloading. Background Technology

[0002] With the rapid development of artificial intelligence technology, various AI models, such as Large Language Models (LLM), have been widely applied in various fields. As the scale of model parameters explodes, the storage space required for training far exceeds the memory capacity limit of AI accelerators, posing a severe challenge to training large models with limited hardware resources. To address the memory bottleneck problem of AI accelerators, an offloading technology based on heterogeneous storage layers can be adopted. This technology constructs a multi-level storage architecture consisting of AI accelerator memory, host (such as CPU) memory, and external persistent storage (such as NVMe SSD), dynamically moving tensor data that is not currently involved in computation out of the AI ​​accelerator memory.

[0003] In scenarios based on heterogeneous storage hierarchical offloading, since model parameters, optimizer states, and gradient layer state tensors are usually stored in host memory or external persistent storage media, in order to reduce transmission time, during model training, AI accelerators usually perform forward propagation and backward propagation calculations to obtain gradient data for each layer, and then the host (i.e., CPU) updates the model parameters based on the gradient data.

[0004] In response to the above scenarios, during model training, related technologies either require gradient calculation and model parameter updates to be executed sequentially, or although they can be executed in parallel, the scheduling strategy fails to prioritize based on the deterministic requirements of the next iteration. This results in the processor frequently experiencing severe idle waiting during iteration transitions due to unready model parameters. Overall, there are numerous scheduling bubbles, failing to achieve deep overlap between computation and I / O operations, and the overall training efficiency needs further improvement. Summary of the Invention

[0005] In view of this, this application provides a model training method, device, program product and storage medium based on heterogeneous storage offloading.

[0006] According to a first aspect of this application, a model training method based on heterogeneous storage offloading is provided. The method is applicable to a processing node used to perform a model training task, the processing node including a processor and an AI accelerator, the method being executed by the processor, and the method comprising: The AI ​​accelerator is used to perform the current round of backpropagation calculation. After the AI ​​accelerator calculates the gradient of each layer of the model, the gradient is cached in the memory of the AI ​​accelerator. A parameter update task for updating the model parameters of that layer is added to the pre-built task pool. During the current round of backpropagation computation performed by the AI ​​accelerator, the following operations are performed in parallel: Obtain and execute parameter update tasks from the task pool to update the model parameters of the corresponding layer using the calculated gradients; Before the AI ​​accelerator completes the gradient calculation of the first layer of the model, the parameter update tasks in the task pool are scheduled and executed in descending order of the number of model layers. After the AI ​​accelerates the gradient calculation of the first layer of the model, the remaining parameter update tasks in the task pool are scheduled and executed in order of increasing model layer number, so that the model parameters that are needed first in the next round of forward propagation calculation are updated first.

[0007] According to a second aspect of this application, a computer program product is provided, the computer program product comprising a computer program that, when executed, implements the method mentioned in the first aspect above.

[0008] According to a third aspect of this application, an electronic device is provided, the electronic device including a processor, a memory, and a computer program stored in the memory that is executable by the processor, wherein the processor executes the computer program to implement the method mentioned in the first aspect above.

[0009] According to a fourth aspect of this application, a computer-readable storage medium is provided, on which a computer program is stored, which, when executed, implements the method mentioned in the first aspect above.

[0010] By applying the scheme provided in this application, a model parameter update task pool can be constructed. During the backpropagation computation, each time the AI ​​accelerator completes the gradient calculation for a model layer, a new model parameter update task is added to the task pool. While the AI ​​accelerator performs backpropagation computation, the processor consumes parameter update tasks from the task pool in parallel. This scheme breaks the limitation of traditional methods that update parameters in a fixed backpropagation order (from deep to shallow), introducing a phased dynamic scheduling strategy: before backpropagation reaches the first layer, model parameter update tasks are scheduled in descending order of model layer number; after the gradient calculation for the first layer is completed, the remaining model parameter update tasks can be scheduled in ascending order of model layer number. This scheduling logic precisely matches the computational dependencies of the model, ensuring that the shallowest parameters most needed are updated and ready before the next iteration begins, reducing processor waiting time caused by parameter update lag, maximizing the parallel efficiency of the CPU and AI accelerator, and improving the overall throughput of model training by timely releasing GPU memory space to optimize the utilization of heterogeneous storage resources.

[0011] It should be understood that the above general description and the following detailed description are exemplary and explanatory only, and do not limit this application. Attached Figure Description

[0012] To more clearly illustrate the technical solutions in the embodiments of this application, the accompanying drawings used in the description of the embodiments will be briefly introduced below. Obviously, the accompanying drawings described below are only some embodiments of this application. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0013] Figure 1 This is a schematic diagram illustrating the sequential execution of gradient calculation and model parameter update in related technologies.

[0014] Figure 2 This is a schematic diagram illustrating the parallel execution of gradient calculation and model parameter update in related technologies.

[0015] Figure 3 This is a flowchart of a model training method according to an exemplary embodiment of this application.

[0016] Figure 4 This is a schematic diagram of a model parameter update task in a task pool dynamically scheduled based on the next round of forward propagation requirements, which is an exemplary embodiment of this application.

[0017] Figure 5 This is a schematic diagram of CPU and GPU parallel processing in an exemplary embodiment of this application.

[0018] Figure 6This is a schematic diagram of an exemplary embodiment of the present application, illustrating the pruning gradient based on an optimistic pruning machine.

[0019] Figure 7 This is a schematic diagram comparing the differences in the batch data execution order during model training between existing technologies and this application.

[0020] Figure 8 This is a schematic diagram of the logical structure of an electronic device according to an embodiment of this application. Detailed Implementation

[0021] The technical solutions of the embodiments of this application will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of this application, and not all embodiments. Based on the embodiments of this application, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of this application.

[0022] With the rapid development of artificial intelligence technology, various AI models, such as Large Language Models (LLM), have been widely applied in various fields. However, with the explosive growth in the scale of model parameters, the memory pressure on AI accelerators (e.g., GPUs) executing model training tasks has become a core bottleneck restricting large-scale parallel computing. During training iterations, the system needs to maintain massive tensor states, mainly including model parameters, gradients, optimizer states, and intermediate activation values ​​generated by feedforward computation. As the number of model parameters increases from tens of billions to hundreds of billions, the storage space required for training far exceeds the memory capacity limit of AI accelerators, posing a severe challenge to training large models with limited hardware resources.

[0023] To address the memory bottleneck in AI accelerators, a heterogeneous storage-level offloading technology can be employed. This technology dynamically moves tensor data that is not currently involved in computation out of the AI ​​accelerator's memory by constructing a multi-level storage architecture consisting of AI accelerator memory, host (e.g., CPU) memory, and external persistent storage (e.g., NVMe SSDs). Current heterogeneous storage offloading technologies mainly fall into two categories: The first type is a two-level offloading method based on "AI accelerator memory - host memory". This method maintains the model's state tensor in host memory and offloads model parameters, optimizer states, or gradients that are not involved in the current layer's computation to host memory via the PCIe bus. This method increases the scale of model parameters that a node can support by expanding the physical storage boundary.

[0024] The second category is a three-tiered offloading method based on an "AI accelerator memory - host memory - persistent storage mechanism." This method further extends the storage boundary to external persistent media such as NVMe SSDs. Its key feature is utilizing the massive storage space of hard drives to hold hundreds of billions or even trillions of model parameters and optimizer states. When computation reaches the corresponding neural network layer, the required tensor data is retrieved back to the AI ​​accelerator memory or host memory through asynchronous I / O or multi-disk concurrent scheduling mechanisms. This method aims to achieve the deployment and training of large-scale models in general-purpose hardware environments through deep utilization of hardware storage layers.

[0025] In scenarios based on heterogeneous storage offloading, since model parameters, optimizer states, and gradient layer state tensors are typically stored in host memory or external persistent storage, to reduce transmission time, during model training, AI accelerators usually perform forward and backward propagation calculations to obtain gradient data for each layer. The host (i.e., the CPU) then updates the model parameters based on this gradient data. For example, in scenarios where model parameters are stored in external persistent storage, the AI ​​accelerator can complete the current round of forward and backward computations to obtain the gradients for each layer. Then, the CPU retrieves the model parameters for each layer from the external persistent storage and performs the model parameter update operation. Related model training schemes suffer from numerous synchronization barriers and scheduling bubbles, preventing deep overlap between computation and I / O (i.e., reading data from external persistent storage).

[0026] During model training, some scenarios require a global synchronization mechanism to ensure training stability. A typical application is gradient clipping based on the global norm. After backpropagation calculates the gradients of each layer, directly using them for parameter updates can lead to uncontrolled parameter update step sizes when the gradient magnitude is too large (i.e., "gradient explosion"), causing the model to skip the optimal solution. Gradient clipping constrains the gradient magnitude through a preset threshold, avoiding excessively large single parameter update magnitudes and ensuring model convergence stability. Norm-based clipping is a common method, treating all parameter gradients as a global vector and calculating the L2 norm. If the norm exceeds the threshold, the gradient vector is scaled proportionally while maintaining the gradient direction. Because this clipping method relies on global gradient norm calculation, the CPU can only start the model parameter update task after the AI ​​accelerator has completed gradient calculations for all layers.

[0027] For example, such as Figure 1The diagrams illustrate the execution order of gradient calculation and model parameter update in the relevant technologies. Taking the scenario where model parameters are stored on an external storage device as an example, the CPU only begins to read model parameters from the external storage device (I / O) and update the model parameters after the GPU has completed the forward propagation (FWD) and backward propagation (BWD) calculations for all layers. Clearly, the CPU spends a significant amount of time idling.

[0028] Of course, some scenarios may not require the aforementioned global synchronization mechanism. For example, if gradient pruning is not needed, the GPU and CPU can achieve partial overlap to some extent. Figure 2 As shown in the figure, taking the scenario where model parameters are stored on an external storage device as an example, after the GPU completes the forward propagation (FWD) and backward propagation (BWD) calculations for some layers of the model and obtains the gradients of some layers, the CPU begins to execute the model parameter updates for these layers, that is, reading the model parameters of these layers from the external storage device (I / O) and updating the model parameters of these layers.

[0029] For scenarios where a global synchronization mechanism is not required (e.g., training tasks that do not require gradient clipping), a certain degree of parallel processing can be achieved between the GPU and CPU. Figure 2 As shown in the diagram, taking the scenario where model parameters are stored on external storage as an example, after the GPU completes the forward propagation (FWD) and backward propagation (BWD) calculations for some layers of the model and obtains the gradients for those layers, the CPU begins to execute the model parameter updates for these layers. This involves reading the model parameters from the external storage device (I / O) and updating them. This allows for partial overlap between gradient calculation and parameter updates. However, the order of model parameter updates typically follows the inherent order of backward propagation (from deep to shallow layers), failing to prioritize based on the dependencies of the next round of forward propagation calculations. This can lead to shallow layer parameters not being updated in a timely manner, causing CPU waiting. Furthermore, the interaction with external storage still relies on traditional I / O mechanisms, resulting in significant context switching overhead and failing to achieve deep overlap between computation and I / O operations, thus limiting further improvements in overall training efficiency.

[0030] In summary, for scenarios based on heterogeneous storage offloading, the relevant technologies either require gradient calculation and model parameter updates to be executed serially during model training, or although they can be executed in parallel, the scheduling strategy fails to prioritize based on the deterministic requirements of the next iteration. This results in the processor frequently experiencing severe idle waiting during iteration transitions due to unready model parameters. Overall, there are many scheduling bubbles, and deep overlap between computation and I / O operations is not achieved, so the overall training efficiency needs to be improved.

[0031] Based on this, this application provides a model training method based on heterogeneous storage offloading. It constructs a model parameter update task pool. During the backpropagation computation, the AI ​​accelerator adds a new model parameter update task to the task pool after completing the gradient calculation for each layer of the model. While the AI ​​accelerator performs backpropagation computation, the processor consumes parameter update tasks from the task pool in parallel. By analyzing the deterministic needs during the training process, the traditional first-in-first-out (FIFO) order is broken, and a demand-aware scheduling algorithm is adopted to prioritize updating and backpropagating the model parameters with the earliest demand time in the next round of forward propagation from the task pool. This proactive scheduling method ensures that key parameters can initiate I / O transmission as soon as their gradients are ready, minimizing the idle time of the processor waiting for model parameter backpropagation.

[0032] The model training method of this application embodiment is applicable to processing nodes that perform model training tasks. These nodes can be constructed using a heterogeneous computing architecture consisting of general-purpose processors (such as CPUs) and dedicated AI accelerators (such as GPUs, TPUs, or NPUs). In some scenarios, the processing node also includes external storage devices, such as SSDs, independent of the aforementioned processors and AI accelerators. The processor is primarily responsible for scheduling and control flow management of the training tasks, while the AI ​​accelerator focuses on high-throughput matrix operations.

[0033] The models mentioned in the embodiments of this application can be of various types, such as models with deep network structures, such as convolutional neural networks (CNN), recurrent neural networks (RNN), and the mainstream Transformer architecture in recent years. These models involve a large amount of gradient calculation and parameter updates during training.

[0034] The following combination Figure 3 , Figure 4 and Figure 5 The model training method of the embodiments of this application will be described, such as... Figure 3 As shown, the method may include the following steps: S302. Utilize the AI ​​accelerator to perform the current round of backpropagation calculation, and after the AI ​​accelerator calculates the gradient of each layer of the model, cache the gradient in the memory of the AI ​​accelerator, and add a parameter update task to the pre-built task pool to update the parameters of the model at that layer. In step S302, in the current round of model iteration, after completing the forward propagation calculation, the processing node can initiate the backpropagation calculation process. For example, an AI accelerator (such as a GPU or NPU) can be used to execute the current round of backpropagation calculation. This calculation process is triggered layer by layer. When the gradient calculation of a certain layer (e.g., layer L) is completed, the AI ​​accelerator can cache the gradient data in its local on-chip memory. Simultaneously, the AI ​​accelerator can send a lightweight task generation signal to the host processor (such as a CPU) to trigger the model parameter update task generation logic. Figure 4 As shown, after receiving this signal, the processor can register and insert a corresponding "model parameter update task" into the pre-built "task pool". This task is essentially a structure containing the gradient data memory address, references to the corresponding model layer parameters, and metadata for update operations (such as optimizer calculations), used to update the model parameters of that layer. By caching the gradients of each layer, a foundation can be provided for the dynamic scheduling of subsequent tasks.

[0035] S304. During the current round of backpropagation calculation performed by the AI ​​accelerator, the following operations are performed in parallel: Parameter update tasks are retrieved from the task pool and executed to update the model parameters of the corresponding layer using the calculated gradients; wherein, before the AI ​​accelerator completes the gradient calculation of the first layer of the model, the parameter update tasks in the task pool are scheduled and executed sequentially in descending order of model layer number; after the AI ​​accelerator completes the gradient calculation of the first layer of the model, for the remaining parameter update tasks in the task pool, the parameter update tasks in the task pool are scheduled and executed sequentially in ascending order of model layer number, so that the model parameters that are needed first in the next round of forward propagation calculation are updated first.

[0036] In step S304, the processor no longer passively waits for the backpropagation to complete. Instead, it works in parallel with the AI ​​accelerator (such as a GPU), consuming parameter update tasks generated by the accelerator in real time from the task pool. For example, the processor can maintain one or more worker threads that continuously poll or monitor the status of the task pool. When a new task is detected, the worker thread can acquire the task and perform parameter update operations, such as calling the step function of the optimizer (such as SGD or Adam) to update the model parameters of the corresponding layer using gradients cached in the AI ​​accelerator's memory.

[0037] To minimize processor idling time, this embodiment employs a two-stage, out-of-order dynamic scheduling strategy. The entire scheduling logic is cleverly divided into two stages, with the switching critical point being whether the AI ​​accelerator has completed the gradient calculation for the first layer of the model. In the first stage, i.e., during the window period when backpropagation has just begun and the first layer gradient has not yet been calculated, the processor follows a scheduling order of "model layer number from largest to smallest." That is, the processor will prioritize selecting and executing tasks with larger model layers (e.g., layers L, L-1, and L-2) from the task pool.

[0038] Once the AI ​​accelerator has completed the gradient calculations for the first layer and added the corresponding model parameter update tasks to the task pool, it can switch to the second stage. At this point, such as Figure 4 As shown, the processor can schedule the remaining, unexecuted model parameter update tasks (usually concentrated in the shallow layers of the model) in the task pool according to the order of model layer number from smallest to largest. Considering the strict dependencies in the forward propagation computation of deep neural networks: the next iteration of forward propagation must start from the first layer of the model, and the output of each layer is the input of the next layer. If the shallow layer parameter updates lag, when a new round of forward propagation starts, the computation flow will be forced to block when it reaches the layer with unupdated parameters, waiting for the model parameter update thread to complete its work. This will lead to idle computing resources and delays. By prioritizing the updating of shallow layer parameters in the second phase, it can be ensured that when the current round of backpropagation ends and the next round of forward propagation is about to begin, the most urgently needed model parameters have been updated first.

[0039] For example, after completing the current round of forward propagation calculation, the GPU begins the current round of backward propagation calculation, calculating the gradients from layer L to layer 1 in sequence and caching them. For instance, after completing the gradient calculation for layer L, the CPU starts using the gradient of layer L to update the model parameters of layer L, and then uses the gradient of layer L-1 to update the model parameters of layer L-1. When the GPU has calculated the gradient of layer 1, for example, if the CPU is currently updating the model parameters of layer 5, after updating the model parameters of layer 5, the CPU will not continue to update the model parameters of layer 4, but will first update the model parameters of layer 1 so that the GPU can use the updated model parameters of layer 1 for the next round of forward propagation calculation. After completing the gradient calculation of layer 1, the model parameters are updated in the order of layer 2, layer 3, and layer 4.

[0040] This "look-ahead" update strategy effectively hides the parameter update delay at the end of the backpropagation stage, making the pipeline of forward and backward propagation more tightly connected. This can reduce the synchronization waiting time between the processor and the accelerator and improve the overall training throughput.

[0041] like Figure 5 The diagram illustrates the execution sequence in an exemplary embodiment of this application. Taking a scenario where model parameters are stored on an external storage device as an example, after the GPU completes the forward propagation (FWD) and backward propagation (BWD) calculations for some layers of the model and obtains the gradients for those layers, the CPU begins updating the model parameters for those layers. This involves reading the model parameters from the external storage device (I / O) and updating them. By scheduling the execution order of the model parameter update tasks based on the needs of the next round of forward propagation calculations, it ensures that the model parameters most needed are updated first. This allows the forward computation (F) and model parameter updates to be executed in parallel even in the next round, maximizing the parallel efficiency of the GPU and CPU.

[0042] To achieve this complex two-stage sorting, the task pool's data structure cannot be a simple first-in-first-out (FIFO) queue. Instead, it can be designed as a priority queue or an unordered list in conjunction with a specific scheduler.

[0043] By prioritizing the updating and backhaul of the underlying parameters with the earliest required time points in the next round of forward propagation from the task pool, this proactive scheduling method ensures that critical parameters can initiate I / O transmission as soon as their gradients are ready, minimizing the idle time of the processor waiting for underlying data backhaul.

[0044] In the model training process, related technologies either update model parameters immediately after gradient calculation without caching the gradients, which, due to the lack of gradient storage, forces model parameter updates to strictly follow the backpropagation order and prevents dynamic scheduling of these tasks; or they allocate separate storage spaces for model parameters and gradients, doubling the memory usage of AI accelerators and limiting the size of trainable models and batch sizes. To address these technical bottlenecks, some embodiments employ an in-situ memory reuse mechanism based on lifecycle mutual exclusion. Specifically, during the current round of forward propagation calculation using the AI ​​accelerator, corresponding storage space can be pre-allocated in the AI ​​accelerator's memory for each network layer to store the model parameters required for that layer. In the subsequent backpropagation phase, after the AI ​​accelerator calculates the gradients layer by layer, it does not request new AI accelerator memory regions but directly writes the calculated gradient data into the physical address where the model parameters were originally stored, thus overwriting the existing parameter data with the gradients. The core logic of this design lies in utilizing the strict mutual exclusion of parameters and gradients in the physical lifetime. That is, when the backward gradient calculation of a certain layer is completed, the original parameters of that layer are no longer dependent on subsequent calculations in the current iteration cycle. Therefore, the buffer occupied by it can be safely converted into a storage container for the output gradient of that layer.

[0045] By employing this in-situ overwrite method, this embodiment maximizes memory utilization without increasing the memory footprint of the AI ​​accelerator. Simultaneously, the strategy of caching gradient data instead of consuming it immediately breaks down the synchronization barrier between gradient calculation and parameter updates in traditional training. This expands the producer-consumer queue to full capacity, providing a foundation for dynamically prioritizing model parameter update tasks based on the needs of the next round of forward propagation. This decouples parameter update tasks from the heavy computational flow, allowing for flexible scheduling within the idle window of the entire iteration cycle. This improves the parallel efficiency between the AI ​​accelerator and the processor, as well as the overall training throughput.

[0046] To achieve in-situ reuse of model parameter storage space for gradient data and avoid the additional overhead of data movement, in some embodiments, for any layer of the model, after the AI ​​accelerator completes the gradient calculation for that layer, the address mapping table in the memory management unit can be updated. This unbinds the physical memory block that originally pointed to the model parameters of that layer from the logical address space of the parameters and rebinds it to the logical address space used to index the gradient of that layer. Then, based on the updated mapping relationship, the AI ​​accelerator can be controlled to directly execute memory write instructions to write the calculated gradient data into that physical memory block, thereby achieving in-situ overwriting of parameters by the gradient at the physical level. By modifying the page table or mapping descriptor rather than moving the data itself, the semantic transformation of the storage object is achieved, avoiding the huge overhead of traditional memcpy operations. This "pointer-level" switching consumes almost no memory bandwidth, reduces latency during backpropagation, and ensures that the gradient calculation pipeline is not interrupted. Furthermore, since the original physical blocks are reused, there is no need to reserve extra redundant space for gradients in the memory of the AI ​​accelerator, eliminating the risk of peak memory overflow caused by temporary buffer allocation. This makes it possible to train ultra-large-scale parameter models on AI accelerators with limited memory capacity, while providing transparent and efficient interface support for the subsequent asynchronous scheduler to access gradient data through a unified logical address.

[0047] In the model training process of related technologies, global gradient pruning operations typically require the AI ​​accelerator to complete gradient calculations for all layers of the model before the data is sent back to the processor. Only then can the processor aggregate and calculate the global gradient norm, execute the pruning logic, obtain the pruned gradients for each layer, and perform subsequent model parameter update operations. This synchronization mechanism requires gradient calculation and model parameter updates to be executed sequentially, introducing waiting delays, severely disrupting the continuity of pipeline parallelism, causing computing resources to be idle during communication gaps, and reducing overall training throughput. Therefore, for scenarios where gradient calculation and model parameter updates are executed in parallel, gradient pruning operations based on the global gradient norm are usually not feasible.

[0048] In the distributed training process of deep learning models, gradient explosion often leads to excessively large parameter update steps, resulting in training divergence. Traditional fixed-threshold gradient clipping methods are difficult to adapt to the dynamic changes in gradient norm at different training stages. If the threshold is set too high, it cannot effectively suppress spikes; if it is set too low, it will distort the gradient direction and hinder convergence. Furthermore, executing gradient statistics and clipping sequentially will prolong critical path latency. To solve these problems, this embodiment introduces a global gradient norm prediction and asynchronous clipping mechanism based on an exponential moving average model. Specifically, when scheduling and executing parameter update tasks in the task pool sequentially, the system can obtain the actual global gradient norm of the previous iteration and the predicted global norm determined in the previous iteration. Then, it uses a preset weight factor to weight the two and dynamically calculates the predicted global gradient norm of the current iteration using an exponential moving average algorithm. For any parameter update task to be executed, the system directly determines the clipping coefficient based on the predicted global gradient norm, clips the gradients required for the task in real time, and uses the clipped gradients to update the model parameters. By leveraging an exponential moving average model to smoothly track historical gradient statistics, the pruning threshold can adaptively follow the changing trend of the gradient norm. This achieves precise suppression of gradient spikes without needing to store complete historical gradient data, avoiding the limitations of fixed thresholds under complex loss surfaces. Furthermore, since the prediction norm of the current round can be directly derived from the historical data of the previous round, the system can predetermine the pruning coefficient without waiting for the gradient calculations of all levels in the current round. This breaks the serial dependency of "aggregating statistical norms first, then performing pruning" in the traditional process, enabling parallel execution of gradient calculation tasks and model parameter update tasks. This effectively masks the computational time consumed in the parameter update stage, improving the pipeline utilization of heterogeneous accelerators and overall training throughput.

[0049] For example, the global gradient norm of the current iteration can be determined based on formula (1): Formula (1) in, Let the norm of the predicted global gradient in the t-th iteration be denoted as . These are preset weighting factors; This represents the actual global gradient norm in the (t-1)th iteration; Let represent the global gradient norm predicted in the (t-1)th iteration.

[0050] While optimistic pruning based on prediction improves parallel efficiency, deviations between predicted and true values ​​may exist, potentially affecting model convergence accuracy. To balance training efficiency and model quality, in some embodiments, after the AI ​​accelerator completes the current round of backpropagation, the actual global gradient norm for the current iteration can be determined based on the gradients of each layer of the model cached in the AI ​​accelerator. The prediction deviation is then determined based on the predicted and actual global gradient norms. For example, after the AI ​​accelerator completes the current round of backpropagation, the actual global gradient norm can be calculated based on all cached gradient data and compared with previous predictions to determine the prediction deviation. By setting a preset threshold, the reliability of the prediction can be intelligently assessed. If the deviation is within an acceptable range, the prediction is accurate, and training continues. If the deviation exceeds the preset threshold, it indicates that inaccurate predictions may have misled parameter updates. In this case, a correction process can be triggered to correct the updated model parameters. This ensures high-speed parallelism while retaining the ability to correct for anomalies, guaranteeing the robustness of the training process.

[0051] For example, such as Figure 6 As shown, during the backpropagation computation, after the GPU completes the gradient calculation for one layer, the CPU can prune the gradient of that layer using an optimistic gradient pruning mechanism, and then use the pruned gradient to update the model parameters of that layer. Once the GPU has completed the gradient calculation for all layers, it can obtain the actual global gradient norm based on the gradients of all layers. Then, by comparing the actual global gradient norm with the predicted global gradient norm, it can determine if there are any anomalies with large prediction biases. If so, the model parameters are reloaded and updated. After the update, the updated model parameters can be directly used for the next round of forward propagation computation without needing to read the model parameters from external storage devices (i.e., no I / O). This allows for deep overlap between computation and I / O operations, improving overall training efficiency.

[0052] In training models, related technologies typically employ a "batch-by-batch, full-layer" serial processing mode. This means the AI ​​accelerator must sequentially complete the forward and backward propagation calculations for each batch of data across all layers of the model. Only after all batches' gradients have been calculated are the full gradients aggregated, and model parameter updates executed. This approach results in a significant delay in the readiness time of the full gradients, requiring the entire network's backpropagation to complete before model parameter updates can begin. This severely limits the overlap between gradient calculation and model parameter updates, causing idle computing resources during the waiting window. To address these issues, some embodiments introduce a training mechanism based on layer-by-layer gradient accumulation and execution order reconstruction. Specifically, after obtaining the sample data corresponding to the current training task, it can be divided into multiple batches. Task instructions are then sent to the AI ​​accelerator to break the traditional serial logic and adopt a "single-layer continuous processing mode." During forward propagation, the forward calculations for each layer of the model are executed continuously for these multiple batches. Only after all batches of the current layer have been calculated is the continuous forward calculation of multiple batches for the next layer triggered. After completing the forward computation of all batches at the last layer of the model, the backpropagation phase can begin. Similarly, multiple batches of backpropagation are performed consecutively for each layer of the model, and the gradients generated by each batch are accumulated and aggregated in situ in the AI ​​accelerator's memory to obtain the aggregated gradient for that layer, which is then cached. Since backpropagation follows the order from the output layer to the input layer, when the aggregated gradient of the last layer is calculated and cached, all gradients for that layer are ready. Therefore, there is no need to wait for the gradient calculations of the penultimate and preceding layers to complete; the aggregated gradient can be immediately retrieved from the AI ​​accelerator's memory, and the corresponding model parameter update task can be initiated. Simultaneously, the AI ​​accelerator executes the backpropagation and gradient aggregation of the remaining layers in parallel. This strategy, by advancing the time point when all gradients are ready from the "end of backpropagation across the entire network" to the "end of backpropagation of the first layer," expands the overlap window between the model parameter update task, the I / O operations of the underlying storage system, and the gradient calculations of the remaining layers. This provides sufficient scheduling margin for the throughput of the underlying storage system, effectively widening the overlap space between I / O and computation, and improving the overall training throughput in a heterogeneous computing environment.

[0053] In some embodiments, when adopting the aforementioned "single-layer continuous processing mode," the AI ​​accelerator can be designed to perform forward and backward propagation with multiple batches in reverse order of computation between adjacent layers of the model, i.e., using a stack-like "last-in, first-out" scheduling logic. Specifically, assume that the training data for the current training task is divided into batch 1, batch 2, and batch 3. For any layer of the model, during the forward propagation computation, the computation of that layer can be performed sequentially on batch 1, batch 2, and batch 3 to obtain their respective activation values ​​in that layer. Then, the forward computation of all batches in the next layer is performed. During the forward computation in the next layer, the processing order is reversed, i.e., the forward computation of batch 3, batch 2, and batch 1 is performed sequentially to obtain their respective activation values ​​in the next layer and cache them. The backpropagation process is similar. After completing the forward propagation calculations for all batches in the last layer, the backpropagation calculations for all batches in the last layer can be initiated. At this point, the previous execution order (e.g., the execution order of the forward propagation in the last layer) can be reversed. For example, assuming the forward propagation calculation order in the last layer is batch 3, batch 2, batch 1, then the backpropagation calculation order in the last layer is batch 1, batch 2, batch 3, obtaining the gradient values ​​for each batch in the last layer. These gradient values ​​can then be used to overwrite the activation values ​​of the corresponding batches, and the backpropagation calculations for all batches in the penultimate layer can begin. In the penultimate layer, the processing order is reversed again, i.e., batch 3, batch 2, batch 1 are reverse-calculated sequentially to obtain their respective gradient values ​​in the penultimate layer, which are then cached and used to overwrite the activation values ​​of the corresponding layers. By reversing the order, the gradient values ​​of each batch in each layer of the model can reuse the activation values ​​of that batch in place.

[0054] For example, such as Figure 7 The diagram shows a comparison of the lifecycle of the state tensor and the model parameter update time window during model training in related technologies and in the model training process of this application. This visually demonstrates the optimization strategy for micro-batch processing timing during model training. The upper part represents the "batch-by-batch full-level processing mode" in related technologies, while the lower part represents the "single-layer continuous processing mode" in this application.

[0055] In related technologies, micro-batch processing is performed sequentially, meaning that batch 2 is processed only after the forward and backward propagation calculations of batch 1 are completed. This causes model parameter updates to wait for the gradient calculations of all micro-batches to be completed before they can begin, resulting in a delayed update window that cannot overlap with gradient calculations. However, this application's embodiment employs a hierarchical aggregation and zigzag flipping strategy to process all micro-batches continuously within the same layer (e.g., the first layer processes batch 1 → batch 2, and the second layer processes batch 2 → batch 1). During backpropagation, the flipping of the micro-batch order enables in-situ reuse of activation value memory. This scheduling method allows parameter updates to begin immediately once the gradient aggregation of the last layer (e.g., the fourth layer) is complete, significantly advancing the update window and allowing it to overlap with the backpropagation calculations of the remaining layers, thereby effectively shortening the training cycle and improving memory utilization.

[0056] In some embodiments, an external storage device (such as an NVMe solid-state drive array, SSD, etc.) independent of the processor and AI accelerator can also be configured in the processing node for full storage of model parameters. When there are parameter update tasks to be executed in the task pool, a dedicated data transmission channel between the external storage device and the processor memory can be established through direct memory access technology. This bypasses the AI ​​accelerator and the complex operating system kernel protocol stack, directly loading the model parameters required by the task into the processor's memory in blocks. After the processor performs in-place update calculations on the loaded parameters using the cached gradients, it writes the updated parameters back to the external storage device through this transmission channel. Direct memory access technology enables zero-copy interaction between the storage device and the processor memory. Since the parameter loading and write-back process does not go through the AI ​​accelerator and does not require the CPU to copy byte by byte, the high-speed interconnect bandwidth of the accelerator is effectively freed up for gradient communication and calculation, while reducing the load on the host CPU. Furthermore, in the scenario of a three-level heterogeneous storage based on "AI accelerator memory - host memory - external storage", this application embodiment can increase the overlap window between I / O (such as reading model parameters from external storage) and computation by introducing a demand-aware dynamic scheduling strategy, thereby improving the overall training efficiency of the model.

[0057] The model training method of this application will be described below with reference to a specific embodiment.

[0058] The technical approach of this embodiment revolves around a demand-driven asynchronous scheduling framework. By deeply reconstructing the tensor lifecycle during large model training, it achieves maximum overlap between computation, communication, and storage I / O. The main process of this asynchronous scheduling framework is as follows: In heterogeneous storage training scenarios, by orchestrating the timing of gradient accumulation at model levels and combining it with an in-place memory reuse mechanism for parameters and gradients, the originally serial I / O transmission and computation process is transformed into a highly parallel asynchronous execution flow, thereby completely masking the I / O latency of external storage using a computation window.

[0059] This embodiment mainly includes the following core technical strategies: (1) Lifecycle-Aware Memory Reuse and Gradient Caching Design. This embodiment utilizes the mutual exclusion property of parameters and gradients in the physical lifecycle during large model training to design an in-situ memory reuse mechanism. After the system completes the forward and backward computation of the current layer, its original parameter buffer is immediately converted into a storage container for the output gradient of that layer. This design expands the producer-consumer queue to full capacity without increasing additional GPU memory usage, providing a physical basis for subsequent asynchronous scheduling. By caching gradients instead of consuming them immediately, parameter update operations can be decoupled from the current iteration computation, allowing update tasks to be flexibly scheduled within the idle window of the entire iteration cycle.

[0060] (2) Hierarchical gradient accumulation and execution order reconstruction techniques. To broaden the overlap between I / O and computation, this embodiment reconstructs the execution timing of micro-batches, adopting a single-layer continuous processing mode, that is, first completing the forward and backward computation of all micro-batches in a single layer and generating aggregated gradients. By reversing the processing order of adjacent layers (Z-shaped flow), this embodiment significantly advances the ready time of the full gradient while achieving in-situ reuse of activation values ​​and reducing I / O overhead. This strategy greatly expands the overlap window of asynchronous parameter updates, providing sufficient scheduling margin for the throughput of the underlying storage system.

[0061] (3) Demand-aware asynchronous priority update scheduling strategy. Under the premise of decoupling gradient caching, this embodiment constructs a dynamic task pool to be updated. The scheduling engine breaks the traditional first-in-first-out order by parsing the deterministic requirements of the training sequence and adopts a demand-aware scheduling algorithm: it prioritizes the underlying parameters with the earliest demand time in the next round of forward propagation from the task pool for priority update and backhaul. This forward-looking scheduling method ensures that key parameters can start I / O transmission as soon as their gradients are ready, minimizing the idle time of the computing unit due to waiting for the underlying data backhaul.

[0062] (4) Optimistic Pruning and Self-Healing Mechanism Based on Prediction Model: To address the global synchronization barrier problem introduced by gradient pruning, this embodiment designs an optimistic pruning mechanism. The system utilizes the spatiotemporal stability of the gradient norm and predicts the global gradient norm of the current iteration using an exponential moving average (EMA) model, enabling local pruning of gradients at each layer immediately after generation and entry into the asynchronous update pipeline. Furthermore, this embodiment also designs an algebraic rollback scheme. When the prediction deviation exceeds a threshold, the parameters are directly corrected by performing a reverse algebraic transformation on the optimizer state, without re-executing the computation task. This technique eliminates pipeline bubbles caused by global synchronization, ensuring extreme temporal overlap between computation and I / O tasks.

[0063] Compared with existing technologies, it has the following beneficial effects: By implementing lifecycle-aware in-situ memory reuse, gradient caching is achieved without adding extra space, enabling single-node hardware to support ultra-large-scale model training and significantly reducing training costs. Through "Z-shaped" hierarchical accumulation and demand-aware scheduling, the system amortizes storage and transfer time within the computation window, effectively reducing idle time caused by computation units waiting for data and improving training throughput. An optimistic pruning mechanism is introduced to deconstruct the global gradient synchronization constraint. Update tasks at each layer can immediately enter the asynchronous pipeline after generation, ensuring continuous utilization of system resources throughout the iteration cycle. The accompanying algebraic rollback mechanism ensures parameter self-healing with low algebraic overhead when prediction errors occur, squeezing out extremely high system concurrency performance while maintaining the rigor of the algorithm logic.

[0064] The solutions in the above embodiments can be freely combined to obtain new solutions when there is no conflict. Due to space limitations, they will not be listed one by one here.

[0065] Furthermore, embodiments of this application also provide an electronic device, such as... Figure 8 As shown, the electronic device 80 includes a processor 81, a memory 82, and computer instructions stored in the memory that can be executed by the processor 81. When the processor 81 executes the computer instructions, it implements the method described in any of the above embodiments.

[0066] Furthermore, this application also provides a computer program product, which includes a computer program that, when executed, implements the method of any of the above embodiments.

[0067] Accordingly, this application also provides a computer storage medium storing a program that, when executed by a processor, implements the method in any of the above embodiments.

[0068] The embodiments of this application may take the form of a computer program product implemented on one or more storage media (including but not limited to disk storage, CD-ROM, optical storage, etc.) containing program code. Computer-usable storage media include permanent and non-permanent, removable and non-removable media, and information storage can be implemented by any method or technology. Information may be computer-readable instructions, data structures, program modules, or other data. Examples of computer storage media include, but are not limited to: phase-change memory (PRAM), static random access memory (SRAM), dynamic random access memory (DRAM), other types of random access memory (RAM), read-only memory (ROM), electrically erasable programmable read-only memory (EEPROM), flash memory or other memory technologies, read-only optical discs, read-only memory (CD-ROM), digital versatile optical discs (DVD) or other optical storage, magnetic tape, disks or other magnetic storage devices, or any other non-transfer medium that can be used to store information accessible by a computing device.

[0069] For the device embodiments, since they basically correspond to the method embodiments, the relevant parts can be referred to in the description of the method embodiments. The device embodiments described above are merely illustrative. The units described as separate components may or may not be physically separate, and the components shown as units may or may not be physical units, that is, they may be located in one place or distributed across multiple network units. Some or all of the modules can be selected to achieve the purpose of this embodiment according to actual needs. Those skilled in the art can understand and implement this without creative effort.

[0070] The user information (including but not limited to user device information, user personal information, etc.) and data (including but not limited to data used for analysis, data stored, data displayed, etc.) involved in this application are all information and data authorized by the user or fully authorized by all parties. The collection, use and processing of the relevant data shall comply with relevant laws, regulations and standards, and corresponding operation entry points shall be provided for users to choose to authorize or refuse.

[0071] It should be noted that, in this document, relational terms such as "first" and "second" are used only to distinguish one entity or operation from another, and do not necessarily require or imply any such actual relationship or order between these entities or operations. The terms "comprising," "including," or any other variations thereof are intended to cover non-exclusive inclusion, such that a process, method, article, or apparatus that comprises a list of elements includes not only those elements but also other elements not expressly listed, or elements inherent to such a process, method, article, or apparatus. Without further limitation, an element defined by the phrase "comprising one..." does not exclude the presence of other identical elements in the process, method, article, or apparatus that includes said element.

[0072] The methods and apparatus provided in the embodiments of this application have been described in detail above. Specific examples have been used to illustrate the principles and implementation methods of this application. The descriptions of the embodiments above are only for the purpose of helping to understand the methods and core ideas of this application. At the same time, for those skilled in the art, there will be changes in the specific implementation methods and application scope based on the ideas of this application. Therefore, the content of this application should not be construed as a limitation of this application.

Claims

1. A model training method based on heterogeneous storage offloading, characterized in that, The method is applicable to processing nodes used to perform model training tasks, the processing node including a processor and an AI accelerator, the method being executed by the processor, and the method comprising: The AI ​​accelerator is used to perform the current round of backpropagation calculation. After the AI ​​accelerator calculates the gradient of each layer of the model, the gradient is cached in the memory of the AI ​​accelerator. A parameter update task for updating the model parameters of that layer is added to the pre-built task pool. During the current round of backpropagation computation performed by the AI ​​accelerator, the following operations are performed in parallel: Obtain and execute parameter update tasks from the task pool to update the model parameters of the corresponding layer using the calculated gradients; Before the AI ​​accelerator completes the gradient calculation of the first layer of the model, the parameter update tasks in the task pool are scheduled and executed in descending order of the number of model layers. After the AI ​​accelerates the gradient calculation of the first layer of the model, the remaining parameter update tasks in the task pool are scheduled and executed in order of increasing model layer number, so that the model parameters that are needed first in the next round of forward propagation calculation are updated first.

2. The method of claim 1, wherein, Before performing the current round of backpropagation computation using the AI ​​accelerator, the method further includes: During the current round of forward propagation calculation using the AI ​​accelerator, for each layer of the model, a corresponding storage space is allocated in the memory of the AI ​​accelerator to store the model parameters required by that layer in the current round of forward propagation. The step of caching the gradient in the memory of the AI ​​accelerator after calculating the gradient of each layer of the model includes: The gradients of each layer of the model are calculated using the AI ​​accelerator. After the gradient of each layer of the model is calculated, the gradient of that layer is stored in the memory of the AI ​​accelerator in the storage space corresponding to that layer, so as to overwrite the model parameters of that layer.

3. The method of claim 2, wherein, The process of calculating the gradients of each layer of the model using the AI ​​accelerator, and storing the gradient of each layer in the memory of the AI ​​accelerator corresponding to that layer after calculation to overwrite the model parameters of that layer, includes: For any layer of the model, after the AI ​​accelerator completes the calculation of the gradient of that layer, the mapping relationship in the address mapping table is updated so as to rebind the physical memory block in the memory of the AI ​​accelerator that was originally used to store the model parameters of that layer to the logical address space used to index the gradient of that layer. Based on the updated address mapping table, the AI ​​accelerator is controlled to perform a memory write operation to write the calculated layer number ladder to the physical memory block and overwrite the model parameters of that layer.

4. The method of claim 1, wherein, The sequential scheduling and execution of parameter update tasks in the task pool includes: Obtain the actual global gradient norm of the previous iteration, and the predicted global norm of the previous iteration; The actual global gradient norm of the previous iteration and the predicted global norm of the previous iteration are weighted using a preset weighting factor to obtain the predicted global gradient norm of the current iteration. For any parameter update task to be executed, the gradients required for the parameter update task are clipped based on the predicted global gradient norm, and the clipped gradients are used to update the model parameters required for the parameter update task, so as to realize the parallel execution of gradient calculation and model parameter update.

5. The method of claim 4, wherein, After updating the model parameters required for the parameter update task using the clipped gradient, the method further includes: After the AI ​​accelerator completes the current round of backpropagation calculation, the actual global gradient norm of the current iteration is determined based on the gradients of each layer of the model cached in the AI ​​accelerator. The prediction bias is determined based on the predicted global gradient norm and the actual global gradient norm; If the prediction deviation is greater than a preset threshold, the updated model parameters are corrected.

6. The method of claim 1, wherein, The gradient includes aggregated gradients, and the method further includes: prior to performing the current round of backpropagation computation using the AI ​​accelerator. Obtain the sample data corresponding to the current round of model training task, and divide the sample data into multiple batches; Issue task instructions to the AI ​​accelerator to cause it to perform the following operations: During the current round of forward propagation computation, for each layer of the model, the multiple batches of forward computation at that layer are executed consecutively; During the current round of backpropagation calculation, for each layer of the model, the backpropagation calculation of the multiple batches at that layer is executed continuously, and the gradients of the multiple batches at that layer are aggregated to obtain the aggregated gradient of that layer, and the aggregated gradient is cached in the memory of the AI ​​accelerator. The step of obtaining and executing the parameter update task from the task pool includes: The aggregated gradient is retrieved from the memory of the AI ​​accelerator, and a model parameter update task is performed based on the aggregated gradient to update the model parameters of the corresponding layer using the aggregated gradient.

7. The method of claim 6, wherein, During the current round of forward and backward propagation computations, the AI ​​accelerator performs multiple batches of computations in opposite order in adjacent layers of the model. This is so that the gradient values ​​of each batch in each layer during the backward propagation computation can reuse the memory of the activation values ​​of that batch in that layer during the forward propagation computation.

8. The method of claim 1, wherein, The processing node also includes an external storage device independent of the processor and the AI ​​accelerator. The model parameters are stored in the external storage device. The step of retrieving and executing the parameter update task from the task pool to update the model parameters using calculated gradients includes: For the parameter update task to be executed, a data transmission channel is established between the external storage device and the processor's memory using direct memory access technology, so as to load the model parameters required for the parameter update task into the processor's memory; The gradient of this layer is used to perform update calculations on the loaded model parameters, and the updated model parameters are written back to the external storage device through the data transmission channel.

9. A computer program product, characterised in that, The computer program product includes a computer program that, when executed, implements the method as described in any one of claims 1-8.

10. An electronic device, comprising: The electronic device includes a processor, a memory, and computer instructions stored in the memory that are executable by the processor, wherein the processor executes the computer instructions to implement the method as described in any one of claims 1-8.

11. A computer readable storage medium, characterized in that, The computer-readable storage medium stores a computer program that, when executed by a processor, implements the method as described in any one of claims 1-8.

Images