PyTorch DDP compatibility for multi-machine sparse training

[DarshanFofadiya]

Context

SparseLab's SparseLinear uses standard nn.Parameter for the values array, which means DDP's all-reduce on gradients should work out of the box for the gradient sync.

The open question is topology synchronization. When a DST algorithm (SET, RigL) mutates the sparse topology every N steps, the set of live (row, col) positions can diverge across DDP ranks if not explicitly synchronized. Each rank sees different mini-batches → different gradient magnitudes → different drop/grow decisions.

What needs to happen

1. Verify that basic DDP gradient sync works with SparseLinear (no topology mutation, just Static sparsity). This should work today — confirm it.
2. Design a topology sync protocol for DST algorithms. Options:
   • Rank-0 broadcast: rank 0 computes the topology update, broadcasts the new col_indices to all ranks. Simple, correct, but serializes the mutation decision.
   • Majority vote: each rank proposes drop/grow candidates, keep the ones a majority agree on. More communication, potentially better topology.
   • Independent + periodic resync: let ranks diverge for K steps, then resync to rank 0's topology. Cheapest communication, but theoretically weaker.
3. Implement the chosen protocol as a method on DynamicSparsityAlgorithm.
4. End-to-end test: train MNIST with 2+ processes via torchrun, verify convergence matches single-process.

Files to read

• sparselab/router.py — DynamicSparsityAlgorithm.update() is where topology mutation happens
• sparselab/nn.py — SparseLinear and its parameters
• sparselab/ops.py — autograd Function

Open question

Is there a standard PyTorch idiom for 'parameters that mutate structurally during training' under DDP? NAS (DARTS-style) has a similar problem. If you know of prior art, please comment.