[Paper] PruneX: A Hierarchical Communication-Efficient System for Distributed CNN Training with Structured Pruning

Source: arXiv - 2512.14628v1

Overview

PruneX tackles a bottleneck that’s becoming more common in large‑scale deep‑learning workloads: the limited bandwidth between nodes in multi‑GPU clusters. By tightly coupling a structured pruning algorithm with the communication hierarchy of the cluster, PruneX slashes the amount of data that needs to be exchanged during distributed CNN training, delivering dramatic speed‑ups without sacrificing model quality.

Key Contributions

• Hierarchical Structured ADMM (H‑SADMM): a novel pruning routine that enforces node‑level structured sparsity before any inter‑node synchronization, making the sparsity pattern easy to compress.
• Dynamic buffer compaction: eliminates zero‑valued entries and their indexing metadata, so only the truly needed numbers travel across the network.
• Leader‑follower execution model: separates intra‑node (high‑bandwidth) and inter‑node (bandwidth‑limited) process groups, allowing dense collective operations on already‑compacted tensors.
• System‑level integration: PruneX is built on top of standard data‑parallel frameworks (e.g., PyTorch DDP) and works with off‑the‑shelf GPUs and interconnects.
• Empirical validation: on ResNet‑50/101 across 64 GPUs, PruneX cuts inter‑node traffic by ~60 % and achieves a 6.75× strong‑scaling speed‑up, beating both a dense baseline and a popular Top‑K gradient compressor.

Methodology

1. Structured Pruning at the Node Level

   • Each GPU first runs H‑SADMM, an ADMM‑based optimizer that forces groups of weights (e.g., entire channels or filter blocks) to become exactly zero.
   • Because the sparsity is structured (regular blocks), the remaining non‑zero weights can be stored in a compact dense tensor without needing per‑element indices.

2. Two‑Tier Communication Graph

   • Intra‑node: GPUs within the same physical server exchange full‑precision gradients using the fast NVLink/PCIe fabric. No compression is needed here.
   • Inter‑node: Only the compacted tensors (already stripped of zeros) are sent across the slower network (e.g., InfiniBand). A lightweight “leader” GPU per node aggregates the compacted data, performs the collective (e.g., AllReduce), and then broadcasts the result back to the “followers”.

3. Dynamic Buffer Compaction

   • Before each inter‑node AllReduce, the system scans the gradient buffer, packs the non‑zero blocks into a contiguous buffer, and records the block layout once per iteration.
   • After the collective, the compacted result is unpacked back into the original gradient layout for the local optimizer step.

4. Integration with Existing Training Loops

   • PruneX plugs into the standard training loop as a drop‑in replacement for the torch.distributed backend.
   • The pruning schedule (how aggressively to prune) can be tuned per‑epoch, allowing a gradual transition from dense to highly sparse models.

Results & Findings

• Model accuracy: After the pruning schedule, the final top‑1 accuracy on ImageNet stayed within 0.5 % of the dense reference, confirming that the structured sparsity did not degrade performance.
• Latency breakdown: Inter‑node communication time dropped from ~45 ms per iteration (dense) to ~18 ms (PruneX), while intra‑node synchronization remained unchanged.
• Scalability: The benefit grew with more nodes because the proportion of traffic that traverses the slower inter‑node links increases in larger clusters.

Practical Implications

• Faster training pipelines: Teams can train larger CNNs on existing GPU clusters without upgrading network hardware, cutting both time‑to‑model and cloud compute costs.
• Energy savings: Reducing data movement translates directly into lower power consumption for the network fabric—an often‑overlooked component of the training carbon footprint.
• Simplified deployment: Since PruneX works with standard dense collectives after compaction, developers don’t need to rewrite kernels or maintain separate sparse‑tensor libraries.
• Better model compression: The structured sparsity produced by H‑SADMM is already friendly to downstream inference optimizations (e.g., channel pruning, hardware accelerators), so the same pruning step serves both training efficiency and deployment compactness.
• Compatibility with existing frameworks: By exposing a thin wrapper around PyTorch’s DistributedDataParallel, PruneX can be adopted in CI pipelines with minimal code changes.

Limitations & Future Work

• Applicability beyond CNNs: The current design leverages the regular grid structure of convolutional filters; extending H‑SADMM to transformers or graph neural networks will require new sparsity patterns.
• Static hierarchy assumption: PruneX assumes a clear separation between intra‑node and inter‑node links. Heterogeneous clusters (e.g., mixed‑precision interconnects, varying bandwidth) may need adaptive leader‑selection strategies.
• Pruning overhead: The ADMM solver adds a modest compute cost per iteration (≈2–3 % of total runtime). Future work could explore lighter‑weight structured pruning heuristics or amortized updates.
• Robustness to extreme sparsity: When pruning becomes too aggressive, the compacted tensors shrink dramatically, potentially causing load‑imbalance across nodes. Adaptive sparsity schedules are an open research direction.

PruneX demonstrates that a co‑design of algorithmic sparsity and system‑level communication can unlock substantial gains in distributed deep‑learning training. As model sizes keep growing and network budgets stay tight, approaches like PruneX are poised to become a core part of the production AI stack.

Authors

• Alireza Olama
• Andreas Lundell
• Izzat El Hajj
• Johan Lilius
• Jerker Björkqvist

Paper Information

• arXiv ID: 2512.14628v1
• Categories: cs.DC
• Published: December 16, 2025
• PDF: Download PDF