通过子网络数据并行，大幅降低大模型训练显存占用。
Model Parallelism With Subnetwork Data Parallelism
- 将模型分块为子网络，在不同设备上并行训练，不传激活值。
- 显存降低28%-60%，在同等计算量下性能持平或更好。
- 适合资源受限环境下训练大模型的研究者和工程师。
大规模预训练大模型对加速器内存要求极高，且通信开销大。本文提出子网络数据并行（SDP），一种分布式训练框架，将模型划分为结构化子网络，在工作节点间并行训练，无需交换激活值。研究了两种互补的掩码策略：反向掩码仅在反向传播中引入稀疏性，保持梯度无偏；前向掩码在前向传播中也移除参数，提升效率并带来额外正则化。进一步探索了神经元级与块级两种子网络构建策略，适用于Transformer与CNN。在从1B参数的LLaMA在FineWeb上预训练到ResNet-18在CIFAR上的实验中，SDP在保持或提升性能的同时，使每设备显存使用减少28%-60%，且在匹配浮点运算量条件下表现优异。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。