Memory optimization method and device for cross-domain distributed model training

Abstract

The application discloses a GPU optimization method and device for cross-domain distributed model training, relates to the technical field of data processing, and is applied to any computing node in a computing cluster. The optimizer state tensor is loaded into a graphics processing unit in advance before entering a parameter updating stage, and after completing model parameter updating based on the optimizer state tensor, the newly generated optimizer state tensor is transferred to the memory of a central processing unit to update the optimizer state tensor stored in the memory and delete the optimizer state tensor in the graphics processing unit. Compared with the scheme in the prior art in which the optimizer state tensor stays in the GPU memory throughout the training period, the optimizer state tensor is dynamically scheduled according to the detection result of the model training stage, and the peak memory occupancy rate of the GPU in the training process is reduced.

Description

Technical Field

[0001] This disclosure relates to the field of data processing technology, specifically to a memory optimization method, apparatus, electronic device, storage medium, and program product for cross-domain distributed model training. Background Technology

[0002] In training large language models with a large number of model parameters, the total storage overhead of model parameters, gradients, and optimizer state tensors (such as the first and second momentum in the Adam optimizer) far exceeds the video memory capacity of a single graphics processing unit (GPU).

[0003] In related technologies, the optimizer state tensor resides in the GPU memory throughout the entire training cycle. However, the optimizer state tensor is only accessed during the actual parameter update phase. It does not participate in the computation at all during the forward and backward propagation phases, which account for the majority of the training time, yet it continuously occupies a large amount of GPU memory resources, resulting in excessively high peak GPU memory usage during the training process. Summary of the Invention In view of this, this disclosure provides a method, apparatus, electronic device, storage medium, and program product for cross-domain distributed model training to solve the problem of excessive peak memory usage of GPUs during existing training processes.

[0004] Firstly, this disclosure provides a memory optimization method for cross-domain distributed model training, applied to any computing node in a computing cluster. Each computing node integrates a central processing unit (CPU) and a graphics processing unit (GPU). The CPU of each computing node stores the full set of model parameters for the model to be trained. During training, each computing node reads the full set of model parameters from its corresponding CPU and loads them into the GPU of its corresponding computing node to perform training operations. Different computing nodes in the computing cluster use different training datasets during training. The method includes: in response to a training start operation, performing an initial training operation on the model based on the training dataset and the loaded full set of model parameters; when the model training enters the parameter update phase, utilizing the initial... The generated optimizer state tensor is used to update the model parameters; the optimized optimizer state tensor with updated model parameters is transferred from the graphics processor (GPU) to the memory of the corresponding central processing unit (CPU) and then deleted from the GPU; for the training of the target epoch after the first completion of the model parameter update, a target timing detection operation is performed before each epoch of model training enters the parameter update phase; when the target timing is detected, the optimized optimizer state tensor in the CPU memory is loaded into the GPU for model parameter update; after the model parameters are updated, the optimized optimizer state tensor stored in the CPU memory is updated using the newly generated optimized optimizer state tensor in the GPU, and the newly generated optimized optimizer state tensor is deleted from the GPU. Secondly, this disclosure provides a memory optimization device for cross-domain distributed model training, applied to any computing node in a computing cluster. Each computing node integrates a central processing unit (CPU) and a graphics processing unit (GPU). The CPU of each computing node stores the full set of model parameters for the model to be trained. During training, each computing node reads the full set of model parameters from its corresponding CPU and loads them into the GPU of its corresponding computing node to perform training operations. Different computing nodes in the computing cluster use different training datasets during training. The device includes: a first training module, used to perform the first round of training on the model based on the training dataset and the loaded full set of model parameters in response to a training start operation; a first update module, used to update the model parameters using the initially generated optimizer state tensor when the model training enters the parameter update phase; a first transfer module, used to transfer the optimizedr state tensor after the model parameter update from the GPU to the memory of the corresponding CPU and delete the optimizedr state tensor from the GPU; a detection module, used to perform a detection operation on the target timing before each round of model training enters the parameter update phase for the target round of training after the first completion of the model parameter update; and a second update module. The third update module is used to load the optimizer state tensor in the central processing unit memory into the graphics processor to update the model parameters when the arrival of the target is detected; the third update module is used to update the optimizer state tensor stored in the central processing unit memory with the newly generated optimizer state tensor in the graphics processor after the model parameters are updated, and to delete the newly generated optimizer state tensor from the graphics processor. Thirdly, this disclosure provides an electronic device, including: a memory and a processor, which are communicatively connected to each other. The memory stores computer instructions, and the processor executes the computer instructions to perform the memory optimization method for cross-domain distributed model training described in the first aspect or any corresponding embodiment.

[0005] Fourthly, this disclosure provides a computer-readable storage medium storing computer instructions for causing a computer to execute the memory optimization method for cross-domain distributed model training described in the first aspect or any corresponding embodiment thereof.

[0006] Fifthly, this disclosure provides a computer program product, including computer instructions for causing a computer to execute the memory optimization method for cross-domain distributed model training described in the first aspect or any corresponding embodiment.

[0007] The memory optimization method for cross-domain distributed model training provided in this disclosure involves the following steps: At the start of model training, the computing node performs the first round of training on the model based on the corresponding training dataset and the loaded full model parameters. During the parameter update phase, the model parameters are updated using the initially generated optimizer state tensor. After the first round of parameter updates, the optimizer state tensor is transferred from the graphics processor (GPU) to the memory of the central processing unit (CPU) corresponding to the computing node, and then deleted from the GPU. In subsequent target rounds of model training, before each round of model training enters the parameter update phase, a target timing detection operation is performed. When the target timing is detected, the optimizer state tensor in the CPU's memory is loaded into the GPU in advance for model training. The system performs a parameter update operation and, after updating the model parameters, updates the optimizer state tensor stored in memory using the newly generated optimizer state tensor, and then deletes the newly generated optimizer state tensor from the graphics processor. This is achieved by pre-loading the optimizer state tensor into the graphics processor before entering the parameter update phase, and then transferring the newly generated optimizer state tensor to the central processing unit's memory after completing the model parameter update based on it to update the optimizer state tensor stored in memory, and finally deleting the optimizer state tensor from the graphics processor. Compared to related technologies where the optimizer state tensor resides in GPU memory throughout the entire training cycle, this system dynamically schedules the optimizer state tensor based on the detection results during model training, reducing the peak GPU memory usage during training. Attached Figure Description

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

[0009] Figure 1 This is a schematic diagram of a scenario corresponding to the memory optimization method for cross-domain distributed model training according to an embodiment of this disclosure; Figure 2 This is a schematic diagram corresponding to the memory optimization method for cross-domain distributed model training according to an embodiment of the present disclosure; Figure 3 This is a structural block diagram of a memory optimization device for cross-domain distributed model training according to an embodiment of the present disclosure; Figure 4 This is a schematic diagram of the hardware structure of an electronic device according to an embodiment of the present disclosure. Detailed Implementation

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

[0011] It is understood that before using the technical solutions disclosed in the various embodiments of this disclosure, users should be informed of the types, scope of use, and usage scenarios of the personal information involved in this disclosure in an appropriate manner in accordance with relevant laws and regulations, and user authorization should be obtained.

[0012] For example, upon receiving a user's active request, a prompt message is sent to the user to explicitly inform them that the requested operation will require the acquisition and use of the user's personal information. This allows the user to independently choose whether to provide personal information to the software or hardware, such as the electronic device, application, server, or storage medium performing the operations of this disclosed technical solution, based on the prompt message.

[0013] As an optional but non-limiting implementation, in response to a user's active request, sending a prompt message to the user can be done via a pop-up window, where the prompt message can be presented in text format. Furthermore, the pop-up window can also include a selection control allowing the user to choose "agree" or "disagree" to provide personal information to the electronic device.

[0014] It is understood that the above notification and user authorization process are merely illustrative and do not constitute a limitation on the implementation of this disclosure. Other methods that comply with relevant laws and regulations may also be applied to the implementation of this disclosure.

[0015] It is understood that the data involved in this technical solution (including but not limited to the data itself, the acquisition or use of the data) shall comply with the requirements of relevant laws, regulations and related provisions.

[0016] According to an embodiment of this disclosure, a memory optimization method for cross-domain distributed model training is provided. It should be noted that the steps shown in the flowchart in the accompanying drawings can be executed in a computer system such as a set of computer-executable instructions. Furthermore, although a logical order is shown in the flowchart, in some cases, the steps shown or described may be executed in a different order than that shown here.

[0017] This embodiment provides a memory optimization method for cross-domain distributed model training, applied to any computing node in a computing cluster. Each computing node integrates a central processing unit (CPU) and a graphics processing unit (GPU). This embodiment does not limit the number of CPUs and GPUs in each computing node. Each computing node's CPU stores the full set of model parameters for the model to be trained. During training, each computing node reads the full set of model parameters from its corresponding CPU and loads them into its corresponding GPU to perform the training operation. Different computing nodes in the cluster use different training datasets. Specifically, the training datasets used by different computing nodes can be stored in their respective CPUs or in a shared CPU storage area within the cluster. Specifically, the storage area can be registered as paged memory using a page-locking mechanism, allowing data transfer between the CPU and GPU to be performed via Direct Memory Access (DMA) to improve data transfer efficiency. Figure 1 As shown in the example, this application takes cross-domain distributed training of the same model by two computing clusters (computing cluster 1 and computing cluster 2) with cross-domain communication as an example. Each computing cluster contains two computing nodes, namely computing node 1 and computing node 2, and computing node 3 and computing node 4. Each computing node contains a graphics processing unit (GPU) and a corresponding central processing unit (CPU). The graphics processing units contained in one computing cluster are defined as GPU1 and GPU2, and the corresponding central processing units are configured as CPU1 and CPU2; the graphics processing units contained in the other computing cluster are defined as GPU3 and GPU4, and the corresponding central processing units are configured as CPU3 and CPU4. The memory optimization method in this application embodiment is explained in conjunction with this example scenario. Figure 2 This is a flowchart of a memory optimization method for cross-domain distributed model training according to an embodiment of the present disclosure, such as... Figure 2 As shown, the process includes the following steps: Step 101: In response to the training start operation, perform the first round of training on the model based on the training dataset and the loaded full model parameters.

[0018] For example, the training dataset and model parameters used in the first round of training can be loaded from the CPU to the GPU before training begins, thereby reducing the data transfer time during the iteration process and improving the overall training efficiency. The full set of model parameters may include the model weight tensor, containing the weight matrices and bias terms of each layer of the network; this type of model parameter does not include the optimizer state tensor.

[0019] Step 102: When the model training enters the parameter update phase, the model parameters are updated using the initially generated optimizer state tensor.

[0020] For example, when model training enters the parameter update phase, the model parameters are updated using the initially generated optimizer state tensor. This embodiment uses the Adam optimizer, whose first-order momentum (exp_avg) and second-order momentum (exp_avg_sq) employ a lazy initialization strategy. Before the first call to optimizer.step(), the optimizer has no momentum tensor storage. After the gradients are obtained through forward and backward propagation in the first round of training, the first entry into the parameter update phase triggers the Adam kernel's lazy initialization. Based on the size of the model parameter tensor within the GPU, the kernel dynamically allocates memory in the GPU memory, generates and initializes the optimizer state tensor composed of the first-order and second-order momentum. Based on this initially generated optimizer state tensor, the adaptive parameter update amount is calculated, completing the model parameter update within the GPU. Specifically, the optimizer internally maintains an empty dictionary `state`, with the key being the model parameters. In each step, it iterates through the parameters, checks if the parameters are in `state`, and if they are not, it obtains the parameter tensor `shape`, creates two zero-based tensors of the same shape `exp_avg` and `exp_avg_sq`, and stores them in `state`. If they already exist, it directly retrieves the existing momentum tensor and uses momentum and gradients to calculate adaptive update values ​​and correct the parameters.

[0021] Step 103: Transfer the optimized state tensor with updated model parameters from the graphics processor to the memory of the corresponding central processing unit and delete the optimized state tensor from the graphics processor.

[0022] For example, after the model parameters are updated, the optimizer state tensor can be migrated to CPU paged memory storage, and the corresponding tensor in GPU memory can be deleted to release the memory. Specifically, the optimizer state tensor storing first-order and second-order momentum in GPU memory is transferred as a whole to the CPU's paged memory for storage; after the data migration is completed, the optimizer state tensor in GPU memory is deleted, and this part of the memory resources is reclaimed to reduce the long-term memory occupation of the GPU. Optionally, CPU paged memory can be replaced with ordinary pageable memory. Although direct DMA transfer cannot be used between the CPU and GPU, and CPU intervention is required for data copying, which reduces the transfer rate, basic offloading functions can still be completed when CPU memory is scarce to optimize GPU memory.

[0023] Step 104: For the target training rounds after the first completion of model parameter update, perform a target timing detection operation before each round of model training enters the parameter update phase.

[0024] For example, in each training round after the initial model parameter update, a target timing detection operation is performed first, just before the parameter update phase begins, after forward propagation and backward gradient calculation. If the target timing is determined to have been reached, the optimizer state tensor is preloaded from the CPU memory to the GPU for model parameter updates. After the model parameter updates are completed, the next training round begins. By adding timing detection before parameter updates, the optimizer state tensor can be preloaded, reducing long-term GPU memory usage and latency caused by cross-device data transfer.

[0025] The target timing can be set after the backpropagation phase is completed, such as when all gradients have been generated. At this point, the detection anchor point is set, and the target timing is detected when the anchor point is reached. Alternatively, the anchor point can be set before calling optimizer.step(). Only when the target timing is determined to have been reached will the optimizer state tensor be loaded from the CPU and step() be executed to complete the parameter update. If the target timing has not been reached, the parameter update process will be skipped directly, which improves the control effect on GPU memory.

[0026] Step 105: When the arrival of the target is detected, the optimizer state tensor in the central processing unit memory is loaded into the graphics processor to update the model parameters.

[0027] For example, when the target timing is determined, the optimizer state tensor, consisting of the first-order and second-order momentum, stored in the central processing unit's memory, is read and loaded into the graphics processing unit's video memory. Based on the loaded optimizer state tensor and the current gradient, Adam optimization calculation is performed to solve for the model parameter update increment and complete the iterative update of the model parameters.

[0028] Step 106: After the model parameters are updated, the optimizer state tensor stored in the central processing unit memory is updated using the newly generated optimizer state tensor in the graphics processor, and the newly generated optimizer state tensor is deleted from the graphics processor.

[0029] For example, after completing the model parameter update calculation, the graphics processing unit (GPU) generates new first-order and second-order momentum tensors that incorporate the gradient information from this round, i.e., the newly generated optimizer state tensor. This set of tensors is transferred to the CPU memory, overwriting and updating the old version of the optimizer state tensor stored in memory, and saving the latest gradient statistics. After the data write-back is completed, the newly generated optimizer state tensor in the GPU memory is deleted, and the corresponding GPU memory resources are reclaimed. Optionally, the optimizer state tensor offloading target can be expanded from CPU memory to NVMe solid-state drives. When CPU memory is insufficient to accommodate all optimizer states, a two-level offloading strategy (GPU → CPU → NVMe) is adopted to obtain greater model training capacity.

[0030] The memory optimization method for cross-domain distributed model training provided in this embodiment loads the optimizer state tensor into the GPU before entering the parameter update phase. After updating the model parameters based on the optimizer state tensor, the newly generated optimizer state tensor is transferred to the CPU's memory to update the optimizer state tensor stored in memory, and the optimizer state tensor in the GPU is deleted. Compared with the related technology where the optimizer state tensor resides in the GPU memory throughout the entire training cycle, this method dynamically schedules the optimizer state tensor based on the detection results during the model training phase, thereby reducing the peak GPU memory usage during training.

[0031] In some optional implementations, each computing node's graphics processor is configured with a corresponding parameter transfer buffer; the method further includes: for any computing node, after completing local training for the target round based on the corresponding training dataset, sharing the full model parameters obtained by its own training with other computing nodes through the corresponding parameter transfer buffer; obtaining the full model parameters obtained by other computing nodes through the sharing operation; and fusing the full model parameters obtained by its own training with the full model parameters obtained by other computing nodes to obtain the full model parameters trained for the current computing cluster.

[0032] For example, each compute node's graphics processing unit (GPU) is allocated an independent parameter buffer. This buffer can be used to temporarily cache the full model parameters to be interacted with, and is isolated from the training compute memory. After the current compute node completes local training for the target round based on the corresponding training dataset, it stores the full model parameters obtained from the local iteration update into the parameter buffer of its local GPU. Based on this buffer, it completes data interaction with other compute nodes in the cluster, sends local model parameters outward, and simultaneously receives and caches the full model parameters uploaded by other nodes. It also performs fusion calculations (such as weighted average, weighted fusion based on sample size, etc.) on the full model parameters it trained and the full model parameters of other compute nodes. Specifically, it can integrate the weight information obtained independently by each compute node to generate the full model parameters after training that are adapted to the entire compute cluster. Optionally, after each round of training, the parameter update amount for each round can be accumulated to the CPU buffer directly through a callback registration mechanism. After N steps of accumulation, cross-domain synchronization is performed, which can avoid snapshot storage overhead and eliminate the need to save old parameter copies.

[0033] In some optional implementations, the method further includes: grouping the model parameters that need to be shared, and using the video memory capacity occupied by the group with the largest number of model parameter elements in all groups as the video memory size of the parameter transfer temporary storage area.

[0034] For example, parameter bucketing technology is used to group model parameters that need to be shared across computing nodes. According to preset bucketing rules (such as by the number of parameter elements, tensor dimension, or network layer division), all model parameters to be shared are divided into multiple parameter groups (parameter buckets). The total number of model parameter elements in each group is counted and the corresponding GPU memory usage is calculated. The group with the most parameter elements and the largest GPU memory usage is selected, and its GPU memory capacity is used as the GPU memory size of the parameter transfer buffer allocated in the graphics processor. When sharing model parameters, only one parameter group is loaded into this parameter transfer buffer each time to complete the transmission and reception. Because the size of the parameter transfer buffer matches the largest parameter bucket, it can accommodate any group of parameters, avoiding transmission buffer overflow. At the same time, the buffer size is dynamically allocated based on the actual usage of the buckets to avoid idle GPU memory resources. Combined with a dedicated transfer buffer to isolate transmitted data from training data, this further optimizes the GPU memory utilization efficiency. Optionally, a temporary storage area can be allocated independently for each bucket. Although this increases the GPU memory overhead, it enables asynchronous overlap of communication and computation between buckets (through interleaved execution of multiple CUDA Streams), which is suitable for scenarios with ample PCIe bandwidth and a need for maximum training throughput.

[0035] In some optional implementations, each computing node's graphics processor is further configured with a corresponding parameter increment transfer buffer; the method further includes: performing parameter increment calculation on the trained full model parameters and the pre-training full model parameters; sharing the calculated parameter increment with other computing clusters under the cross-domain distributed training architecture through the parameter increment transfer buffer; obtaining the parameter increments corresponding to other computing clusters through the sharing operation; fusing the parameter increment corresponding to its own computing cluster with the parameter increments of other computing clusters to obtain the global parameter increment after model training under the cross-domain distributed training architecture; and superimposing the global parameter increment on the full model parameters in the current computing node's central processing unit to obtain the globally shared, updated full model parameters.

[0036] For example, each computing node's graphics processor is additionally allocated a separate parameter increment transfer buffer. This buffer caches parameter increments exchanged across clusters and is isolated from the intra-cluster parameter transfer buffer and training computation memory. After fusing the parameters within the current computing cluster to obtain the full set of trained model parameters, the difference is calculated between these parameters and the baseline full set of model parameters stored in the central processing unit before training. This difference yields the parameter increment representing the weight changes. This parameter increment is stored in the parameter increment transfer buffer and shared with other computing clusters based on the cross-domain distributed architecture. Simultaneously, it receives and caches parameter increments uploaded by other computing clusters. Finally, the parameter increment corresponding to its own computing cluster is fused with the parameter increments of other computing clusters (e.g., weighted averaging, sample-size weighted fusion, or other distributed fusion strategies) to obtain the global parameter increment after model training under the cross-domain distributed training architecture. The original full model parameters stored in the CPU memory are read, and the global parameter increments are added to the original full model parameters element by element to generate global full model parameters that are shared and uniformly updated across all computing clusters. The updated full model parameters are saved in the CPU memory as the benchmark weights for the next round of intra-cluster training and cross-domain incremental synchronization.

[0037] In some optional implementations, the method further includes: grouping the parameter increments that need to be shared, and using the video memory capacity occupied by the group with the largest number of parameter increment elements in all groups as the video memory size of the parameter increment transfer buffer.

[0038] For example, a bucketing mechanism is used to handle parameter increments transmitted across clusters. Specifically, all parameter increments to be shared are divided into multiple parameter increment groups. The total number of parameter increment elements in each group is counted, and the memory usage of each group is calculated. The group with the most increment elements and the highest memory usage is selected. The memory size of the parameter increment transfer buffer in the graphics processor is allocated based on the memory capacity of this group. When parameter increments are exchanged between cross-domain computing clusters, only one set of parameter increments is loaded into the parameter increment transfer buffer at a time to complete the transmission and reception. Since the buffer size matches the largest increment group, it can accommodate all types of increment groups, avoiding transmission buffer overflow. At the same time, the buffer is dynamically configured based on the actual size of the increment, reducing idle memory. Furthermore, the increment transfer buffer is isolated from the training memory and the intra-cluster parameter transfer buffer, further optimizing the utilization of graphics processor memory resources.

[0039] In some optional implementations, the method further includes: constructing a sharded data view; mapping the sharded data view to a storage area of ​​the training dataset of the storage computing cluster and a data range allocated to the current computing node; and reading the training dataset corresponding to the current computing node according to the mapping relationship.

[0040] For example, the complete training dataset is pre-sharded, and a dedicated data interval is assigned to each computing node. A sharded data view is constructed to establish a mapping relationship between computing nodes and their dedicated data intervals, ensuring that each computing node only reads its allocated training data. Specifically, the complete training dataset is sharded, which can be divided into continuous or discrete data intervals based on dimensions such as sample size and storage path. Based on the sharding information, a sharded data view is generated. This view is a set of structured indexes, recording the storage address, start and end indices of the data interval, and the identifier of the computing node to which each data shard belongs. It serves only as a logical index and does not allocate additional physical memory. A mapping and binding relationship is established between the sharded data view and the data interval allocated to the current computing node, locking the data range accessible to the local machine. At training startup, the training data within the corresponding interval in shared storage is read according to this mapping relationship and used as the training dataset for the current computing node. The sharded data view exists only as a logical index and does not generate a copy of the original data. Each node only loads its own sharded data, reducing memory and GPU memory usage, while enabling parallel reading of the dataset by multiple nodes, improving the data loading efficiency of distributed training.

[0041] In some alternative implementations, the method further includes updating the full model parameters stored in the central processing unit using the globally shared, updated full model parameters.

[0042] For example, global parameter increments are superimposed on the full model parameters stored on the CPU before training, generating a fully updated global model parameter set shared by all cross-domain clusters. Transmitting lightweight parameter increments across clusters instead of transmitting the complete full parameter model reduces cross-domain communication pressure.

[0043] In one specific implementation of this application, when any computing node contains multiple GPUs, these GPUs can train the model using the same training dataset and in conjunction with all model parameters. Alternatively, the training dataset corresponding to the computing node can be evenly divided according to the number of GPUs, so that each GPU trains the model based on its assigned subset of training data and in conjunction with all model parameters. After the training of all GPUs on the same computing node is completed, the results can be fused and used as the full set of model parameters obtained from training on the current computing node.

[0044] The memory optimization method for cross-domain distributed model training provided in this application significantly reduces long-term GPU memory usage by dividing mutually exclusive forward / backward propagation, dual training phases for parameter updates, and dynamically scheduling the optimizer's state lifecycle. It loads the GPU momentum tensor only when parameter updates are needed, releasing the corresponding memory in other stages, and implementing this method non-intrusively while remaining compatible with Adam's lazy initialization logic. It constructs a CPU hierarchical persistent cache and a GPU shared bucket intermediate cache system, achieving parallel data sharding computation based on a memory-free sharded view. Cross-domain transmission only involves parameter increments, reducing WAN communication volume. The GPU intermediate cache is uniformly allocated according to the maximum bucket size and serially reused, reducing memory waste. It implements coordinated scheduling of N-step local training and CPU optimizer offloading, accumulating parameter increments through snapshot differential calculation, reducing cross-domain synchronization frequency. It utilizes a shared GPU temporary storage area to complete cross-domain incremental reduction and global model write-back, ensuring CPU and GPU data consistency and guaranteeing model convergence performance while controlling communication overhead.

[0045] Immediately after the parameter update is complete, the momentum tensor is removed from the GPU state dictionary (del state['exp_avg'] and del state['exp_avg_sq']), allowing the GPU memory manager to reclaim its memory instead of holding it during the next forward / backward propagation round. Since the optimizer state (first-order and second-order momentum) resides in GPU memory only during the actual parameter update, and is removed from the GPU optimizer state dictionary and released during the forward / backward propagation phase, which accounts for the majority of the training time, the GPU memory is reclaimed. Preliminary experiments show that, in a cross-domain training environment with a model size of 1.5GB (approximately 0.6B), 4 × NVIDIA RTX 4080 16GB GPUs, and a data parallelism of 2, after enabling the scheme described in the embodiments of this application, the total GPU memory usage decreased from 15719 MiB to 11213 MiB, saving approximately 29.9% (4526 MiB). Among these savings, approximately 816 MiB was independently contributed as one of the innovations (optimizer state unloading), and the total saving was 4526 MiB after contributing as the second innovation (parameter / parameter increment hierarchical buffer).

[0046] Instead of allocating a separate GPU temporary storage area for each bucket, all buckets share a single GPU temporary storage area allocated according to the maximum bucket size. Since bucketing operations are executed serially, only one bucket occupies the temporary storage area at any given time. Therefore, this sharing mechanism reduces the GPU temporary storage memory overhead from "number of buckets × maximum bucket size" to "1 × maximum bucket size" while maintaining functional equivalence.

[0047] After updating the GPU parameters across domains, a process is forced to synchronize only the full model parameters from the device to the host, synchronizing the latest parameter values ​​to the CPU copy to ensure data consistency between the two ends and avoid errors in the subsequent N-step cumulative parameter increment calculation due to the CPU holding an expired parameter snapshot.

[0048] This embodiment also provides a memory optimization device for cross-domain distributed model training. This device is used to implement the above embodiments and preferred embodiments, and details already described will not be repeated. As used below, the term "module" can be a combination of software and / or hardware that implements a predetermined function. Although the device described in the following embodiments is preferably implemented in software, hardware implementation, or a combination of software and hardware, is also possible and contemplated.

[0049] This embodiment provides a memory optimization device for cross-domain distributed model training, applied to any computing node in a computing cluster. Each computing node integrates a central processing unit (CPU) and a graphics processing unit (GPU). The CPU of each computing node stores the full set of model parameters for the model to be trained. During training, each computing node reads the full set of model parameters from its corresponding CPU and loads them into its corresponding GPU to perform the training operation. Different computing nodes in the computing cluster use different training datasets during training. Figure 3 As shown, it includes: The first training module 201 is used to perform the first round of training on the model based on the training dataset and the loaded full model parameters in response to the training start operation. The first update module 202 is used to update the model parameters using the initially generated optimizer state tensor when the model training enters the parameter update stage. The first transfer module 203 is used to transfer the optimized state tensor after the model parameters are updated from the graphics processor to the memory of the corresponding central processing unit and to delete the optimized state tensor from the graphics processor. The detection module 204 is used to perform a detection operation on the target timing before each training round of the model enters the parameter update stage, for the target training round after the first completion of the model parameter update. The second update module 205 is used to load the optimizer state tensor in the central processing unit memory into the graphics processor to update the model parameters when the arrival of the target is detected. The third update module 206 is used to update the optimizer state tensor stored in the central processing unit memory with the newly generated optimizer state tensor in the graphics processor after the model parameters are updated, and to delete the newly generated optimizer state tensor from the graphics processor.

[0050] The memory optimization device for cross-domain distributed model training provided in this embodiment loads the optimizer state tensor into the graphics processor before entering the parameter update stage. After the model parameter update is completed based on the optimizer state tensor, the newly generated optimizer state tensor is transferred to the central processing unit's memory to update the optimizer state tensor stored in memory, and the optimizer state tensor in the graphics processor is deleted. Compared with the related technology, where the optimizer state tensor resides in the GPU memory throughout the entire training cycle, the optimizer state tensor is dynamically scheduled according to the detection results during the model training stage, which reduces the peak memory usage of the GPU during training.

[0051] In some optional implementations, each computing node's graphics processor is configured with a corresponding parameter transfer buffer; the device further includes: a first sharing module, used for any computing node, after completing local training for a target round based on the corresponding training dataset, to share the full model parameters obtained by its own training with other computing nodes through the corresponding parameter transfer buffer; a first acquisition module, used to obtain the full model parameters obtained by other computing nodes through the sharing operation; and a first fusion module, used to fuse the full model parameters obtained by its own training with the full model parameters obtained by other computing nodes to obtain the full model parameters after training corresponding to the current computing cluster.

[0052] In some optional implementations, each computing node's graphics processor is further configured with a corresponding parameter increment transfer buffer; the device further includes: a computing module, used to perform parameter increment calculation on the trained full model parameters and the pre-training full model parameters; a second sharing module, used to share the calculated parameter increments with other computing clusters under the cross-domain distributed training architecture through the parameter increment transfer buffer; a second acquisition module, used to obtain the parameter increments corresponding to other computing clusters through the sharing operation; a second fusion module, used to fuse the parameter increments corresponding to its own computing cluster with the parameter increments of other computing clusters to obtain the global parameter increments after model training under the cross-domain distributed training architecture; and an overlay module, used to overlay the global parameter increments on the full model parameters in the central processing unit of the current computing node to obtain globally shared, updated full model parameters.

[0053] In some optional embodiments, the apparatus further includes: a first determining module, configured to group the model parameters that need to be shared, and use the video memory capacity occupied by the group with the largest number of model parameter elements in all groups as the video memory size of the parameter transfer temporary storage area; and a second determining module, configured to group the parameter increments that need to be shared, and use the video memory capacity occupied by the group with the largest number of parameter increment elements in all groups as the video memory size of the parameter increment transfer temporary storage area.

[0054] In some optional embodiments, the apparatus further includes: a construction module for constructing a sharded data view; a mapping module for mapping the sharded data view to a storage area of ​​the training dataset of the storage computing cluster and a data range allocated to the current computing node; and a reading module for reading the training dataset corresponding to the current computing node according to the mapping relationship.

[0055] In some optional embodiments, the apparatus further includes a fourth update module for updating the full model parameters stored in the central processing unit using the globally shared, updated full model parameters.

[0056] The memory optimization apparatus for cross-domain distributed model training provided in this disclosure can execute the memory optimization method for cross-domain distributed model training provided in any embodiment of this disclosure, and has the corresponding functional modules and beneficial effects for executing the method. Further functional descriptions of the above modules and units are the same as in the corresponding embodiments described above, and will not be repeated here.

[0057] Figure 4 This is a schematic diagram of the structure of an electronic device provided in an embodiment of this disclosure.

[0058] The following is a detailed reference. Figure 4 The diagram illustrates a structural schematic suitable for implementing an electronic device according to embodiments of the present disclosure. The electronic device may include a processor (e.g., a central processing unit, graphics processor, etc.) 501, which can perform various appropriate actions and processes according to a program stored in read-only memory (ROM) 502 or a program loaded from memory 508 into random access memory (RAM) 503. The RAM 503 also stores various programs and data required for the operation of the electronic device. The processor 501, ROM 502, and RAM 503 are interconnected via a bus 504. An input / output (I / O) interface 505 is also connected to the bus 504.

[0059] Typically, the following devices can be connected to I / O interface 505: input devices 506 including, for example, touchscreens, touchpads, keyboards, mice, cameras, microphones, accelerometers, gyroscopes, etc.; output devices 507 including, for example, liquid crystal displays (LCDs), speakers, vibrators, etc.; memory devices 508 including, for example, magnetic tapes, hard disks, etc.; and communication devices 509. Communication device 509 allows electronic devices to communicate wirelessly or wiredly with other devices to exchange data. Although Figure 4 Electronic devices with various devices are shown, but it should be understood that it is not required to implement or have all of the devices shown, and more or fewer devices may be implemented or have instead.

[0060] In particular, according to embodiments of this disclosure, the processes described above with reference to the flowcharts can be implemented as computer software programs. For example, embodiments of this disclosure include a computer program product comprising a computer program carried on a non-transitory computer-readable medium, the computer program containing program code for performing the methods shown in the flowcharts. In such embodiments, the computer program can be downloaded and installed from a network via communication device 509, or installed from memory 508, or installed from ROM 502. When the computer program is executed by processor 501, it performs the functions defined in the memory optimization method for cross-domain distributed model training according to embodiments of this disclosure.

[0061] Figure 4 The electronic device shown is merely an example and should not be construed as limiting the functionality and scope of the embodiments disclosed herein.

[0062] This disclosure also provides a computer-readable storage medium in which the methods described in this disclosure can be implemented in hardware or firmware, or implemented as recordable on a storage medium, or implemented as computer code downloaded over a network and originally stored on a remote storage medium or a non-transitory machine-readable storage medium and subsequently stored on a local storage medium. Thus, the methods described herein can be processed by software stored on a storage medium using a general-purpose computer, a dedicated processor, or programmable or dedicated hardware. The storage medium can be a magnetic disk, optical disk, read-only memory, random access memory, flash memory, hard disk, or solid-state drive, etc.; further, the storage medium may also include combinations of the above types of memory. It is understood that computers, processors, microprocessor controllers, or programmable hardware include storage components capable of storing or receiving software or computer code. When the software or computer code is accessed and executed by the computer, processor, or hardware, the memory optimization method for cross-domain distributed model training shown in the above embodiments is implemented.

[0063] A portion of this disclosure can be applied to computer program products, such as computer program instructions, which, when executed by a computer, can invoke or provide methods and / or technical solutions according to this disclosure through the operation of the computer. Those skilled in the art will understand that the forms in which computer program instructions exist in a computer-readable medium include, but are not limited to, source files, executable files, and installation package files. Accordingly, the ways in which computer program instructions are executed by a computer include, but are not limited to: the computer directly executing the instructions; the computer compiling the instructions and then executing the corresponding compiled program; the computer reading and executing the instructions; or the computer reading and installing the instructions and then executing the corresponding installed program. Here, the computer-readable medium can be any available computer-readable storage medium or communication medium accessible to a computer.

[0064] Although embodiments of the present disclosure have been described in conjunction with the accompanying drawings, those skilled in the art can make various modifications and variations without departing from the spirit and scope of the present disclosure, and such modifications and variations all fall within the scope defined by the appended claims.

Claims

1. A memory optimization method for cross-domain distributed model training, characterized in that, The method is applied to any computing node in a computing cluster, where each computing node integrates a central processing unit (CPU) and a graphics processing unit (GPU). The CPU of each computing node stores the full set of model parameters for the model to be trained. During training, each computing node reads the full set of model parameters from its corresponding CPU and loads them into its corresponding GPU to perform the training operation. Different computing nodes in the computing cluster use different training datasets during training. The method includes: In response to the start of training operation, the model is trained for the first time based on the training dataset and all loaded model parameters. When the model training enters the parameter update phase, the model parameters are updated using the initially generated optimizer state tensor. The optimizer state tensor with updated model parameters is transferred from the graphics processor to the memory of the corresponding central processing unit and then deleted from the graphics processor. For the target training rounds after the initial model parameter update, a target timing detection operation is performed before the model enters the parameter update phase in each training round. When the arrival at the target is detected, the optimizer state tensor in the central processing unit memory is loaded into the graphics processor to update the model parameters; After the model parameters are updated, the optimizer state tensor stored in the central processing unit memory is updated using the newly generated optimizer state tensor in the graphics processor, and the newly generated optimizer state tensor is deleted from the graphics processor.

2. The method according to claim 1, characterized in that, Each of the computing nodes' graphics processors is configured with a corresponding parameter transfer and temporary storage area; the method further includes: For any computing node, after completing the local training for the target round based on the corresponding training dataset, it shares the full set of model parameters obtained from its own training with other computing nodes through the corresponding parameter transfer and temporary storage area. The full set of model parameters trained by other computing nodes can be obtained through a sharing operation; The full set of model parameters obtained from its own training is fused with the full set of model parameters obtained from the training of other computing nodes to obtain the full set of model parameters after training for the current computing cluster.

3. The method according to claim 2, characterized in that, Each of the computing nodes' graphics processors is also configured with a corresponding parameter increment transfer buffer; the method further includes: The parameters of the trained full model are compared with the parameters of the pre-training full model to calculate the parameter increment; The calculated parameter increments are shared with other computing clusters in the cross-domain distributed mode training architecture through the parameter increment transfer and temporary storage area. The parameter increments corresponding to other computing clusters are obtained through shared operations; The parameter increment corresponding to the computing cluster to which it belongs is fused with the parameter increment of other computing clusters to obtain the global parameter increment after model training under the cross-domain distributed mode training architecture. The global parameter increments are used to superimpose the full model parameters in the central processing unit of the current computing node to obtain the updated full model parameters that are shared globally.

4. The method according to claim 3, characterized in that, The method further includes: The model parameters that need to be shared are grouped, and the video memory capacity occupied by the group with the largest number of model parameter elements is used as the video memory size of the parameter temporary storage area. The parameter increments that need to be shared are grouped, and the video memory capacity occupied by the group with the largest number of parameter increment elements is used as the video memory size of the parameter increment temporary storage area.

5. The method according to claim 3, characterized in that, The method further includes: Build a sharded data view; The sharded data view is mapped to the storage area of ​​the training dataset of the storage computing cluster and the data range allocated to the current computing node; Based on the mapping relationship, read the training dataset corresponding to the current computing node.

6. The method according to claim 3, characterized in that, The method further includes: The updated full model parameters shared across the entire domain are used to update the full model parameters already stored in the central processing unit.

7. A memory optimization device for cross-domain distributed model training, characterized in that, The device is applied to any computing node in a computing cluster. Each computing node integrates a central processing unit (CPU) and a graphics processing unit (GPU). The CPU of each computing node stores the full set of model parameters for the model to be trained. During training, each computing node reads the full set of model parameters from its corresponding CPU and loads them into its corresponding GPU to perform the training operation. Different computing nodes in the computing cluster use different training datasets during training. The device includes: The first training module is used to perform the first round of training on the model based on the training dataset and all loaded model parameters in response to the training start operation. The first update module is used to update the model parameters using the initially generated optimizer state tensor when the model training enters the parameter update phase. The first transfer module is used to transfer the optimized state tensor with updated model parameters from the graphics processor to the memory of the corresponding central processing unit and to delete the optimized state tensor from the graphics processor. The detection module is used to perform a detection operation on the timing of the target training round after the first model parameter update is completed. The second update module is used to load the optimizer state tensor in the central processing unit memory into the graphics processor to update the model parameters when the arrival of the target is detected. The third update module is used to update the optimizer state tensor stored in the central processing unit memory with the newly generated optimizer state tensor in the graphics processor after the model parameters are updated, and to delete the newly generated optimizer state tensor from the graphics processor.

8. An electronic device, characterized in that, include: A memory and a processor are interconnected, the memory stores computer instructions, and the processor executes the computer instructions to perform the memory optimization method for cross-domain distributed model training as described in any one of claims 1 to 6.

9. A computer-readable storage medium, characterized in that, The computer-readable storage medium stores computer instructions for causing the computer to execute the memory optimization method for cross-domain distributed model training as described in any one of claims 1 to 6.

10. A computer program product, characterized in that, Includes computer instructions for causing a computer to execute the memory optimization method for cross-domain distributed model training as described in any one of claims 1 to 6.

Images