AEGIS: Scaling Long-Sequence Homomorphic Encrypted Transformer Inference via Hybrid Parallelism on Multi-GPU Systems

1. 研究背景与问题 (Problem)

背景：
全同态加密（FHE）允许在加密数据上直接进行计算，为机器学习即服务（MLaaS）提供了隐私保护。然而，将 FHE 应用于 Transformer 模型（特别是长序列输入）时，面临着巨大的计算和内存开销。

核心挑战：
随着序列长度的增加，加密后的权重和激活值（Activations）会迅速膨胀，超出单张 GPU 的显存容量，迫使系统必须采用多 GPU 并行执行。然而，现有的并行策略在加密场景下失效，主要存在以下三个问题：

1. 明文并行策略失效： 传统的张量并行（Tensor Parallelism, TP）和数据并行（Data Parallelism, DP）假设张量可以任意切分且通信高效。但在 FHE 中，由于密文打包（Packing）机制，Token 被编码在固定的槽位（Slots）中，无法任意切分。直接应用明文策略会导致巨大的显存冗余（复制整个密文）或无效计算。
2. RNS 并行（Limb Parallelism）的通信瓶颈： 现有的同态加密加速器常利用剩余数系统（RNS）将模数链（Modulus Chain）切分到不同设备（Limb Parallelism）。然而，Transformer 中的矩阵乘法和非线性操作需要频繁的密钥切换（Key Switching）和旋转（Rotation），这导致设备间需要频繁同步 RNS 模数分量。实验表明，这种通信开销远超计算开销，甚至导致多 GPU 比单 GPU 更慢（例如在 2 张 GPU 上性能下降 2.25 倍）。
3. 通信放大效应： 加密通信的体积是明文的数百倍。例如，传输 1024 个 Token 在明文下仅需约 2MB，而在加密下（CKKS 方案）可能膨胀至 1.3GB。现有的集合通信原语（如 AllReduce, AllGather）在加密场景下效率极低。

结论： 现有的解决方案要么忽略了应用层的结构（仅关注加密底层），要么忽略了加密的约束（直接套用明文并行），导致无法在长序列场景下实现可扩展的推理。

---

2. 方法论：AEGIS 框架 (Methodology)

作者提出了 AEGIS（Application-Encryption Guided Inference System），这是一个编译器 - 运行时协同设计的框架，旨在通过应用层与加密层感知的混合并行来解决上述问题。

核心设计思想

AEGIS 不再将应用数据流和加密底层操作割裂处理，而是联合分析 Transformer 的数据依赖和 CKKS 的 RNS 耦合，进行统一的调度。

关键技术组件

1. 依赖感知的负载划分策略 (Dependency-Aware Workload Partitioning)
AEGIS 提出了两种协同的放置策略，以最小化通信：

• 模数相干放置 (Modulus-coherent placement)： 针对加密底层。分析密钥切换和旋转链中的依赖关系，将必须共同参与同一组操作的模数分量（Polynomials）强制放置在同一个 GPU 上。这消除了跨设备的密钥切换同步需求。
• Token 相干放置 (Token-coherent placement)： 针对应用层。分析 Token 级别的依赖（如 FFN 层的输出组、Attention 的聚合）。将属于同一个 Token 的不同密文切片（Slices）尽可能放置在同一个设备上，避免在设备间传输部分 Token 结果，从而减少 AllReduce 等集合通信的频率。
• 优先级： 模数依赖优先，Token 相干次之。

2. 前馈网络 (FFN) 的并行划分

• 在 FFN 层，AEGIS 确保属于同一输出组的权重切片位于同一设备。
• 每个设备持有部分加密激活值，并行进行本地多项式乘法，然后通过一次 AllGather 获取缺失的激活值，最后进行本地非线性计算（如 GELU）。这避免了明文 TP 中常见的冗余复制。

3. 自注意力机制 (Self-Attention) 的并行划分

• 分阶段执行： 将注意力计算分为部分计算、分布式 Softmax 和输出聚合。
• 优化通信： 在计算 $Q \times K^T$ 后，使用 ReduceScatter 将中间结果重新分布，使得每个设备可以独立进行 Softmax 的多项式近似计算，无需跨设备同步。最后通过 AllGather 获取权重进行 $A \times V$ 计算。

4. 通信隐藏与算子重排序 (Communication Hiding via Operator Reordering)

• 问题： 传统的执行顺序要求等待所有密文分量就绪后才能发起通信，导致 GPU 计算流阻塞。
• 解决方案： AEGIS 在编译期对多项式算子进行重排序。例如，在计算旋转（Rotation）时，交错不同偏移量的计算顺序，使得部分聚合结果更早可用。
• 效果： 将集合通信（如 ReduceScatter, AllGather）与剩余的计算任务重叠（Overlap），从而在不增加额外传输的情况下隐藏通信延迟。

5. 编译流程
AEGIS 包含三个阶段的编译过程：

1. 应用 IR 生成： 分析 Token 和序列级依赖。
2. 多项式 IR 生成： 分析 RNS 级依赖，插入加密集合通信原语，并根据硬件成本模型进行算子重排序。
3. 执行计划生成： 生成确定性的多 GPU 执行计划，运行时无需动态同步。

---

3. 主要贡献 (Key Contributions)

1. AEGIS 框架： 首个针对多 GPU 长序列加密 Transformer 推理的编译器 - 运行时协同框架，联合优化了应用数据流和底层多项式执行。
2. 依赖感知划分策略： 提出了“模数相干”和"Token 相干”的放置规则，解决了应用语义与加密约束之间的不匹配问题，显著减少了不必要的跨设备通信。
3. 确定性通信隐藏机制： 开发了多项式算子重排序机制，无需运行时启发式算法即可将通信与计算重叠，有效掩盖了设备间延迟。
4. 性能突破： 在保持高缩放效率的同时，大幅降低了显存占用和通信开销。

---

4. 实验结果 (Results)

实验在 NVIDIA RTX A6000 (2 GPU) 和 NVIDIA A100 (4 GPU) 上进行，使用 BERT-Base 模型和 SST-2 任务，输入序列长度从 128 到 2048。

• 通信开销降低：

  • 在 2048 Token 输入下，相比最先进的设计（如 HEBooster, Cinnamon, Hydra），AEGIS 在 FFN 层减少了 57.9% 的 GPU 间通信，在自注意力层减少了 81.3% 的通信。
  • 在 2 GPU 上，FFN 通信从数十 GB 降至 788.65 MB，自注意力降至 2.46 GB。

• 扩展效率 (Scaling Efficiency)：

  • 在 4 GPU 上，AEGIS 实现了 96.62% 的端到端扩展效率（接近线性加速）。
  • 相比之下，Limb 并行（HEBooster）效率极低（甚至低于 100%），Tensor 并行（Hydra）受限于显存和通信，效率约为 75%。

• 端到端加速比：

  • 在 4 GPU 上，相比单 GPU 基线，AEGIS 实现了 3.86 倍 的端到端加速。
  • 相比直接应用明文张量并行（在加密下效率极低），AEGIS 快 1.86 倍。

• 显存效率：

  • 在 4 GPU 上，相比复制中间密文的基线（如 Hydra），AEGIS 实现了 69.1% 的单设备显存减少。
  • 支持在 48GB 显存的 GPU 上处理 2048 Token 的长序列，而基线方法（如 Hydra, Cinnamon）在 512 或 1024 Token 时即发生显存溢出（OOM）。

---

5. 意义与影响 (Significance)

1. 理论突破： 证明了单纯复制明文并行策略或单纯依赖底层 RNS 并行都无法解决长序列加密推理问题。必须建立应用层与加密层协同的并行模型。
2. 实用价值： AEGIS 使得在商用 GPU 集群上高效、低成本地运行长序列隐私保护 Transformer 推理成为可能，为 MLaaS 中的隐私计算提供了可行的技术路径。
3. 系统指导： 该工作表明，未来的同态加密系统应将打包（Packing）、放置（Placement）和通信调度视为一个跨层优化问题，而非孤立处理。

总结： AEGIS 通过精细化的依赖分析和算子重排序，成功解决了多 GPU 环境下长序列加密 Transformer 推理的“通信墙”和“显存墙”问题，实现了接近线性的扩展效率和显著的显存节省，是该领域的重要进展。