Asynchronous training method and device based on 3D parallel and distributed parameter server and storage medium

Abstract

The invention relates to an asynchronous training method and device based on 3D parallel and distributed parameter servers and a storage medium. A training process of parameter server nodes in the method comprises the following steps: S101, initializing global model weights, wherein the global model weights are distributed on a plurality of parameter server nodes; s102, a current local model weight is received, initialization of the local model weight is achieved through a pre-selected 3D parallel training method, and updating of the local model weight is achieved by executing local training of preset time duration through all training tasks; s103, updating the current global model weight based on the local model weight; and S104, issuing the global model weight to the training task, and returning to S102. Compared with the prior art, the method has the advantages that the applicability of distributed model training to heterogeneous equipment clusters is improved, and meanwhile, the training efficiency and stability during processing of the ten-thousand-card-level clusters are guaranteed.

Description

Technical Field

[0001] The present invention relates to the technical field of distributed training, and in particular to an asynchronous training method, device, and storage medium based on 3D parallelism and a distributed parameter server. Background Art

[0002] At present, the training scale of large language models (LLMs) has reached the level of tens of thousands of cards. Most existing training frameworks perform well in the training of homogeneous clusters at the level of tens of thousands of cards. However, due to the large differences in the hardware performance of heterogeneous clusters, problems such as communication delays between nodes and difficulties in parameter synchronization are likely to occur during training. Existing training frameworks have deficiencies in stability and low training efficiency when facing heterogeneous clusters.

[0003] Therefore, how to improve the applicability of distributed model training for heterogeneous device clusters while ensuring the training efficiency and stability when dealing with clusters at the level of tens of thousands of cards has become a problem to be solved in this field. Summary of the Invention

[0004] The purpose of the present invention is to overcome the above-mentioned deficiencies existing in the prior art and provide an asynchronous training method, device, and storage medium based on 3D parallelism and a distributed parameter server to improve the applicability of distributed model training for heterogeneous device clusters while ensuring the training efficiency and stability when dealing with clusters at the level of tens of thousands of cards.

[0005] The purpose of the present invention can be achieved by the following technical solutions:

[0006] According to a first aspect of the present invention, there is provided an asynchronous training method based on 3D parallelism and a distributed parameter server. The method is implemented based on multiple parameter server nodes and multiple heterogeneous clusters that communicate with each other. The local model training on the heterogeneous clusters is divided into multiple different training tasks. The training process of the parameter server nodes specifically includes the following steps: S101, initialize the global model weights, and the global model weights are distributed on the multiple parameter server nodes; S102, receive the current local model weights, and the initialization of the local model weights is achieved through a preselected 3D parallel training method, and the update of the local model weights is achieved through local training of each training task for a preset time length; S103, update the current global model weights based on the local model weights; S104, send the global model weights to the training tasks and return to S102.

[0007] As a preferred technical solution, the training process of the training task specifically includes the following steps: S201, initialize the local model weights and pull the current global model weights; S202, within a preset time length, perform multiple steps of local training to update the local model weights, and the preset time length is determined through negotiation by the parameter server nodes; S203, after the end of the time length, collect the current complete local model weights and upload the local model weights to all parameter server nodes through a two-phase commit protocol; S204, after the global model weights are updated, download the latest global model weights and return to S202.

[0008] As a preferred technical solution, if a failure occurs during the upload or download process, skip the global model weight update for the current round and return to S202.

[0009] As a preferred technical solution, the global model weights are evenly distributed among the multiple parameter server nodes through hashing.

[0010] As a preferred technical solution, in S103, the parameter server node updates the global model weights based on the local model weights and a preselected update strategy, and the update strategy is one of a synchronous update strategy, a partial synchronous update strategy, and an asynchronous update strategy.

[0011] The synchronous update strategy includes that when all training tasks upload their corresponding local model weights, the global model weights are updated; the partial synchronous update strategy includes that when a preset proportion of training tasks upload their corresponding local model weights, the global model weights are updated; the asynchronous update strategy includes that according to the moment when the local model weights are first received and a preset first time window, receive all the local model weights uploaded within the first time window and update the global model weights according to all the local model weights.

[0012] As a preferred technical solution, the parameter server node determines whether there is a training task joining or exiting based on the heartbeat situation from the training task. The specific process includes: when a new heartbeat is received, the parameter server node considers that a new training task has joined the training process; when no heartbeat is received within three consecutive preset second time windows, the parameter server node considers that the corresponding training task has exited.

[0013] As a preferred technical solution, when it is detected that a training task has exited, it automatically restarts using a preset fault tolerance script and rejoins the training process after fixing the error.

[0014] According to a second aspect of the present invention, there is provided an asynchronous training device based on 3D parallel and distributed parameter servers, including a memory, a processor, and a program stored in the memory, and when the processor executes the program, the method described above is implemented.

[0015] According to a third aspect of the present invention, there is provided a storage medium, on which a program is stored, and when the program is executed, the method described above is implemented.

[0016] Compared with the prior art, the present invention has the following beneficial effects:

[0017] 1. The present invention adopts distributed parameter server nodes, disperses computing tasks on different nodes, can reduce the memory and communication requirements of the parameter server, and can be extended to model training tasks on a scale of hundreds of billions. At the same time, the training cluster is divided into different training tasks, so that the cluster scale supported by this method can reach tens of thousands of cards and supports the training of heterogeneous chip clusters, effectively improving the stability of training.

[0018] 2. The present invention initializes the local model weights through a 3D parallel training method, supports differential parallelism of training tasks, and the parallel strategies of each training are independent of each other to maximize the training efficiency of hardware devices.

[0019] 3. In the present invention, the parameter server node judges whether there is a training task joining or exiting according to the heartbeat situation from the training task, supports the dynamic joining and exiting of training tasks, and cooperates with the fault tolerance strategy to effectively increase the effective training time of the model. BRIEF DESCRIPTION OF THE DRAWINGS

[0020] Figure 1 It is a schematic diagram of the 3DPS framework structure and communication in an embodiment of the present invention;

[0021] Figure 2 It is a schematic diagram of the training process of a parameter server node provided in an embodiment of the present invention;

[0022] Figure 3 It is a schematic diagram of the training process of a training task provided in an embodiment of the present invention;

[0023] Figure 4 It is a schematic diagram of the interaction process between the parameter server node and the training task provided in an embodiment of the present invention. DETAILED DESCRIPTION OF THE EMBODIMENTS

[0024] 3D parallelism refers to the simultaneous application of data parallelism, tensor parallelism, and pipeline parallelism in model training, which is a commonly used distributed model training strategy in large-scale language model training. Existing distributed model training methods based on 3D parallelism require full synchronization at each training step, which poses extremely high requirements on the stability of the cluster. However, in the daily operation and maintenance of large-scale clusters, due to the complexity and uncontrollability of hardware nodes, node failures occur frequently on a daily basis. Under the full synchronization scheme, any single-point failure will cause the training task to be interrupted and restarted, significantly reducing the training efficiency.

[0025] In addition, existing distributed parameter server (PS) frameworks, such as DiLoCo and EDiT, have made some progress in large-scale language model training, but there are still some significant limitations. Among them, DiLoCo assumes that the devices within all worker nodes are homogeneous, which limits its applicability in heterogeneous device clusters, and DiLoCo shows diminishing returns when there are more than 8 worker nodes, with limited scalability. At the same time, although DiLoCo performs well in communication efficiency, its distributed computing characteristics lead to a decrease in the FLOP and data efficiency of the model, which may become a bottleneck in large-scale training. In addition, the models supported by DiLoCO only support the training tasks of models with a size of hundreds of megabytes. While EDiT supports heterogeneous devices to a certain extent, its hierarchical distributed strategy may face efficiency problems when dealing with large-scale heterogeneous clusters, and it requires the same parallel strategy to be adopted for different heterogeneous clusters, making it difficult to fully utilize the hardware performance. Although EDiT reduces the communication overhead through hierarchical parameter synchronization, communication is still a potential bottleneck in large-scale heterogeneous clusters. In addition, EDiT introduces a pseudo-gradient penalty strategy to solve the training instability caused by different progress of different worker nodes, but this method may not be flexible enough in some cases.

[0026] Based on this, in order to expand the heterogeneous device support ability of existing distributed training methods and improve the training efficiency and flexibility of training methods, the present invention provides an asynchronous training method based on 3D parallelism and a distributed parameter server (PS), that is, an asynchronous training method based on the 3DPS framework. This method can support heterogeneous chips both between worker nodes and within worker nodes, enabling it to operate efficiently in complex large-scale heterogeneous clusters, fully utilize different types of computing resources, and support more task types and architectures.

[0027] The present invention will be described in detail below with reference to the accompanying drawings and specific embodiments. This embodiment is implemented on the premise of the technical solution of the present invention, and gives detailed implementation manners and specific operation processes, but the protection scope of the present invention is not limited to the following embodiments.

[0028] Example:

[0029] like Figure 1 As shown, the 3DPS framework includes multiple parameter server nodes and multiple heterogeneous clusters that communicate with each other, and the model parameters to be trained are distributed on multiple parameter server nodes to reduce the memory and communication overhead of parameter updates of the parameter servers. At the same time, the 3DPS framework divides the model training on the heterogeneous cluster into different training tasks (Group), and the training tasks regularly upload their locally updated weights to the parameter server nodes, and perform global updates of the parameters on the parameter server nodes. The 3DPS framework provided in this embodiment supports synchronous or asynchronous parameter updates between different training tasks to maximize training efficiency. Finally, the 3DPS framework supports dynamic joining and exiting of training tasks, which greatly improves the stability of Wanka cluster training.

[0030] Parameter Server

[0031] This method can realize a decentralized parameter server, in which the model weights are evenly distributed on different nodes through hashing, and are responsible for calculating the gradient updates of their respective weights to balance memory and communication overhead. Specifically, the parameter server considers the computing power overhead of different training tasks (such as different training tasks are trained with different chips), and performs a weighted average and gradient update on the local gradient updates uploaded according to the number of training samples processed. Furthermore, the parameter server supports two scenarios: fault recovery during idle time and fault recovery during global synchronization. If a failure occurs when the parameter server is not interacting with a training task, you can simply switch to a backup machine to restart the service without affecting the training task; otherwise, the training task will skip this round of global updates and continue local training, and perform a global update after the parameter server is back online. Furthermore, the parameter server supports hot expansion, and the expansion process is unaware of the training task.

[0032] Dynamic joining and exiting of training tasks (elastic expansion mechanism)

[0033] This method takes into account the differences in hardware design and computing power of different chip clusters, allowing the use of different 3D hybrid parallel training methods.

[0034] Furthermore, if a training task exits due to an error (software or hardware failure), it will not affect the training of other tasks, and the task can rejoin the overall training process after the error is fixed. Specifically, the training task periodically sends heartbeats to the parameter server node, and the parameter server node determines whether a training task has joined or exited based on the heartbeat from the training task. The specific process includes: when a new heartbeat is received, the parameter server node considers that a new training task has joined the training process; when the heartbeat of the training task is not received within three consecutive preset second time windows, the parameter server node considers that the corresponding training task has exited.

[0035] Furthermore, the training task supports fault recovery and automatic restart. When it is detected that a certain training task fails and exits, the fault-tolerant script supports automatic restart of the task. Specifically, the parameter server node and the training task write the latest checkpoint to the disk. If the framework detects a fault in a certain service, it restarts the corresponding service and loads the corresponding checkpoint to recover the service.

[0036] The training method based on the 3DPS framework mainly includes two parts, namely the training process of the parameter server node and the training process of the training task. In this embodiment, it is agreed that Θ t represents the global model weight of the model at the t-th step (i.e., the iteration round), represents the model weight of the model saved in the i-th parameter server at the t-th step (saved in the parameter server, where represents the complete local model weight of the model saved in the training task p at the t-th step (saved in the parameter server), represents the local model weight sent from the training task p to the i-th parameter server at the t-th step, where

[0037] Figure 2 Fig. shows one of the training processes of the parameter server node. There are multiple parameter server nodes, and each node maintains a part of the global model weight, denoted as θ i . The parameter server node is mainly responsible for the gradient update of the global model and sends the updated model weight back to the training task. The parameter server node maintains the Group ID of the currently joined training task and supports the dynamic joining and exiting of the training task. Its main training process is as follows

[0038] Step S101, initialize the global model weight. Specifically, according to the predefined model partitioning scheme, each node i loads a part of the weight of the corresponding model, where Θ 0 represents the initial weight of the global model.

[0039] Step S102, receive the current local model weight. Specifically, the parameter server node listens to the service interface, receives the model weight uploaded by the training task, and saves it in the cache.

[0040] Step S103, update the current global model weight based on the current local model weight. Specifically, according to the synchronous update strategy, partial synchronous update strategy, or asynchronous update strategy selected by the user, the gradients of the model weights are weighted and summed for global update. Exemplarily, the gradient is expressed as:

[0041]

[0042] Where λ p is the calculation weight corresponding to the training task p in the weighted summation process, expressed as:

[0043]

[0044] Where token p represents the number of samples processed by the p-th training task in the current round.

[0045] Thus, for the local model weights uploaded by the training tasks, the calculation is performed according to the following formula:

[0046]

[0047] Where N is the total number of training tasks, and α is the learning rate, representing the speed of gradient descent.

[0048] Furthermore, the synchronous update strategy includes that after all training tasks upload their corresponding local model weights, the global model weights are updated; the partial synchronous update strategy includes that after a preset proportion of training tasks upload their corresponding local model weights, the global model weights are updated; the asynchronous update strategy includes that according to the time when the local model weights are first received and a preset first time window, all local model weights uploaded within the first time window are received, and the global model weights are updated according to all the received local model weights.

[0049] Step S104, the global model weights are sent to the training tasks, and step S102 is returned. Specifically, the parameter server node notifies the training tasks to download the latest global model weights and return to step S102 to repeat the update process.

[0050] Figure 3 shows the training process of one of the training tasks. The training task is responsible for the training task of the local model, and regularly uploads the locally updated weights to the parameter server node, and waits for the parameter server node to perform a global update and then pulls the latest weights to the local for local training. Each training task can run on different chip clusters and adopt its own optimal 3D parallel training method. The training samples used by each training task during local training are pre-acquired, and its main training process is as follows:

[0051] Step S201, initialize the local model weights and pull the current global model weights. Specifically, each training task initializes the model according to the local 3D parallel training method and pulls the latest global model weights Θ t-1 .

[0052] Step S202: Within a preset time length, perform multiple steps of local training to update the local model weights. Specifically, the preset time length is actually a synchronization time interval negotiated by the parameter server nodes. Within this synchronization time interval, each training task performs n steps of local training, thereby updating the local model weights. Compared with the existing method that requires communication with the parameter server after each step of local training, this step can ensure that communication with the parameter server is carried out after n steps of local model training, effectively avoiding network latency and improving training efficiency in heterogeneous clusters and environments with poor networks. Moreover, it communicates the local model weights instead of the model gradients, resulting in a relatively low overall communication latency.

[0053] Step S203: After the negotiated time length ends, collect the current complete local model weights and upload the local model weights to all parameter server nodes through a two-phase commit protocol. Specifically, a specified worker node collects a complete set of model weights and uploads the local model weights to all parameter servers according to the hash to ensure that the local model weights are correctly submitted to all parameter server nodes. If a parameter server node fails at this time, the training task will skip the weight upload in this round, return to Step S202, and wait for the parameter server node to come online again to submit the local model weights.

[0054] Step S204: After the global model weights of the parameter server nodes are updated, download the latest global model weights Θ t , and return to Step S202.

[0055] In the training method based on the 3DPS framework provided in this embodiment, the decoupled parameter server nodes and training tasks each maintain their own task states: the parameter server nodes have four states: idle, receiving weights, globally updating, and distributing weights; the training tasks have three states: local training, uploading weights, and distributing weights. The two only interact through the http protocol during the weight upload and distribution processes, and the service at one end is not affected by the service at the other end. The interaction process is as Figure 4 shown.

[0056] Furthermore, the present invention also provides an asynchronous training device based on 3D parallel and distributed parameter servers, including a memory, a processor, and a program stored in the memory. When the processor executes the program, it implements one or more steps of the foregoing method. The specific implementation process is basically the same as the execution process of the foregoing method and will not be elaborated here.

[0057] Furthermore, the present invention also provides a storage medium, on which a program is stored, and when the program is executed, one or more steps of the foregoing method are implemented. The specific implementation process is basically the same as the execution process of the foregoing method, and will not be elaborated herein.

[0058] The preferred specific embodiments of the present invention have been described in detail above. It should be understood that those of ordinary skill in the art can make many modifications and variations based on the concept of the present invention without creative efforts. Therefore, all technical solutions that can be obtained by those skilled in the art in the technical field of the present invention through logical analysis, reasoning, or limited experiments based on the concept of the present invention on the basis of the prior art should fall within the protection scope determined by the claims.

Claims

1. An asynchronous training method based on 3D parallel and distributed parameter servers, characterized in that, The method is implemented based on multiple parameter server nodes and multiple heterogeneous clusters that communicate with each other. The local model training on the heterogeneous clusters is divided into multiple different training tasks. The training process of the parameter server nodes specifically includes the following steps: S101, Initialize the global model weights, and the global model weights are distributed on the multiple parameter server nodes; S102, Receive the current local model weights. The initialization of the local model weights is implemented through a preselected 3D parallel training method, and the update of the local model weights is achieved by performing local training for a preset time length for each training task; S103, Update the current global model weights based on the local model weights; S104, Send the global model weights to the training tasks, and return to S102.

2. The asynchronous training method based on 3D parallel and distributed parameter servers according to claim 1, wherein The training process of the training tasks specifically includes the following steps: S201, Initialize the local model weights and pull the current global model weights; S202, Perform multiple steps of local training within a preset time length to update the local model weights. The preset time length is determined through negotiation by the parameter server nodes; S203, After the end of the time length, collect the current complete local model weights and upload the local model weights to all parameter server nodes through a two-phase commit protocol; S204, After the global model weights are updated, download the latest global model weights and return to S202.

3. The asynchronous training method based on 3D parallel and distributed parameter servers according to claim 2, wherein If a failure occurs during the upload or download process, skip the global model weight update for the current round and return to S202.

4. The asynchronous training method based on 3D parallel and distributed parameter servers according to claim 1, wherein The global model weights are evenly distributed on the multiple parameter server nodes through hashing.

5. The asynchronous training method based on 3D parallel and distributed parameter servers according to claim 1, wherein In S103, the parameter server nodes update the global model weights based on the local model weights and a preselected update strategy, and the update strategy is one of a synchronous update strategy, a partial synchronous update strategy, and an asynchronous update strategy.

6. The asynchronous training method based on 3D parallel and distributed parameter servers according to claim 5, wherein The synchronous update strategy includes that when all training tasks have uploaded their corresponding local model weights, the global model weights are updated; The partial synchronous update strategy includes that when a preset proportion of training tasks have uploaded their corresponding local model weights, the global model weights are updated; The asynchronous update strategy includes that according to the moment when the local model weights are first received and a preset first time window, receive all the local model weights uploaded within the first time window, and update the global model weights according to all the local model weights.

7. The asynchronous training method based on 3D parallel and distributed parameter servers according to claim 1, wherein The parameter server nodes judge whether there are training tasks joining or exiting according to the heartbeat conditions from the training tasks. The specific process includes: When a new heartbeat is received, the parameter server node considers that a new training task has joined the training process; When no heartbeat is received within three consecutive preset second time windows, the parameter server node considers that the corresponding training task has exited.

8. The asynchronous training method based on 3D parallel and distributed parameter servers according to claim 7, characterized in that, After detecting that a training task has exited, automatically restart it using a preset fault tolerance script and rejoin the training process after fixing the error.

9. An asynchronous training device based on 3D parallel and distributed parameter servers, comprising a memory, a processor, and a program stored in the memory, characterized in that, When the processor executes the program, it implements the method according to any one of claims 1-8.

10. A storage medium having a program stored thereon, characterized in that, When the described program is executed, it implements the method according to any one of claims 1-8.

Images