CacheGen: KV Cache Compression and Streaming for Fast Large Language Model Serving

研究背景

  1. 分布式机器学习训练中的数据传输：集体通信

  2. 集体通信原理

    • 输入

      拓扑结构（topology）：表示网络中节点（如服务器、GPU）和连接方式（如网络链路）的布局。

      需求（demand）：the amount of data each GPU wants to send to other GPUs in the topology

    • 输出

      路径（routes）：在网络中为每个通信任务选择合适的传输路径，确保数据能够高效地从一个节点传递到另一个节点。
      调度（schedule）：为各个任务分配时间和资源，使得任务可以在不同的时间点按计划执行，避免资源冲突。
      通过上面两个输出，来最大化带宽利用率/最小化任务完成时间，或者两者都有优化。

  3. 问题：现有集体通信方案的局限性

    在应对大规模集群时可能会为了扩展性而牺牲质量。

相关工作

  • 精确建模

    如MSCCL：精确地建模系统及其计算方面，努力获得接近最优的解决方案。
    但当规模扩大到更大时（例如 30 到 60 个机箱的分布式训练任务），这些方法的计算复杂度和处理难度显著增加，导致它们变得不可行或难以有效应用。

  • 启发式建模

    如TACCL：为了简化问题，牺牲了相当大的性能。
    将路由（数据包的传输路径）和调度（每个数据块在每条链路上的发送时机）分开处理；
    性能显著低于经过人工设计的解决方案，具体差距为两倍或更多。

本文工作

  • TE-CCL：一种基于流量工程的集体通信方法。具备更好的可扩展性和更佳的性能。
  • 实验：在 two-chassis NDv2 topology（一种具体的网络拓扑结构）上，TE-CCL相较于当前最先进的解决方案，能够提升 2倍以上 的性能。此外，TE-CCL 还可以在更大规模的拓扑结构上继续保持良好的扩展性。

TE-CCL

  1. 基于混合整数线性规划（MILP）的优化方法
  2. 在特定通信条件下，TE-CCL将整数变量去除，问题由此转变为线性规划
  3. 在 TE-CCL 这个整体方法或框架中，包含了能够处理多租户和异构拓扑的算法