Fokam et al. (2024) — ARXIV-2410.05985
Key takeaways
Introduces PD-ASGD, which decouples the forward and backward passes of backpropagation into separate threads and applies lock-free layer-wise parameter updates as soon as per-layer gradients are available. This simultaneously increases hardware utilization (higher ratio of forward to backward threads) and reduces parameter staleness compared to standard ASGD.
- Bias bound (Section 5.1 / Appendix C). The stationary distribution of PD-ASGD is within a bounded KL divergence from that of standard SGD, where the bound scales as O(tau^2 * eta^2) with staleness tau and learning rate eta. Holds when: Requires staleness-induced noise small relative to mini- batch noise; holds in the continuous-time SDE limit.
- Staleness reduction formula. Layer-wise updates reduce average parameter staleness from tau_block = betaTM to tau_layer = betaT(M-1)/2, a factor of approximately 2/(M-1) improvement. Holds when: M-layer network, backward pass time beta*T per layer; improvement grows with network depth M.
- Empirical speedup. PD-ASGD achieves up to 2.14x wall-clock speedup over LPPSGD and 1.34x over DDP at matched accuracy on CIFAR-10/100 with ResNet architectures. Holds when: 3 A100 GPUs; ResNet-18 and ResNet-34; CIFAR-10 and CIFAR-100.
What the evidence does not cover
- Convergence guarantees are to a stationary distribution rather than exact local optima; the approximation quality depends on staleness-induced noise being small relative to mini-batch noise.
- Layer-wise lock-free writes can cause inconsistent reads mid-forward-pass where some layers see old and some see new parameters; the paper acknowledges this but omits it from the main theoretical model.
- Experiments are limited to vision classification and one sequence task; performance on transformer-based LLM training is untested.
- Implementation requires C++ LibTorch multithreading; integration with standard Python PyTorch distributed training stacks is non-trivial.
Standing in the anthology
Read — the reading is NOTE-150. Arrived in the imported batch, which brought in the asynchronous and bounded-delay branch of distributed training.