Model Parallelism With Subnetwork Data Parallelism
Fuente:
arXiv
Gespeichert in:
| Hauptverfasser: | , , , , |
|---|---|
| Format: | Preprint |
| Veröffentlicht: |
2025
|
| Schlagworte: | |
| Online-Zugang: | |
| Tags: |
Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
|
| _version_ | 1866916069480857600 |
|---|---|
| author | Singh, Vaibhav Khalid, Zafir Cagnasso, Pietro Oyallon, Edouard Belilovsky, Eugene |
| author_facet | Singh, Vaibhav Khalid, Zafir Cagnasso, Pietro Oyallon, Edouard Belilovsky, Eugene |
| contents | Pre-training large neural networks at scale imposes heavy memory demands on accelerators and often requires costly communication. We introduce Subnetwork Data Parallelism (SDP), a distributed training framework that partitions a model into structured subnetworks trained across workers without exchanging activations. We study two complementary masking regimes: backward masking, which applies sparsity only in the backward step to retain unbiased gradients, and forward masking, which also removes parameters in the forward pass to deliver stronger efficiency gains while providing additional regularization. We further explore two subnetwork construction strategies: neuron level and block level, applied across both transformers and CNNs. In experiments spanning 1B LLaMA pre-training on FineWeb to ResNet-18 on CIFAR, SDP reduces per device memory usage by 28%-60% while maintaining or improving performance under FLOP-matched settings. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2507_09029 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Model Parallelism With Subnetwork Data Parallelism Singh, Vaibhav Khalid, Zafir Cagnasso, Pietro Oyallon, Edouard Belilovsky, Eugene Machine Learning Artificial Intelligence Pre-training large neural networks at scale imposes heavy memory demands on accelerators and often requires costly communication. We introduce Subnetwork Data Parallelism (SDP), a distributed training framework that partitions a model into structured subnetworks trained across workers without exchanging activations. We study two complementary masking regimes: backward masking, which applies sparsity only in the backward step to retain unbiased gradients, and forward masking, which also removes parameters in the forward pass to deliver stronger efficiency gains while providing additional regularization. We further explore two subnetwork construction strategies: neuron level and block level, applied across both transformers and CNNs. In experiments spanning 1B LLaMA pre-training on FineWeb to ResNet-18 on CIFAR, SDP reduces per device memory usage by 28%-60% while maintaining or improving performance under FLOP-matched settings. |
| title | Model Parallelism With Subnetwork Data Parallelism |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2507.09029 |