CEINR: Critical Element-aware In-Network Recovery for Distributed Machine Learning

Distributed Machine Learning (DML) accelerates training by sharing workloads across multiple workers, but it is vulnerable to worker failures during training. Existing recovery solutions typically restore failed workers from local checkpoints or remote hosts. When local recovery states are unavailable, however, fetching model parameters from remote hosts incurs high recovery latency and heavy network overhead because large amounts of data must travel over long paths. To address this issue, we propose Critical Element-aware In-Network Recovery (CEINR), a programmable-switch-assisted recovery solution for DML. CEINR moves recovery states from remote hosts into the network by maintaining them during normal training and enabling failed workers to retrieve recovery parameters directly from nearby Top-of-Rack switches. To fit limited switch memory, CEINR further introduces a critical element identification mechanism that keeps only essential model parameters updatable while leaving non-critical parameters fixed. Evaluations on AlexNet and ResNet-18 show that CEINR cuts recovery time by up to 93.15% and recovery traffic by up to 96.00% while preserving high model accuracy.