Optimizing Dynamic Neural Networks with Brainstorm

文章标题: 使用 Brainstorm 优化动态神经网络
作者/机构: Weihao Cui (上海交通大学), Zhenhua Han (微软亚洲研究院), Lingji Ouyang (中国科学技术大学), Yichuan Wang (上海交通大学), Ningxin Zheng (微软亚洲研究院), Lingxiao Ma (微软亚洲研究院), Yuqing Yang (微软亚洲研究院), Fan Yang (微软亚洲研究院), Jilong Xue (微软亚洲研究院), Lili Qiu (微软亚洲研究院), Lidong Zhou (微软亚洲研究院), Quan Chen (上海交通大学), Haisheng Tan (中国科学技术大学), Minyi Guo (上海交通大学)

A1 主要贡献

本文针对动态神经网络（Dynamic Neural Networks, DNNs）的优化问题，提出了一个名为 Brainstorm 的深度学习框架。

核心问题：现有的深度学习框架主要为静态神经网络设计，其算子执行顺序对所有输入都是确定的。然而，动态神经网络在推理时会根据输入稀疏地激活不同的子网络，这种动态性带来了新的优化机会。现有框架无法抓住这些机会，关键在于表达粒度的错位（misaligned expression granularity）：框架只能在张量（tensor）级别追踪数据流，而动态性（如 Mixture-of-Experts (MoE) 模型中的“令牌”分发）通常发生在子张量（sub-tensor）级别。因此，现有编译器无法理解和追踪这种细粒度的数据流，也就无法收集用于动态优化的运行时配置文件。

研究目标：构建一个能够有效优化动态神经网络执行的框架，通过统一动态性的表达方式，使其易于被追踪和分析，从而利用运行时动态性分布来专门化模型执行。

创新点与主要贡献：
1. 识别了新的优化空间：利用动态神经网络中激活分布的非均匀性，通过收集动态性的统计配置文件（statistical profiles）来专门化模型执行，从而开辟了一个新的动态优化空间。
2. 指出了核心挑战：明确了优化动态神经网络的主要障碍是现有深度学习框架中张量级编程模型与追踪所需的细粒度数据流之间的粒度错位。
3. 提出了统一的编程抽象：
* Cell：一种新的数据抽象，允许模型开发者描述动态性发生的数据粒度（例如，张量内的一个令牌或一个图像块）。
* Router：一个统一的接口，让开发者表达 Cell 应如何被动态分派到不同的计算分支。Brainstorm 负责高效执行这些路由行为。
这一设计使得 Brainstorm 能够以正确的粒度收集细粒度数据流的运行时配置文件。
4. 提出了四种动态优化策略：基于可追踪的 Cell 级别数据流分析，实现了多种动态优化：
* 动态水平融合（Dynamic Horizontal Fusion）：分析路由到各分支的 Cell 数量，将多个并行分支融合成针对高频负载优化的 GPU 核函数。
* 配置文件引导的模型放置（Profile-Guided Model Placement）：通过跨层的 Cell 级别分析，优化并行分支的分布式放置，以最小化跨 GPU 的通信。
* 推测式路由（Speculative Routing）：利用分支激活的配置文件，推测性地启动分支算子，以隐藏路由开销。
* 推测式权重预加载（Speculative Weight Preloading）：推测性地预加载分支权重，以节省 GPU 内存。

通过在 PyTorch 上实现 Brainstorm，并对6个主流动态神经网络进行评估，实验结果表明，与现有先进解决方案相比，Brainstorm 带来了最高 11.7倍 的加速（平均 3.29倍），或减少了 42% 的内存消耗。


图 1: 动态神经网络在令牌级、补丁级和像素级进行路由的示例。

A3 背景知识与动机

动态神经网络（Dynamic Neural Networks）。为了模仿人脑的工作方式，机器学习社区积极研究动态神经网络的设计。各种动态性被提出来以适应不同输入的模型结构和参数。图1展示了动态神经网络的代表性模式。构建动态神经网络最常见的方式是使用路由机制，自适应地将（部分）输入分派到不同的子网络。一个通常被称为路由器的功能，会预测输入值应该经过哪个子网络。许多路由策略已被提出用于不同任务，例如top-k路由器【索引3，Outrageously large neural networks: The sparsely-gated mixture-of-experts layer，2017，arXiv】。不同分支中的子网络可能具有不同的权重、架构或参数数量，以更好地适应被路由的输入。例如，MoE网络训练并行的专家网络，并将输入令牌分派到不同的专家子网络中，每个子网络都期望专门处理特定类别的输入【索引14，Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity，2022，JMLR】、【索引15，Gshard: Scaling giant models with conditional computation and automatic sharding，2020，arXiv】。ClassSR【索引10，Classsr: A general framework to accelerate superresolution networks by data characteristic，2021，CVPR】根据超分辨率的难度将图像块路由到异构分支。Skip-Conv【索引23，Skip-convolutions for efficient video processing，2021，CVPR】将新的像素路由到计算中，并跳过前一帧的重复像素。模型开发者通常使用一个张量来存储多个令牌/图像块/像素，并使用像einsum【索引24，Zur Elektrodynamik bewegter Körper. (German) [On the electrodynamics of moving bodies]，1905，Annalen der Physik】这样的数据移动算子来编程子张量级别的动态性。

动态优化机会。在编程语言领域，利用程序动态性的统计配置文件进行即时（JIT）优化已经被广泛研究【索引25，Profile guided optimization in llvm，2022，http://clang.llvm.org】、【索引26，Autofdo: Automatic feedback-directed optimization for warehouse-scale applications，2016，CGO】，例如HotSpot JVM会推测性地裁剪掉在收集的运行中从未执行过的路径【索引16，Java hotspot vm，2022，http://oracle.com】。然而，现有深度学习框架的优化主要集中在静态神经网络上，它们错过了由神经网络动态性带来的大量动态优化机会。图2展示了四种动态神经网络的路由分布。图2a和图2b分别是将令牌和图像块分派到不同分支的动态神经网络。我们观察到它们的令牌/图像块分布是不平衡的：一些分支接收的数据明显多于其他分支。这为调整高效的GPU核函数以适应其形状与负载分布提供了机会，这可能带来超过10倍的加速。此外，这些并行分支可以被水平融合以进行并发执行（§4.1）。我们还通过分析多层相关性的统计数据发现了优化机会。图2c展示了TaskMoE【索引27，Beyond distillation: Task-level mixture-of-experts for efficient inference，2021，arXiv】的多层相关性，即从第0层的某个专家路由到第1层另一个专家的令牌比例。我们发现连续两层的分支激活是相关的，例如，在第0层的专家0被激活后，第1层的专家14/15有很高的概率被激活。通过将相关的专家共同部署在同一个GPU上，最多可以节省87%的GPU间通信（§4.2）。图2d显示了DynamicRouting【索引28，Learning dynamic routing for semantic segmentation，2020，CVPR】中选定路由器的分支激活情况，该模型有186个路由器，训练用于将图像转发到三个分支中的一个或两个。我们的测量显示，它花费了超过44%的时间在路由上。然而，许多路由器在不同运行时倾向于激活相同的分支，呈现出有偏的分布。例如，路由器3有很高的概率选择分支1和分支2。这为推测性执行创造了机会，例如，跳过路由计算以减少路由开销（§4.3），或机会性地将权重预加载到GPU内存中（§4.4）。此外，我们发现许多动态神经网络可以同时被多种动态优化方法优化。这些优化的关键要求是能够在动态性发生的粒度上收集统计配置文件，而这在现有的深度学习框架中尚未被探索。

编程模型错位。编程模型的错位是现有框架追踪动态性配置文件的主要障碍。如图1所示，语言任务通常在输入句子的令牌粒度上进行路由；视觉任务在输入图像的图像块上进行路由；视频模型根据帧间相似性部分重用先前的像素。所有动态性都发生在句子、图像或帧的张量内部。现有框架使用静态数据流图来优化模型，该图只表达了张量和算子之间的关系。它们无法在运行时收集必要的配置文件。没有模型开发者的明确指定，它们无法理解什么是令牌以及它们是如何被动态分派的，更不用说追踪如图2c所示的复杂令牌级数据流。此外，张量级编程只能应用算子级优化（例如，算子融合），而无法优化更细粒度的数据移动或计算。这些挑战促使Brainstorm提出一个原则性的设计，让模型开发者暴露需要被追踪的信息，并利用收集到的配置文件进行动态优化。

图 2: 四种动态神经网络的路由分布。

A2 方法细节

3. Cell 和 Router 作为核心抽象

统一模型表达。为了让模型开发者能以可追踪的方式表达动态神经网络，Brainstorm 使用 Cell 和 Router 统一了模型表达，以在正确的粒度上构建动态性。

Cell：定义动态性粒度。为了让模型开发者定义动态性发生的数据粒度，Brainstorm 通过一个名为 Cell 的数据抽象来增强传统张量。Cell 是在多个分支间进行动态分派的基本单位。模型开发者可以使用 brt.annotate_cell API 来标注任何张量，以指定张量中 Cell 的粒度（brt 是 Brainstorm 的包名）。

brt.annotate_cell(tensor, dims, shape)

模型开发者需要指定在哪个维度（dims）和以何种形状（granularity）的值进行路由。图3展示了三个分别在令牌、图像块和像素粒度上路由 Cell 的例子。第一个例子路由一个包含三个令牌的张量，这些令牌位于第0维（dims=(0)），每个令牌由一个768个浮点数值的向量表示（shape=(1,768)）。第二个和第三个例子分别在2D图像张量（dims=(0,1)）中路由32x32的图像块（shape=(32,32)）和1x1的像素（shape=(1,1)）。

Router：统一动态分派接口。为了动态分派 Cell，Brainstorm 引入了一个统一的 Router API，它通过 router_fn 支持自定义规则，以决定 Cell 在多个分支间的动态放置。Router 和 router_fn 的 API 定义如下²。

class Router(nn.Module):
    def __init__(self, router_fn): ...
    def forward(self, tensor, **kwargs): ...
def router_fn(tensor, **kwargs) -> Routes: ...

router_fn 将标注了 Cell 的张量作为输入，并生成一个特殊的张量 Routes，其值表示 Cell 应该去往哪个分支。Routes 的形状与源张量中 Cell 的布局相同。例如，图3中的第二个例子有 6×4 个图像块，因此 router_fn 也应生成 6×4 个 Routes。在做路由决策时，可以在 kwargs 中设置辅助输入。在模型的前向传播过程中，Router 将输入张量送入 router_fn 以获取 Cell 的路由决策，然后将 Cell 分派到相应的分支。将现有的动态神经网络代码移植到 Brainstorm 很容易，例如，我们仅修改了12行代码就将 SwitchTransformer【索引14，Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity，2022，JMLR】的官方 PyTorch 实现移植到了 Brainstorm。

解耦控制流与执行。Brainstorm 的 Router 抽象将决定 Cell 应如何动态分派的控制流与其执行解耦。最优的执行策略根据运行时配置文件而大相径庭。Brainstorm 使模型开发者从挑战性的执行优化中解放出来。他们只需要专注于设计路由逻辑，而将执行优化留给 Brainstorm。router_fn 给出的 Routes 由 JIT Profiler 收集以获取统计配置文件。Brainstorm 的动态优化会分析这些统计数据，以找到最高效的执行策略（§4）。

高效的路由操作。Router 背后是一系列高效的 GPU 操作，以实现 router_fn 指定的路由行为。当接收 Cell 的分支位于同一个 GPU 上时，Brainstorm 使用一个高效的数据重排 GPU 核函数来生成多个张量，每个张量包含路由到对应分支的 Cell。与现有解决方案大量使用计算算子（如 einsum）进行细粒度动态数据重排不同，Brainstorm 使用 GPU 核函数直接移动数据以避免不必要的计算。当 Cell 分布到多个 GPU 时，Brainstorm 有一个稀疏通信原语来高效地分散和收集 Cell。与现有深度学习框架中常用的 all-to-all 原语【索引22，Tutel: Adaptive mixture-of-experts at scale，2022，arXiv】、【索引29，Pytorch: An imperative style, high-performance deep learning library，2019，NeurIPS】相比，当 Cell 不均匀地路由到多个 GPU 时，Brainstorm 的稀疏通信更高效，因为它避免了因填充（padding）而产生的不必要通信（详见§6的实现细节）。

图 3: 分别在令牌级、补丁级和像素级路由 Cells 的示例。router_fn 生成路由决策，指示 Cells 应路由到哪个分支 ID（-1 表示丢弃），这些决策由 JIT 分析器收集用于动态优化。

与带控制流的中间表示（IR）的比较。与现有深度学习框架将控制流和数据流混合在一起的中间表示（IR）不同，Brainstorm 选择了 Router 的解耦设计。Brainstorm 的数据流图将 Router 的复杂控制流隐藏在 router_fn 之后。一个 Router 可以被看作是一个数据分发算子，动态地将张量的 Cell 分派到多个分支。这极大地简化了 Cell 级数据流的追踪和分析，因为编译器不再需要从数据流图中分离出与动态性相关的算子，这对于深度学习框架来说是困难的【索引30，Janus: fast and flexible deep learning via symbolic graph execution of imperative programs，2019，NSDI】、【索引31，Terra: Imperative-symbolic co-execution of imperative deep learning programs，2021，NeurIPS】、【索引32，Pytorch: An imperative style, high-performance deep learning library，2019，NeurIPS】。实际上，相比于了解路由逻辑是如何构建的，编译器更需要知道关于路由决策的统计信息，而这足以被 Brainstorm 的 Router 捕获。

增强现有控制流。此外，Brainstorm 还用 Cell 级路由能力增强了现有 IR 中的控制流算子。Brainstorm 的 Router 本身可以被看作是一个 switch-case 算子，将 Cell 路由到不同的分支以有条件地应用不同的函数。结合 while-loop 算子，一个动态神经网络可以将一些 Cell 路由回循环入口进行下一次迭代，并将其他的 Cell 送到输出，这在语言任务的自回归解码中很常用。

4. 动态优化

基于配置文件的优化。Brainstorm 分析收集到的程序执行配置文件来提升运行时性能。与分析程序函数或代码块调用的传统动态优化不同，优化动态神经网络的关键是分析和利用 Cell 级数据流，以使模型执行专门适应运行时的动态性分布。本节介绍我们为动态神经网络确定的四种动态优化。利用 Brainstorm 的 Cell 和 Router 抽象，未来可能实现更多优化。表1列出了进行每种动态优化所需的统计信息。
Table 1: Brainstorm中不同动态优化策略使用的统计信息。

4.1 动态水平融合

基本概念。水平融合是一种编译器优化，它将模型的并发分支融合成一个融合算子，以提高 GPU 计算单元（CU）的利用率并减少启动开销。现有方法【索引33，Rammer: Enabling holistic deep learning compiler optimizations with rtasks，2020，OSDI】、【索引34，Horizontally fused training array: An effective hardware utilization squeezer for training novel deep learning models，2021，MLSys】不适用于动态神经网络，因为它们假设一个静态数据流图，其所有分支都被相同的输入激活。Brainstorm 引入了一种动态水平融合优化，支持动态和稀疏激活的分支，使它们可以在 GPU 上同时执行。

实现机制。特别地，正如我们在图2中所示，动态神经网络的 Cell 分布可能非常不平衡。即使批量大小很大，通过动态水平融合接收少量 Cell 的分支，仍然可以加速模型执行。Brainstorm 利用从 Router 收集的配置文件来提取每个分支的统计负载，即有多少 Cell 被路由到每个分支。Brainstorm 找到 Cell 负载分布的多个百分位点（例如，50%、90%、100%），并为这些形状调整 GPU 核函数。所有调整后的核函数被融合成一个算子。在推理时，Brainstorm 将每个分支的输入填充到最接近的已调整核函数。这需要在运行时追踪动态的 Cell 级数据流，我们将在§5.2中解释 Brainstorm 是如何实现的。请注意，动态融合的 GPU 核函数只使用已激活分支的权重，无需将所有分支的权重加载到 GPU 内存中。

示例说明。图4展示了一个在四个并行分支间路由112个 Cell 的例子。只有其中三个分支（在运行时才知道）被激活。在水平融合之前，三个激活的分支必须顺序执行，这可能无法充分利用 GPU CU。在将所有分支融合成一个 GPU 核函数后，GPU 可以以更高的 CU 利用率同时执行激活的分支。每个分支都使用填充最少的已调整核函数执行，以实现最高效的执行。例如，融合后的核函数包含两个为32个 Cell 和64个 Cell 调整的 Conv 3x3 核函数，网络中的前两个 Conv 3x3 分支通过分别填充4个和2个 Cell 来使用它们。

图 4: 多个分支的算子被水平融合成一个核函数。只有被激活的算子会被执行。每个分支使用最接近的已调优核函数以实现最小的填充。

4.2 配置文件引导的模型放置

跨层相关性。人脑的大脑皮层被组织成不同的区域，其中实现某一功能的神经元位置相近【索引35，Connectivity-driven white matter scaling and folding in primate cerebral cortex，2010，PNAS】。通过分析统计路由决策，我们在人工设计的动态神经网络中观察到了类似的效果。如图2c所示，来自两层的专家有很高的概率一起被激活。这些高度相关的专家之间的 Cell 级通信量高于其他专家。图5展示了一个例子，通过分析多层相关性，Brainstorm 可以将相关的子网络共同放置在同一个 GPU 上，以减少 GPU 间的通信。请注意，除了在运行时收集的动态 Cell 级数据流外，多层相关性还需要分析静态 Cell 级数据流以推断正确的放置约束。我们在§5.1中的分析表明，一个句子张量的每个 Cell 都依赖于前一个 MoE 层的所有 Cell。这意味着存在一个放置约束，即一个句子的所有 Cell 都应被收集到同一个 GPU 上，以便其自注意力算子能生成正确的输出。这提出了一个挑战，需要同时进行动态和静态 Cell 级数据流分析来理解 Cell 的层间相关性。我们在§5.1中解释了 Brainstorm 的静态 Cell 级数据流分析。

图 5: 配置文件引导的模型放置。该示例显示，默认放置有 90% 的跨 GPU 流量，而优化后的放置将其减少到 10%。

单层分布。除了跨层分析，我们发现像图2a所示的单层 Cell 分布也可以帮助模型放置。一些分支可能比其他分支接收更多的 Cell。可以将重负载分支与轻负载分支共同放置，以平衡整体通信，避免在某些 GPU 上出现停顿。

4.3 推测式路由

问题与机会。模型开发者通常构建涉及控制流的路由逻辑，这可能需要CPU处理并产生CPU-GPU同步开销。与它们的理论性能（基于FLOPs）相比，路由开销可能主导推理延迟。我们的测量显示，MSDNet【索引1，Multiscale dense networks for resource efficient image classification，2017，arXiv】和DynamicRouting【索引28，Learning dynamic routing for semantic segmentation，2020，CVPR】分别花费了65%和44%的时间在路由上。我们发现这些模型在推理时选择分支通常具有有偏的概率。我们对Brainstorm的Router配置文件的分析显示，许多Router是高度可预测的。Brainstorm仅通过选择最常出现的分支，就能以超过90%的准确率预测DynamicRouting【索引28，Learning dynamic routing for semantic segmentation，2020，CVPR】的决策（§7.4.6）。

实现机制。如图6所示，Brainstorm可以预先预测Router的路由决策（基于统计配置文件）并跳过router_fn以隐藏路由开销。为了保证正确性，Brainstorm使用一个并行线程来检查router_fn的结果。当发生误判时，模型执行将被回滚以重新执行正确的分支，误判开销可以忽略不计（§7.3）。

图 6: 推测式路由：跳过路由计算并推测性地启动概率最高的分支；当预测错误时，回滚并重新启动到正确的分支。

4.4 推测式权重预加载

问题与机会。为了在有限的GPU内存上运行大型模型的推理，通常需要GPU内存和主机内存之间交换层的权重以减少GPU内存需求【索引36，Nvidia cuda: Unified memory programming，2022，http://docs.nvidia.com】。为了隐藏内存迁移延迟，现有解决方案需要知道层的执行顺序，以便在执行前一层时以流水线方式预加载必要的权重【索引37，Swapadvisor: Pushing deep learning beyond the gpu memory limit via smart swapping，2020，ASPLOS】、【索引38，Capuchin: Tensor-based gpu memory management for deep learning，2020，ASPLOS】。然而，动态神经网络没有静态的层执行顺序。动态激活分支的执行只有在做出路由决策时才知道。这使得现有解决方案难以预加载动态层的权重。

实现机制。如图7所示，类似于推测式路由，Brainstorm利用分支激活分布的统计配置文件，推测性地预加载那些有高概率被激活的分支的权重。当预测性预加载未命中时，它会回退到按需加载，开销可以忽略不计（§7.3）。

图 7: 推测性地预加载概率最高的权重；当预测错误时，回退到按需加载。

5. 追踪 Cell 级数据流

追踪的必要性。为了实现§4中的优化，理解 Cell 如何在网络中传输至关重要，这样编译器才能利用 Cell 级数据流来优化模型执行。在动态神经网络中，存在两种类型的 Cell 级数据流：（1）存在于大多数静态算子（如Conv2D）中的静态数据流，对所有输入都是固定的；以及（2）由 Router 在运行时决定的动态数据流。前者用于理解 Cell 在静态层间的关系；后者用于识别 Cell 在分支间的路由。

5.1 静态 Cell 级数据流

追踪方法：符号执行。以张量为中心的数据流图只保留了张量之间的关系，而没有 Cell 的信息。为了追踪静态算子的所有可能的 Cell 级数据流，Brainstorm 在预编译（ahead-of-time compiling）阶段使用 Cell 级的符号执行来提取更细粒度的关系。对于一个标注了 Cell 的张量，Brainstorm 会初始化该张量的符号版本，其中 Cell 是符号。属于同一个 Cell 的张量值共享相同的符号。Brainstorm 利用算子的张量表达式（在深度学习编译器中广泛使用【索引39，Tvm: An automated end-to-end optimizing compiler for deep learning，2018，OSDI】）来构建算子的计算逻辑。通过检查符号计算的结果，Brainstorm 能够理解 Cell 在静态算子中是如何传输的。

示例分析。图8a展示了三个例子，一个包含多个 Cell 的张量与一个常量矩阵进行矩阵乘法。该张量有两个标注为 A 和 B 的 Cell。第一个例子保留了 Cell 的位置；第二个例子重排了 Cell；第三个例子在输出中混合了所有的 Cell。这个例子表明，当张量值不同时，静态 Cell 级数据流可能会变化。张量级数据流分析很难获得这种更细粒度的关系。图8b展示了两个 MoE 层之间的自注意力算子的静态 Cell 级数据流。因为在自注意力算子中，两个张量之间存在矩阵乘法，并且两个张量都包含 $X_i$ 的 Cell，所以这个自注意力算子会混合来自输入 X 的所有 Cell 来生成输出 Y。通过对 Cell 进行符号执行，我们可以推导出 X 和 Y 中 Cell 之间的关系，即 Y 中的每个 Cell 都源自于 X 中的所有 Cell。

分析的重要性。静态 Cell 级数据流分析对于推导 Cell 的跨层关系是必要的，这在与数据移动相关的优化中非常重要。它允许 Brainstorm 在 Cell 级别探索数据移动，打破了在优化多 GPU 执行时张量级数据移动的限制。例如，如果 Cell 只是被重排而没有混合（例如，图8a中的前两种类型），框架在基于数据局部性在多个 GPU 之间分派 Cell 时就有更大的自由度，以获得更好的性能。对于基于 MoE 的模型，因为令牌在自注意力层中被混合了，这引入了一个约束，要求在进行自注意力之前将一个句子的所有令牌聚合到同一个 GPU 上以推导出输出。正如我们在§4.2中所示，这个要求为优化中如何动态放置 Cell 创造了约束，而这只有在分析了静态 Cell 级数据流之后才知道。

图 8: Cell 级别的不同类型的静态数据流。

5.2 动态 Cell 级数据流

追踪方法：Router 抽象。在 Brainstorm 中，模型开发者使用 Router 来表达动态性。路由逻辑定义在 router_fn 中，它在运行时生成 Cell 的路由决策。Brainstorm 的 Router 抽象使得追踪必要信息变得容易。与传统编程语言的动态优化类似，Brainstorm 专注于收集路由决策的统计配置文件，而不关心它们是如何生成的。

实现机制。如果启用了 Cell 级分析，每当一个 Router 被调用时，Brainstorm 会将其路由决策记录到一个缓冲区中。Brainstorm 有一个单独的线程将该缓冲区流式传输到一个配置文件中。Brainstorm 支持多级分析。一些优化只需要 Router 的局部统计配置文件（例如，Cell 的分支负载）。一些优化需要跨多层的 Cell 级数据流，因此需要直接转储原始决策。作为控制信号，路由决策比动态神经网络中的其他数据张量要小得多。我们在§7.3的评估表明，分析开销可以忽略不计。

6. 实现

代码概览。我们在 PyTorch 上实现了 Brainstorm，代码量约为13,000行：3,000行用于 Brainstorm 核心抽象，3,000行用于动态优化，3,000行 C++ 代码用于核函数调度和稀疏 Cell 通信，以及1,500行用于支持动态优化的自动转换。

系统架构。图9总结了 Brainstorm 的架构。除了现有框架中广泛使用的 Tensor 和 Operator，Brainstorm 引入了 Cell 和 Router 以统一的抽象来表达动态神经网络（§3）。编程的动态神经网络将由编译器通过静态和动态优化进行优化（§4）。Brainstorm 的动态优化需要静态和动态的 Cell 级数据流分析（§5）。Brainstorm 首先以预编译的方式推断静态算子中的静态 Cell 级数据流（§5.1）。在执行编译后的模型时，JIT 分析器会收集 Router 的配置文件以进行进一步的动态数据流分析（§5.2）。

图 9: Brainstorm 的系统架构。阴影部分是 Brainstorm 为动态神经网络引入的组件。

高效的 Cell 路由。Brainstorm 负责动态的 Cell 分派，该分派过程感知应用的动态优化，让模型开发者专注于设计路由算法。对于单GPU上的 Cell 路由，我们使用一个自定义的GPU核函数根据路由决策在张量内部重新排列 Cell。我们借鉴了Tutel【索引22，Tutel: Adaptive mixture-of-experts at scale，2022，arXiv】针对MoE模型的思想，通过一个自定义的GPU核函数并行地为所有分支重新排列 Cell。但我们的实现是通用的，适用于所有动态神经网络，而不仅仅是MoE模型。此外，我们的实现能够感知应用的动态优化。例如，一个动态水平融合的算子可能包含不同大小的GPU核函数，因此需要可变的填充。对于跨多个GPU的 Cell 路由，我们提供了一个更灵活的稀疏通信原语。如图10左侧所示，模型开发者通常结合密集的all-to-all原语和置换操作来进行分布式 Cell 路由。其效率受限于平衡的路由。通过Brainstorm的稀疏通信，它只传输 Cell 而没有额外的填充。稀疏通信的底层实现是一系列点对点通信。然而，它可以适应动态优化的要求，并提供最高效的通信机制。

图 10: 用于分布式 Cell 路由的稀疏 All-to-All。它节省了额外填充带来的冗余通信。

核函数融合的过多候选。Brainstorm将多个分支融合成一个核函数，每个分支包含几个潜在的候选。在运行时，Brainstorm根据分派的Cell触发合适的候选。然而，从分析中得出的过多核函数候选会导致在使用自动调优工具【索引39，Tvm: An automated end-to-end optimizing compiler for deep learning，2018，OSDI】搜索它们时产生相当大的时间开销。为了避免这种情况，Brainstorm只融合每个分支有限的一组候选。同时，如果融合的分支是同构的（即算子相同，只有权重不同），核函数候选会在分支之间共享。例如，由于SwitchTransformer的专家网络使用相同的前馈层，Brainstorm只需要六个候选核函数就可以优化每层256个专家的执行（§7.4.1）。

优化遍（Optimization Passes）。Brainstorm 中的大多数自动转换都是用 torch.fx 实现的。利用 torch.fx 追踪的数据流图，Brainstorm 使用从 Router 收集的统计配置文件来操作数据流图以进行优化。例如，在动态水平融合中，我们用生成的多形状融合核函数替换数据流图中的多分支算子，并更改 Router 以在路由 Cell 时将张量填充到支持的形状。对于推测式路由，我们重新排序数据流图中的算子以跳过和回滚路由逻辑。对于推测式权重加载，我们收集分支的参数，并插入额外的算子用于在运行时加载和卸载它们。配置文件引导的模型放置是一个例外情况，因为模型权重的加载超出了 torch.fx 的范围。在推理之前，Brainstorm 会加载由统计配置文件推导出的放置计划所给定的相应权重。在运行时，Brainstorm 的 Router 会根据放置计划翻译 router_fn 给出的路由决策，将 Cell 路由到适当的设备。

选择动态优化。对于一个给定的动态模型，我们使用基于规则的策略来选择动态优化。当单个分支无法饱和GPU核心时，使用动态水平融合。对于多GPU推理，使用配置文件引导的模型放置。当路由器阻塞GPU核函数提交时，启用推测式路由和权重预加载。当GPU内存有限并使用分页时，使用推测式权重预加载。表2列出了应用于每个模型的动态优化。例如，LiveSR是一个轻量级的超分辨率模型，单个分支可能无法饱和GPU。因此我们对其应用动态水平融合。此外，基于MoE的模型通常是需要多GPU部署的大型语言模型，因此我们应用了放置优化。TaskMoE的输入足够大，可以实现高GPU利用率，因此不需要水平融合。
Table 2: 基准测试规格。（Fusion：动态水平融合；Place：配置文件引导的放置；Route：推测式路由；Load：推测式权重预加载。）

A4 实验环境

  • 硬件配置:
    • 单GPU平台: AMD-EPYC-7V13 CPU，1x NVIDIA A100 (80GB) GPU。
    • 多GPU平台: Intel Xeon E5-2690 v4 CPU，8x NVIDIA V100 (32GB) GPU。
  • 软件配置:
    • 代码实现: 基于 PyTorch 实现，约13000行代码。
    • 依赖库: CUDA 11.3, cuDNN 8.6 (A100) / 8.2 (V100)。
  • 模型与数据集:
    • SwitchTransformer (Switch): NLP模型，在 MNLI 数据集【索引40，GLUE: A multi-task benchmark and analysis platform for natural language understanding，2018，arXiv】上进行评估。
    • TaskMoE: NLP模型，使用合成数据集评估。
    • SwinV2-MoE: 视觉模型，在 ImageNet22k 数据集【索引42，Imagenet: A large-scale hierarchical image database，2009，CVPR】上评估。
    • LiveSR: 视觉模型，在 Iowa-DOT 数据集【索引43，The 5th ai city challenge，2021，CVPR Workshops】上评估。
    • DynamicRouting (DRouting): 视觉模型，在 Cityscapes 数据集【索引44，The cityscapes dataset for semantic urban scene understanding，2016，CVPR】上评估。
    • MSDNet: 视觉模型，在 Imagenet 数据集【索引42，Imagenet: A large-scale hierarchical image database，2009，CVPR】上评估。
  • 基准（Baselines）:
    • PyTorch: 作为所有实验的通用基准，所有被评估模型的官方实现均基于 PyTorch。
    • Tutel: 专门为 MoE 模型设计的优化库，仅在 MoE 模型上与 Brainstorm 进行比较。
    • 所有实验中，Brainstorm 和基准都使用了相同的静态优化（如垂直核函数融合），以公平比较动态优化的增益。

A4 实验结果

Brainstorm (BRT) 在六个代表性的动态神经网络上进行了评估，与 PyTorch 原生优化和模型特定优化（如 Tutel）进行了比较。总体而言，Brainstorm 实现了最高 11.7倍 的加速（平均 3.29倍）或减少了 42% 的 GPU 内存使用。

1. Brainstorm 抽象的有效性 (§7.2)
* 表达能力: 将现有模型移植到 Brainstorm 只需少量代码修改（6-24行），证明了其抽象的简洁性 (Table 3)。
* 追踪开销: 启用 Cell 级数据流追踪带来的平均开销低于 1.0%，几乎可以忽略不计 (Figure 11)。
* Cell 路由效率: Brainstorm 的稀疏通信原语（Sparse All-to-All）比 PyTorch 的密集 all-to-all 通信快 1.88倍 至 2.78倍，因为它避免了不平衡路由中的额外填充开销 (Figure 12)。

2. 微基准测试 (§7.3)
* 动态水平融合: 与 PyTorch 顺序执行相比，动态水平融合最高可带来 41.8倍 的加速。其中，基于统计数据调优的核函数带来了 13.1倍 的加速，而分支的并发执行进一步带来了 3.18倍 的加速 (Figure 13)。
* 配置文件引导的放置: 通过优化模型放置减少跨 GPU 通信，通信延迟获得了 2.45倍至 6.65倍 的加速 (Figure 14)。
* 推测式优化: 在预测命中时，推测式路由能完全隐藏路由延迟，推测式权重预加载能隐藏权重加载延迟。预测失败时，性能与默认执行相当，开销可忽略。在微基准测试中，推测式权重预加载将内存需求减少了8倍 (Figure 15)。

3. 端到端模型执行 (§7.4)
* SwitchTransformer: Brainstorm (BRT) 比 PyTorch (Torch) 和 Tutel 分别快 3.63倍 和 3.33倍。通过为不同负载定制核函数，Brainstorm 避免了 Tutel 因大量填充而导致的计算浪费，并在256个专家时避免了 Tutel 的内存溢出问题 (Figure 16)。
* LiveSR: 动态水平融合 (BRT+HF) 相比基线 (BRT) 实现了高达 8.62倍 的加速。其中，仅使用统计调优核函数的垂直融合 (BRT+VF) 带来了 3.5倍 加速，水平融合进一步带来了 1.79倍 至 2.48倍 的增益 (Figure 17)。
* TaskMoE: Brainstorm 的稀疏通信 (BRT) 比 PyTorch 带来了 1.17倍 的吞吐量提升。在此基础上，配置文件引导的放置 (BRT+P) 进一步带来 1.34倍 的吞吐量提升，通过减少 42%~87% 的跨 GPU 通信实现 (Figure 18)。
* SwinV2-MoE: Brainstorm 的高效 Router 分别比 DeepSpeed-MoE 和 Tutel 带来了高达 5.04倍 和 1.52倍 的吞吐量提升。配置文件引导的放置在该模型上端到端增益不大，但在单层评估中显示出最高 1.26倍 的潜力，表明对更大型 MoE 模型有价值 (Figure 19, 20)。
* MSDNet: 推测式路由 (BRT+SP) 带来了高达 8.44倍 的加速。在此基础上，动态水平融合 (BRT+HF) 进一步带来高达 1.57倍 的增益，总加速比达到 11.7倍 (Figure 21)。
* DynamicRouting: 推测式路由 (BRT+SP) 实现了 90%~95% 的预测准确率，带来高达 1.7倍 的加速 (Figure 22)。推测式权重预加载将推理延迟降低了高达 1.97倍，同时将 GPU 内存消耗减少了 43.5% (Figure 23)。

A7 补充细节

8. 讨论

处理分布漂移。动态优化策略的分析是离线进行的，依赖于具有统计代表性的分析数据。如果实际运行中的数据分布与分析数据发生漂移，可能会导致 Brainstorm 的优化效果下降甚至产生负增益。如图24所示，这种影响取决于模型和漂移的程度。图24评估了分布漂移对动态水平融合的影响。Brainstorm 基于收集的配置文件只为4个和27个图像块的情况调整了 Conv2D 核函数。因此，当一个分支接收到多于4个图像块时，需要填充到27个图像块并使用非最优的27图像块核函数运行。初始分派为每个分支4个图像块，这样就不需要填充。为了模拟增加的分布漂移，我们将一些分支的负载增加到8个图像块，这种情况在配置文件中出现频率较低，因此未被 Brainstorm 调整。我们将分布漂移率定义为接收到的图像块数量与已调整形状（本实验中为4和27）不同的分支的比例。在图24中，我们发现随着漂移率的增加，Brainstorm 的动态水平融合 BRT+HF 的加速比从 4.65倍 降至 2.11倍（与仅应用垂直融合相比）。这是由于在接收8个图像块的分支上进行填充所造成的计算浪费。优化策略需要持续监控 Brainstorm 收集的配置文件，并在分布发生漂移时触发重新优化。重新优化需要时间（通常是几分钟），例如搜索新的放置方案和调整新的 GPU 核函数。因此，在冷启动或重新优化期间，模型执行不使用动态优化。目前，Brainstorm 专注于实现动态神经优化的机制。我们希望能够启发更多先进的解决方案来应对分布漂移的鲁棒性问题。

图 24: 当分支负载分布与用于调整 GPU 核函数的统计数据发生漂移时，Brainstorm 动态水平融合的加速比。

更多动态优化机会。Brainstorm 也可以应用于训练。在微调基于 MoE 的大型语言模型时，专家激活的统计数据可以像推理时一样被利用，例如，重新安排跨 GPU 的专家放置以减少通信量。此外，神经架构搜索中的许多算法也设计了动态架构（例如，DARTS【索引49，Single path oneshot neural architecture search with uniform sampling，2020，ECCV】、SPOS【索引50，DARTS: differentiable architecture search，2019，ICLR】），其激活只有在运行时才知道。它们训练的后期阶段可能会显示出更稳定的分支激活，这可能被 Brainstorm 利用。要支持训练，还需要解决一些工程问题。首先，当前实现中缺少训练中自动微分所需的反向传播。其次，一些算子可能会使 Brainstorm 的动态优化追踪失效。例如，BatchNorm 在训练中执行跨 Cell 计算，与推理时的 Cell 级计算不同，这需要手动指定。Brainstorm 也可以应用于动态稀疏性，即对不同输入使用不同的值/块级稀疏模式（例如，Longformer）。为了优化它们的执行，Brainstorm 需要在细粒度上收集模式统计。然后我们可以为不同的稀疏模式编译多个专门的 GPU 核函数（例如，使用 SparTA【索引51，SparTA: Deep-Learning model sparsity via Tensor-with-Sparsity-Attribute，2022，OSDI】），并在运行时激活最高效的一个。

A5 结论

本文通过收集和分析运行时配置文件，将模型执行专门化以适应动态性分布，从而为动态神经网络识别了一个新的动态优化空间。我们提出了 Brainstorm，这是第一个优化动态神经网络执行的深度学习框架。Brainstorm 的核心是 Cell 和 Router，它们让模型开发者在动态性的粒度上表达动态神经网络，从而可以追踪动态优化所需的信息。模型开发者可以专注于设计动态模型架构，而将优化工作留给 Brainstorm 框架。在 Brainstorm 中，我们提出了四种利用不同粒度运行时配置文件的动态优化。我们的评估表明，Brainstorm 可以将流行的动态神经网络加速高达 11.7倍（平均 3.29倍），或将 GPU 内存使用量减少 42%。

A6 附录

A. 工件附录

摘要。Brainstorm 通过 Cell 和 Router 抽象统一了动态神经网络的编程，为动态神经网络开辟了新的动态优化空间。本工件在单 GPU 和多 GPU 环境下复现了评估中的主要结果。请遵循 http://README.md 中的说明来复现主要结果。

范围。本工件将验证以下声明：
1. Brainstorm 抽象的有效性：通过复现图12的实验，可以验证 Brainstorm 抽象的有效性。
2. 微基准测试：通过复现图13-15的实验，可以用微基准测试验证所提出的动态优化。
3. 端到端模型执行：通过复现图16-23的实验，可以验证§7中声称的 Brainstorm 的端到端延迟。

内容。在本工件中，我们将复现图12-23。每个图都有一个 shell 脚本来自动复现和可视化评估结果。此外，我们还提供了一个托管在 Github Container Registry 上的预构建 Docker 镜像。用户可以使用此镜像快速启动一个容器，该容器已预配置好实验环境。

托管。本工件托管在 https://github.com/RaphaelHao/brainstorm/tree/osdi2023ae。要获取代码，请克 隆 Brainstorm 仓库并切换到 osdi2023ae 分支。

要求。
1. 硬件要求：图 13、15-17 和 21-23 需要一台配备 NVIDIA A100 (80GB) GPU 的服务器；图 12、14 和 18-20 需要一台配备八块 NVIDIA V100 GPU 的服务器。
2. 软件要求：请使用 docker 构建 docker/Dockerfile.update 来设置单 GPU 和多 GPU 实验的环境。还提供了一个一键式脚本 python scripts/docker_gh_build.py --type latest 来构建镜像。
3. CUDA 驱动：版本需大于 11.3。