🤖 machine learning

Model Parallelism With Subnetwork Data Parallelism

本文介绍了子网络数据并行（Subnetwork Data Parallelism, SDP），这是一种分布式训练框架，它将模型划分为结构化的子网络并在各工作节点间进行训练，且无需交换激活值，在多种架构和规模下实现了显著的内存减少（28%-60%），同时保持或提升了性能。

原作者： Vaibhav Singh, Zafir Khalid, Pietro Cagnasso, Edouard Oyallon, Eugene Belilovsky

发布于 2026-06-02

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Vaibhav Singh, Zafir Khalid, Pietro Cagnasso, Edouard Oyallon, Eugene Belilovsky

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图教一个庞大、聪明但极其饥渴的机器人（一个大型 AI 模型）如何说话或识别图像。问题在于，这个机器人需要一个既大又昂贵的厨房，只有极少数人能负担得起建造这样一个厨房。在 AI 的世界里，这个“厨房”就是强大图形卡（GPU）上的计算机内存（RAM）。

这篇论文介绍了一种训练这些巨型机器人的新方法，称为子网络数据并行（Subnetwork Data Parallelism, SDP）。下面通过简单的类比来解释它的工作原理。

旧方法：“全量副本”问题

传统上，为了训练一个大型 AI，研究人员使用一种叫做**数据并行（Data Parallelism, DDP）**的方法。

类比： 想象你有一支由 8 名厨师（GPU）组成的团队，正在尝试烹饪一场盛大的宴会。在旧方法中，每一位厨师都会得到一份完整的、全套的食谱以及所有的食材。他们都烹饪整道菜肴，品尝，然后向彼此大声喊出他们的笔记，以达成共识，从而改进下一轮的食谱。
问题： 这是极其浪费的。每位厨师都需要一个巨大的厨房，才能容纳他们自己那份食谱和食材的副本。如果食谱太大，厨房就会“爆炸”（内存耗尽），导致厨师们根本无法烹饪。

新方法：“专业化团队”（SDP）

作者提出了子网络数据并行（SDP）。与其给每位厨师整本食谱，不如将食谱拆分为特定的章节，并将不同的章节分配给不同的厨师。

类比： 现在，厨师 A 只负责开胃菜，厨师 B 负责汤品，而厨师 C 负责主菜。
- 无需共享整道菜： 至关重要的一点是，他们不需要把整道菜在彼此之间来回传递。他们只讨论自己正在处理的特定食材。
- 厨房缩小了： 因为厨师 A 只需要持有开胃菜的食谱，所以他们的厨房可以小得多。这使得你可以将一场盛大的宴会装进一个更小、更便宜的厨房里。
- 结果： 你可以训练一个更大的机器人（或者训练得更快），而不需要一个超级昂贵的厨房。

两种拆分工作的方式

论文测试了两种将这些“子食谱”分配给厨师的不同方式：

前向掩码（“切分食材”法）：
- 工作原理： 厨师在开始烹饪之前，就直接扔掉了他们不应该使用的食材。他们只烹饪被分配到的那部分菜肴。
- 优势： 这最节省空间，因为他们甚至不需要存储那些未使用的食材。这就像是做汤时，只购买你需要的胡萝卜和洋葱，完全忽略购物清单上的其他内容。
- 代价： 由于他们忽略了部分食谱，他们必须多做几次批次，以确保学到完整的全貌。
后向掩码（“读完整本书，只写你的笔记”法）：
- 工作原理： 厨师阅读了整本食谱（完整模型）以理解上下文，但在记录如何改进菜肴的笔记时，他们只写下自己负责的部分。
- 优势： 这种方式在数学上更“诚实”，因为厨师仍然理解整道菜肴，但他们只在记录笔记时节省空间。这是一种更安全、更稳定的学习方式。

他们发现了什么？

研究人员在两种截然不同的 AI 类型上测试了这些方法：

语言模型（LLMs）： 比如会写文章的“LLaMA”模型。他们尝试训练具有 5 亿和 10 亿个“神经元”（参数）的模型。
- 结果： 他们减少了 28% 到 60% 的内存需求。在某些情况下，他们可以将原本需要巨大厨房的模型放入更小的厨房中，且不会损失 AI 的语言表达能力。
- 加成： 他们发现 SDP 可以完美地与其他的内存节省技巧（如“激活检查点”和“FSDP”）结合使用。这就像堆叠乐高积木；你可以将 SDP 与其他方法结合，将内存占用率大幅降低至 85%。
图像分类器： 识别图片的模型（如 ResNet 和 Swin Transformers）。
- 结果： 他们在标准图像数据集（CIFAR）上训练了这些模型。即使在内存显著减少（有时低至原始容量的 40%）的情况下，该 AI 识别图片的能力仍与全尺寸版本一样出色。在某些情况下，“专业化团队”的方法实际上帮助了 AI 更好地学习，起到了一种“正则化”（一种防止过度思考的方法）的作用。

核心结论

这篇论文并不声称发明了某种新型 AI，也不声称发明了某种用于医院或自动驾驶汽车的新型 AI 用法。相反，它解决的是一个物流组织问题。

把它想象成一种新的建筑施工队组织方式。与其给每个工人一套完整的摩天大楼蓝图（这会占用太多的口袋空间），不如只给每个工人负责的那一层楼的蓝图。他们依然在建造同一栋摩天大楼，但他们需要的口袋空间更小，并且可以用同样的资源建造更高的建筑。

关键要点：

SDP 将一个大型 AI 模型拆分为较小的部分，分布在不同的计算机上。
它消除了每台计算机都需要持有整个模型的必要性，从而节省了 28%–60% 的内存。
它能与现有的 AI 模型（如 LLaMA 和 ResNet）兼容，无需改变其使用方式。
它可以与其他内存节省技巧结合使用，以节省更多空间（高达 85%）。
尽管使用了更少的内存，但 AI 的表现与传统方法一样好（甚至更好）。

技术摘要：子网络数据并行 (Subnetwork Data Parallelism, SDP)

问题陈述

深度神经网络的快速扩展对加速器提出了前所未有的内存需求，通常需要分布式训练策略。虽然数据并行 (DDP) 被广泛使用，但它在每个 GPU 上复制完整的模型，导致高昂的内存开销和昂贵的梯度同步成本。相反，模型并行（如流水线并行、张量并行）将模型拆分到不同设备上以容纳更大的架构，但会引入显著的激活值通信成本，并经常面临流水线气泡（pipeline bubbles）和负载不均衡的问题。

目前的解决方案如全分片数据并行 (FSDP) 和 ZeRO 通过对参数、梯度和优化器状态进行分片来降低单设备内存，但它们仍然会产生大量的通信开销，尤其是在梯度同步期间。在不牺牲收敛质量或不需要高带宽互连的情况下，减少激活值带宽和单节点内存使用量仍是一个亟待解决的挑战。

方法论：子网络数据并行 (SDP)

作者提出了 子网络数据并行 (SDP)，这是一种分布式训练框架，它将模型划分为结构化的子网络，并在不交换激活值的情况下在工作节点（workers）之间进行训练。与按顺序拆分计算的流水线方法不同，SDP 为每个工作节点分配一个结构完整的模型部分（即子网络），该部分保留了从输入到损失的完整路径，从而实现独立的梯度计算。

核心机制

子网络构建：
- 块级 (Block-Level, B-SDP)： 移除整个架构块（例如 Transformer 块或 ResNet 基本块）。
- 神经元/通道级 (Neuron/Channel-Level, N-SDP)： 在全连接层中移除特定的神经元，或在卷积层中移除特定的通道，并确保相邻层之间的一致性。
- 工作节点通过逐步平均（stepwise averaging）来同步重叠的参数。
掩码机制 (Masking Regimes)：
论文研究了两种互补的掩码策略，用以定义训练期间如何利用子网络：
- 前向掩码 (Forward-Masking)： 模型使用掩码后的参数进行评估 ( $m_{fwd} = m$ )。这移除了被掩码组件的参数、激活值和梯度，从而产生显著的内存和计算节省。然而，梯度反映的是掩码后的前向传播过程。
- 后向掩码 (Backward-Masking)： 前向传播使用完整模型 ( $m_{fwd} = m_{uni}$ )，而稀疏性仅应用于反向传播和聚合阶段 ( $m_{bwd} = m$ )。这保留了无偏的梯度估计（具有理论依据），并提供了一个原则性的基准，尽管它主要在梯度和优化器状态而非激活值上节省内存。
理论基础：
作者提供了在 $L$ -光滑设置下后向掩码机制的收敛性分析。他们证明了收敛速率受“掩码误差”控制，该误差定义为均匀掩码与特定掩码分布之间的距离。分析将收敛质量与掩码图的谱间隙（spectral gap）联系起来，表明通过周期性平均实现的重叠参数分配可以维持部分同步。

核心贡献

新范式： 引入 SDP 作为一种内存高效的分布式训练范式，将子网络训练分布到各个节点，在保持或提高性能的同时，将单设备内存降低了 28%–60%。
理论保证： 提供了一个将后向掩码 SDP 的收敛性与掩码图上的谱间隙条件联系起来的理论基础，证明了在特定掩码条件下可以保持无偏梯度。
可组合性： 证明了 SDP 与现有的节内存技术是正交的。它通过物理移除模型组件，同时缩小了参数、梯度、优化器状态和激活值的规模。这使得 SDP 可以与 FSDP (ZeRO-3) 和 激活值检查点 (Activation Checkpointing, AC) 进行清晰的组合。
泛化能力： 在包括大型语言模型 (LLaMA) 和图像分类模型 (ResNet, Swin Transformer) 在内的多种架构上验证了 SDP，表明它是一种“掉入式”（drop-in）的常规数据并行替代方案。

实验结果

大语言模型 (LLaMA)

实验在 FineWeb 数据集上针对 134M、500M 和 1B 参数量的 LLaMA 模型进行。

内存效率： 在 1B 规模下，具有 69% 活跃组件 ( $C=0.69$ ) 的 B-SDP 与 DDP 相比，降低了 28% 的峰值 GPU 内存，同时实现了相等或更低的验证损失。在 500M 规模下，具有 50% 活跃组件的 B-SDP 降低了 40% 的内存。
下游性能： 在五个下游基准测试（ARC-E, BoolQ, HellaSwag, OBQA, SciQ）上，所有 SDP 变体在宏平均指标上与 DDP 的表现差距在 ~1 个百分点以内，部分变体（B-SDP）在特定任务上甚至显示出轻微的提升。
可组合性： 将 B-SDP 与 FSDP 及激活值检查点结合使用，使 1B 模型的单设备峰值内存使用量达到 9.7 GB，相比 DDP (70.9 GB) 减少了 85%，且严格低于仅使用 FSDP+AC 的内存 (11.0 GB)。
通信： SDP 将每步的通信量减少了约 28%（移动 2,253 MB，而 DDP 为 3,143 MB）。在跨节点设置中，B-SDP 与 FSDP 的组合比 DDP 缩短了 19.9% 的实际训练时间（wall-clock time）。

图像分类

实验在 CIFAR-10 和 CIFAR-100 上使用 ResNet-18、WideResNet-18 和 Swin-Tiny 进行。

ResNet-18： SDP 配置仅使用 DDP 内存预算的 40% 即可达到具有竞争力的或更高的准确率。N-SDP（神经元级）表现出强大的正则化效果，在 64% 内存使用率下优于 DDP。
Swin Transformer： B-SDP 在使用 68% 内存的情况下，在 CIFAR-10 上达到了与 DDP 相当的准确率，并在 CIFAR-100 上以 32% 的内存减少实现了 2% 的准确率提升。
稀疏性极限： 虽然前向掩码（B-SDP/N-SDP）在适度稀疏时表现良好，但 后向掩码 (Bb-SDP) 在极高稀疏度水平（例如 $C=3/8$ ）下表现更为优异，防止了其他变体中出现的模型崩溃现象。

重要性与主张

论文声称 SDP 提供了一个实用的“第三维度”并行，与数据并行和模型并行互补。其主要意义在于其能够：

解耦内存与带宽： 通过消除工作节点之间交换激活值的需求，SDP 缓解了经常限制流水线并行和张量并行的带宽瓶颈。
实现 FLOP 匹配的扩展： SDP 允许在相同的硬件预算下训练更大的模型或更长的序列。单次迭代收敛速率的轻微下降，可以通过增加迭代次数（由于减少了单步计算和通信，从而实现 FLOP 匹配）来抵消。
提供掉入式解决方案： SDP 作为标准数据并行的掉入式替代方案，除了掩码逻辑外，无需对底层分布式基础设施进行任何更改。

作者强调，在内存和带宽成为主要约束的场景下，SDP 特别具有价值，它为在内存受限的加速器上训练大型模型提供了一条可扩展的路径，且无需复杂的流水线调度或全分片的通信开销。