A method for network failure recovery for distributed model training

Abstract

This invention discloses an on-network fault recovery method for distributed model training. By introducing a fault-tolerant aggregation mechanism based on window constraints and integrity constraints into the parameter server, the model parameters for the current training round can still be updated when a training node fails. Simultaneously, the model recovery parameters are stored in the top-of-rack switch (ToR switch) and updated during data forwarding. This allows the ToR switch to maintain a copy of the model recovery parameters corresponding to the current training round. Thus, when a training node fails, the corresponding ToR switch directly sends the model recovery parameters to the failed node, thereby accelerating the on-network fault recovery process for the training node.

Description

Technical Field

[0001] This invention belongs to the field of programmable switches and distributed machine learning fault tolerance technology, specifically relating to an on-network fault recovery method for distributed model training. Background Technology

[0002] With the continuous development of machine learning technology, it has been widely applied in fields such as computer vision, natural language processing, and speech recognition. At the same time, the scale of models and the amount of training data continue to grow, placing higher demands on the computing resources and training efficiency for model training. Distributed machine learning has gradually become an important technical means for large-scale model training.

[0003] In distributed machine learning, centralized parameter server architecture is widely used. In this architecture, each training node independently performs local computation and generates gradient information. The parameter server aggregates the gradients from each training node and distributes the updated model parameters back to the training nodes for the next round of training. Through multi-node parallel computation, this architecture can improve model training efficiency and resource utilization to a certain extent.

[0004] However, in actual training, training nodes may fail due to hardware malfunctions, software anomalies, or network instability. When a training node fails, the training process typically needs to be paused or wait for the failed node to recover, thus impacting training efficiency. Existing training node failure recovery methods mostly rely on parameter servers, requiring the failed node to retrieve the model parameters needed to restore the training state. As the scale of deep learning model parameters continues to increase, this recovery process often involves the transmission of large amounts of data across the network, leading to a significant increase in network communication overhead and introducing high recovery latency. These problems are particularly pronounced when the network path between the parameter server and the failed node is long or when network congestion is severe.

[0005] In recent years, with the development of programmable switches and on-network computing technologies, network devices have gradually acquired certain data processing and storage capabilities, providing new technical approaches for performing some computing tasks within the network and shortening data transmission paths. However, limited by the on-chip storage resources of programmable switches, how to improve the fault recovery efficiency of training nodes in distributed machine learning while ensuring model training accuracy remains a subject for further research. In particular, in technical solutions that involve network devices in fault recovery, ensuring that the model parameters recovered from the faulty node are consistent with those used by the normal training nodes in the current system, and ensuring that the training process after recovery can continue in coordination with the overall training progress, remains a technical problem that needs further resolution. Summary of the Invention

[0006] In view of this, the present invention provides an on-network fault recovery method for distributed model training. By introducing a fault-tolerant aggregation mechanism based on window constraints and integrity constraints into the parameter server, the model parameters for the current training round can still be updated when a training node fails. The model recovery parameters are stored in the top-of-rack switch (ToR switch) and updated during data forwarding, enabling the ToR switch to maintain a copy of the model recovery parameters corresponding to the current training round. When a training node fails, the corresponding ToR switch directly sends the model recovery parameters to the failed node, thereby accelerating the on-network fault recovery process of the training node.

[0007] This invention provides an in-network fault recovery method for distributed model training, which specifically includes the following steps:

[0008] Step 1: Construct a distributed machine learning system for training the model to be trained, and deploy an on-network fault recovery system in the distributed machine learning system. The distributed machine learning system includes a parameter server, training nodes, and network devices. The network devices include a ToR switch, and the ToR switch is equipped with a recovery storage unit for storing the model recovery parameters in a stateful manner; complete the initialization of the distributed machine learning system.

[0009] Step 2: During the training process, each training node sends the locally calculated gradient to the parameter server. After receiving the gradient data, the parameter server determines the validity of the gradient data based on window constraints and integrity constraints, and only performs aggregation on the valid gradient data to complete the model parameter update for the current round.

[0010] Step 3: The parameter server distributes the updated model parameters to each training node through the ToR switch. When the ToR switch forwards the model parameters, it performs an on-network update operation of the recovery storage unit in the data plane, so that the ToR switch gradually updates and eventually maintains the model recovery parameter copy corresponding to the current training round.

[0011] Step 4: When the parameter server detects a training node failure, it sends a fault recovery trigger signal to the ToR switch corresponding to the faulty training node according to the mapping relationship between the training node and the network topology. The ToR switch is triggered to read the model recovery parameters corresponding to the current training round from the recovery storage unit according to the faulty training node identifier in the signal, and send them to the faulty training node. After receiving the signal, the faulty training node recovers its local model and re-participates in the subsequent training process based on the recovered model state.

[0012] Furthermore, the recovery storage unit is implemented based on the data plane register of the programmable switch.

[0013] Furthermore, the recovery storage unit is organized in the form of a logical array and accessed by an index. The index is determined by the task identifier and sequence number carried in the data packet, which is used to distinguish and store different training tasks and different model recovery parameters.

[0014] Furthermore, the fault recovery trigger signal includes a recovery identifier, a fault training node identifier, and a target switch identifier. The recovery identifier is used to indicate whether a training node fault has occurred, the fault training node identifier is used to indicate the training node that has failed, and the target switch identifier is used to indicate the ToR switch that performs the fault recovery operation.

[0015] Furthermore, the method for determining the validity of gradient data based on window constraints and integrity constraints is as follows:

[0016] By setting a gradient reception cutoff condition to limit the gradient reception range of the current training round as a window constraint, the gradient data that meets the window constraint is the valid gradient data that can participate in the aggregation of this round. During the gradient submission process of the current round, the gradient data submitted by the training node that has not experienced a failure and has completed the gradient submission of the current round is determined to meet the integrity constraint, and is therefore the valid gradient data.

[0017] Furthermore, the gradient reception cutoff condition is a set time threshold or a reception quantity threshold.

[0018] Furthermore, the online update operation is implemented inline within the data packet forwarding path.

[0019] Furthermore, after receiving a data packet, the ToR switch determines whether the data packet is an aggregated data packet. If it is not an aggregated data packet, it forwards it to the next-hop network node; if it is an aggregated data packet, it determines whether the data packet has a recovery flag set.

[0020] If no recovery flag is set, the aggregated data packet is forwarded to each training node; if a recovery flag is set, the process continues to determine whether the current switch is the ToR switch corresponding to the faulty node.

[0021] If the ToR switch is not the one corresponding to the faulty node, the aggregated data packet will be forwarded to each training node. If it is the ToR switch corresponding to the faulty node, the switch will perform a recovery operation and send the data in the recovery storage unit to the faulty training node.

[0022] Beneficial effects:

[0023] This invention aims to address the pain points of inconsistent model states and low recovery efficiency caused by training node failures in distributed model training. It optimizes the system from two dimensions: fault tolerance during training and rapid fault recovery. During training, a fault-tolerant aggregation mechanism based on window constraints and integrity constraints is introduced on the parameter server side to ensure that the model parameters can still be updated in the current round even when some nodes fail, avoiding parameter inconsistencies or training rollbacks caused by incomplete gradients. In terms of fault recovery, model recovery parameters are stored in the ToR switch and dynamically updated with data forwarding, allowing the switch to hold a copy of the parameters for the current round. When a node fails, the switch directly issues recovery data, significantly reducing communication overhead and recovery latency, and substantially improving the system's efficiency and robustness. Attached Figure Description

[0024] Figure 1 This is a schematic diagram of the distributed machine learning system framework used in the on-network fault recovery method for distributed model training provided by the present invention.

[0025] Figure 2 The flowchart of the ToR programmable switch processing in the on-network fault recovery method for distributed model training provided by the present invention. Detailed Implementation

[0026] The present invention will be described in detail below with reference to the accompanying drawings and embodiments.

[0027] This invention provides an in-network fault recovery method for distributed model training. Its core idea is to introduce a fault-tolerant aggregation mechanism into the parameter server, enabling the updating of model parameters for the current training round and the generation of consistent model parameters even when a training node fails. Simultaneously, on the network side, the programmable data plane capability of the ToR switch is utilized to cache and continuously update model recovery-related parameters in-network. This allows the corresponding ToR switch to directly provide the necessary model parameters for recovery to the failed training node when it fails, thereby accelerating the fault recovery process and improving the continuity of the distributed training system.

[0028] This invention provides an in-network fault recovery method for distributed model training, which specifically includes the following steps:

[0029] Step 1: Construct a distributed machine learning system for training the model to be trained, and deploy an on-grid fault recovery system within the distributed machine learning system. The system framework is as follows: Figure 1As shown, the distributed machine learning system includes a parameter server, training nodes, and network devices. The network devices include multiple ToR switches, each of which is equipped with a recovery storage unit for storing model recovery parameters. This recovery storage unit is implemented based on the data plane register of the programmable switch and is used to perform stateful storage of model recovery parameters during model training; thus completing the initialization of the distributed machine learning system.

[0030] The recovery storage units are organized in the form of logical arrays and accessed by index. The index is determined by the task identifier and sequence number carried in the data packet, which is used to distinguish and store different training tasks and different model recovery parameters.

[0031] It should be noted that, Figure 1 The diagram shown is merely an exemplary system architecture of the present invention. The present invention does not limit the specific number of ToR switches, training nodes, and parameter servers, nor does it limit the one-to-one connection method between switches and training nodes. The present invention is also applicable to large-scale distributed training networks containing multiple racks, multiple ToR switches, and even more training nodes.

[0032] Step 2: During the training process of the distributed model, each training node sends the gradients calculated locally to the parameter server. After receiving the gradient data from each training node, the parameter server determines the validity of the gradient data based on window constraints and integrity constraints. It only performs aggregation on the valid gradient data that meets the window constraints and integrity constraints, thus completing the model parameter update for the current round.

[0033] The specific method for determining the validity of gradient data of each training node based on window constraints and integrity constraints is as follows: the gradient reception cutoff condition is set to limit the gradient reception range of the current training round, which serves as a window constraint. Gradient data that meets the window constraint is valid gradient data that can participate in the current round of aggregation. During the current round of gradient submission, the gradient data submitted by the training node that has not experienced a failure and has completed the current round of gradient submission is determined to meet the integrity constraint, and is therefore valid gradient data that can participate in the current round of aggregation.

[0034] Specifically, window constraints can be established by setting a time threshold or a reception quantity threshold as a gradient reception cutoff condition.

[0035] Step 3: The parameter server distributes the updated model parameters to each training node through the ToR switch. When forwarding model parameters, the ToR switch performs an on-network update operation on the recovery storage unit in the data plane. This update process continues as the model parameters are continuously distributed, enabling the ToR switch to gradually update and eventually maintain the copy of the model recovery parameters corresponding to the current training round.

[0036] In particular, the online update operation is completed inline in the data packet forwarding path without the need for additional control interaction, thereby reducing the communication overhead during the data recovery and maintenance process.

[0037] Specifically, the ToR switch processes the received data packets as follows: Figure 2 As shown, the specific method is as follows:

[0038] The ToR switch determines whether the data packet is an aggregated data packet. If it is not an aggregated data packet, it forwards it directly to the next-hop network node. If it is an aggregated data packet, it further determines whether the data packet has a recovery flag set. If no recovery flag is set, the aggregated data packet is forwarded to each training node. If a recovery flag is set, it continues to determine whether the current switch is the ToR switch corresponding to the faulty node. If it is not the ToR switch corresponding to the faulty node, the aggregated data packet is also forwarded to each training node. If it is the ToR switch corresponding to the faulty node, the switch performs a recovery operation, sending the data in the recovery storage unit to the faulty training node, and the process ends.

[0039] Step 4: When the parameter server detects a training node failure, it sends a fault recovery trigger signal to the ToR switch corresponding to the failed training node based on the mapping relationship between the training node and the network topology. After receiving the fault recovery trigger signal, the ToR switch reads the model recovery parameters corresponding to the current training round from the recovery storage unit according to the fault training node identification information carried in the signal, and sends them to the failed training node. After receiving the model recovery parameters, the failed training node recovers its local model and re-participates in the subsequent training process based on the recovered model state.

[0040] The fault recovery trigger signal includes a recovery identifier, a fault training node identifier, and a target switch identifier. The recovery identifier indicates whether a training node fault has occurred, the fault training node identifier indicates the training node that has failed, and the target switch identifier indicates the ToR switch that performs the fault recovery operation.

[0041] Example:

[0042] This embodiment employs an on-network fault recovery method for distributed model training provided by the present invention, achieving on-network acceleration of the fault recovery process for training nodes. The specific process includes:

[0043] S1. Build and initialize the distributed machine learning system, and set up recovery storage units in each ToR switch to store model recovery parameters. The recovery storage units are implemented based on the data plane registers of the programmable switch and are used to store the model recovery parameters in a stateful manner during model training.

[0044] The distributed machine learning system deploys an on-network fault recovery system. This system includes a parameter server, training nodes, and network devices, including multiple ToR switches. Each ToR switch connects to several training nodes in its corresponding rack and communicates with the parameter server via the data center network. To enable on-network maintenance of model recovery parameters, each ToR switch is equipped with a recovery storage unit.

[0045] The recovery storage unit stores the model recovery parameters corresponding to the current training epoch. It can be implemented using a logical array structure with a size of L, where each storage unit corresponds to a different model parameter item. The recovery storage unit can be accessed based on the parameter index information carried in the data packet, thereby enabling the location and updating of model parameters.

[0046] It should be noted that, in some implementations, the recovery storage unit can also distinguish and store model parameters for different training tasks based on task identifiers (such as jobID), thereby supporting parameter isolation and management in multi-task scenarios; the implementation of this invention does not depend on multi-task scenarios, and its core mechanism is also applicable to single training task environments.

[0047] S2. During model training, each training node performs model training based on local data and sends the calculated gradients to the parameter server. The gradient data is forwarded to the parameter server via a ToR switch. During this process, the ToR switch only performs packet forwarding operations.

[0048] S3. The parameter server processes the received gradient data and updates the model parameters based on a fault-tolerant aggregation mechanism. This mechanism combines window constraints and integrity constraints to determine the validity of the gradient data from each training node. The window constraint limits the gradient reception range for the current training round by setting a gradient reception cutoff condition to determine the set of valid training nodes participating in this round of aggregation. The integrity constraint ensures that the gradient contribution of a single training node is a complete and valid contribution; nodes not meeting the integrity condition do not participate in model parameter aggregation. After meeting the cutoff condition, the parameter server generates the model parameters for the current training round based on the gradients of the valid training nodes.

[0049] S4. The parameter server sends the updated model parameters to each training node. During the forwarding process of the model parameter data packet through the ToR switch, the switch performs an on-network update operation on the corresponding position in the recovery storage unit according to the parameter index information carried in the data packet, thereby gradually maintaining the model recovery parameters of the current training round.

[0050] During model training, the parameter server sends the aggregated model parameters to each training node. When the model parameter data packet is forwarded by the ToR switch, the switch performs an on-network update operation on the recovery storage unit in the data plane, so that each ToR switch gradually maintains a copy of the model recovery parameters consistent with the current training round.

[0051] S5. When a training node failure is detected, the parameter server sends a fault recovery trigger signal to the ToR switch corresponding to the failed training node based on the mapping relationship between the training node and the network topology.

[0052] When a training node fails, the parameter server detects the failure and, based on the mapping relationship between the training node and the network topology, sends a fault recovery trigger signal packet to the corresponding ToR switch. The fault recovery trigger signal packet carries a recovery identifier, a faulty training node identifier, and a target switch identifier, and is sent to the corresponding ToR switch through the network.

[0053] After receiving the fault recovery trigger signal, the S6 and ToR switches read the model recovery parameters corresponding to the current training round from the recovery storage unit according to the fault training node identification information carried in the signal packet, and send them to the fault training node.

[0054] S7. After receiving the model recovery parameters, the faulty training node recovers the local model to keep its state consistent with the current training round. After recovery, the training node continues to participate in the subsequent training process.

[0055] In this embodiment, when a data packet arrives at the ToR switch, the type of the data packet is first determined. If the data packet is not a model parameter data packet sent by the parameter server, it is regarded as a gradient data packet or other ordinary data packet sent by the training node. The ToR switch performs a normal forwarding operation and sends it to the next-hop network node without updating the recovery storage unit.

[0056] If the data packet is a model parameter data packet sent by the parameter server, it is further determined whether the data packet carries a recovery identifier. If it does not carry a recovery identifier, it indicates that the data packet is in the normal training phase and is used to distribute the model parameters of the current training round to the training nodes. At this time, while the ToR switch is forwarding the data packet, it performs an update operation on the corresponding position in the recovery storage unit according to the parameter index information carried in the data packet, thereby realizing the on-network maintenance of the model recovery parameters.

[0057] If the model parameter data packet carries a recovery identifier, it indicates that the current data packet is used to trigger the fault recovery process. At this time, the ToR switch further determines whether it is the target ToR switch. If it is not the target ToR switch, it performs normal forwarding processing on the data packet; if it is the target ToR switch, it reads the corresponding model recovery parameters from the recovery storage unit according to the fault training node identifier carried in the data packet, and sends the model recovery parameters to the fault training node, thereby completing the provision of recovery data.

[0058] Through the above processing flow, the ToR switch can update and maintain the model recovery parameters online during the normal training phase, and perform recovery operations according to the trigger signal of the parameter server during the fault recovery phase, thereby achieving rapid online fault recovery of training nodes without affecting the normal training process.

[0059] Furthermore, window constraints are used to determine the set of valid training nodes participating in model parameter aggregation in the current training round. Window constraints are implemented by setting gradient reception cutoff conditions, which include time thresholds or reception quantity thresholds. When the cutoff conditions are reached, the parameter server ends the gradient reception process of the current training round and determines whether each training node should participate in the current round of model parameter aggregation based on its gradient submission status in this stage.

[0060] Furthermore, the parameter server performs integrity checks on the gradient data of each training node, and only determines the gradients that satisfy the integrity constraints of the current training round as valid gradients and participates in model parameter aggregation. The integrity constraints are used to ensure that the gradient contribution of a single training node in the current round is a complete and valid contribution. When a training node fails during the gradient submission process and fails to complete the gradient submission for the current round, its submitted gradients that do not satisfy the integrity constraints will not participate in the model parameter aggregation for this round.

[0061] Furthermore, the fault-tolerant aggregation mechanism is used to update the model parameters of the current training round based on the gradients of the remaining valid training nodes when some training nodes fail. By combining window constraints and integrity constraints, the aggregation operation is only performed on the gradients of training nodes that meet the participation conditions, thereby avoiding the problem of inconsistent model parameters caused by incomplete gradients participating in aggregation and preventing the training process from backtracking.

[0062] Furthermore, the recovery storage unit is used to store model recovery parameters and maintains stateful storage during model training. The recovery storage unit is organized in the form of a logical array and accessed through an index, where the index is determined based on the task identifier jobID and sequence number seq carried in the data packet, which is used to distinguish and store different training tasks and different model recovery parameters.

[0063] Furthermore, the model recovery parameters stored in the recovery storage unit are pre-selected model parameters that significantly influence model training performance. During model training, only some model parameters are updated, while the remaining model parameters remain fixed. By training and recovering only a subset of model parameters while ensuring model training accuracy, the scale of parameters used for fault recovery is reduced, allowing the model recovery parameters to be stored in the programmable switch. The parameter selection strategy can be implemented based on existing methods, and this invention does not limit its specific implementation.

[0064] Furthermore, during the forwarding of model parameter data packets sent by the parameter server, the ToR switch performs an on-network update operation on the recovery storage unit. The update operation is completed inline in the data packet forwarding path without the need for additional control interaction, thereby reducing the communication overhead during the recovery data maintenance process.

[0065] Furthermore, when the parameter server detects a training node failure, it constructs a fault recovery signal packet and sends the packet to the ToR switch corresponding to the failed training node to trigger fault recovery. The fault recovery signal packet includes at least a recovery identifier (recovery) to indicate whether a training node failure has occurred; a failed training node identifier (failed_worker_id) to indicate the failed training node; and a target switch identifier (target_switch_id) to indicate the ToR switch performing the fault recovery operation.

[0066] In summary, the above are merely preferred embodiments of the present invention and are not intended to limit the scope of protection of the present invention. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the scope of protection of the present invention.

Claims

1. A method for on-network fault recovery for distributed model training, characterized in that, Specifically, the following steps are included: Step 1: Construct a distributed machine learning system for training the model to be trained, and deploy an on-net fault recovery system in the distributed machine learning system. The distributed machine learning system includes a parameter server, training nodes and network devices. The network devices include a ToR switch. The ToR switch is equipped with a recovery storage unit for storing model recovery parameters in a stateful manner. Complete the initialization of the distributed machine learning system; Step 2: During the training process, each training node sends the locally calculated gradient to the parameter server. After receiving the gradient data, the parameter server determines the validity of the gradient data based on window constraints and integrity constraints, and only performs aggregation on the valid gradient data to complete the model parameter update for the current round. Step 3: The parameter server distributes the updated model parameters to each training node through the ToR switch. When the ToR switch forwards the model parameters, it performs an on-network update operation of the recovery storage unit in the data plane, so that the ToR switch gradually updates and eventually maintains the model recovery parameter copy corresponding to the current training round. Step 4: When the parameter server detects a failure in a training node, it sends a failure recovery trigger signal to the ToR switch corresponding to the failure training node according to the mapping relationship between the training node and the network topology. The ToR switch is triggered to read the model recovery parameters corresponding to the current training round from the recovery storage unit according to the failure training node identifier in the signal, and send them to the failure training node. After receiving the faulty training node, it restores the local model and re-participates in the subsequent training process based on the restored model state.

2. The on-network fault recovery method according to claim 1, characterized in that, The recovery storage unit is implemented based on the data plane register of a programmable switch.

3. The on-network fault recovery method according to claim 1, characterized in that, The recovery storage unit is organized in the form of a logical array and accessed by an index. The index is determined by the task identifier and sequence number carried in the data packet, and is used to distinguish and store different training tasks and different model recovery parameters.

4. The on-network fault recovery method according to claim 1, characterized in that, The fault recovery trigger signal includes a recovery identifier, a fault training node identifier, and a target switch identifier. The recovery identifier is used to indicate whether a training node fault has occurred, the fault training node identifier is used to indicate the training node that has failed, and the target switch identifier is used to indicate the ToR switch that performs the fault recovery operation.

5. The on-network fault recovery method according to claim 1, characterized in that, The method for determining the validity of gradient data based on window constraints and integrity constraints is as follows: By setting a gradient reception cutoff condition to limit the gradient reception range of the current training round, as a window constraint, the gradient data that meets the window constraint is the valid gradient data that can participate in the aggregation of this round. During the current round of gradient submission, the gradient data submitted by the training node that has not experienced a failure and has completed the current round of gradient submission is determined to satisfy the integrity constraint, i.e., it is valid gradient data.

6. The on-network fault recovery method according to claim 5, characterized in that, The gradient reception cutoff condition is a set time threshold or a reception quantity threshold.

7. The on-network fault recovery method according to claim 1, characterized in that, The online update operation is implemented inline within the data packet forwarding path.

8. The on-network fault recovery method according to claim 1, characterized in that, After receiving a data packet, the ToR switch determines whether the data packet is an aggregated data packet. If it is not an aggregated data packet, it forwards it to the next-hop network node; if it is an aggregated data packet, it determines whether the data packet has a recovery flag set. If no recovery flag is set, the aggregated data packet is forwarded to each training node; if a recovery flag is set, the process continues to determine whether the current switch is the ToR switch corresponding to the faulty node. If the ToR switch is not the one corresponding to the faulty node, the aggregated data packet will be forwarded to each training node. If it is the ToR switch corresponding to the faulty node, the switch will perform a recovery operation and send the data in the recovery storage unit to the faulty training node.

Images