Proposal: overlap NVMe read and write with computing. about patrickstar HOT 2 CLOSED

嗯嗯，这样是可以减少一些移动开销。下面对 OS 划分策略稍微进行了一下计算：

这里有一个地方需要区分，就是 OS 实际上是分为 3 部分的：

• OS_cpu：一直都放在 cpu 上的 OS。满足：

  • M(OS_cpu) + M(P) + max(M(A, t)) <= M(GPU) + M(CPU)

  即：

  • M(OS_cpu) <= M(GPU) + M(CPU) - M(P) - max(M(A, t))

  为了尽可能利用 CPU，同时减少 CPU 和 NVMe 之间的移动，我们会让 OS_cpu 尽可能大，我们不妨让上式取等号：

  • M(OS_cpu) = M(GPU) + M(CPU) - M(P) - max(M(A, t)) [1]

• OS_move：可以移动的 OS；这部分 OS 会在开始 adam 计算的时候放在 CPU 上，而在正反向计算过程中，需要逐渐 offload 到 CPU 上去。满足：

  • M(OS_cpu) + M(OS_move) + M(P) + max(M(A, t)) > M(GPU) + M(CPU)
  • M(OS_cpu) + M(OS_move) + M(P) + M(A, t_adam) <= M(GPU) + M(CPU) [2]

  将 [1] 带入 [2] 中，并且考虑到 M(A, t_adam) = 0，有：

  • M(OS_move) <= max(M(A, t))

  您这里提到的优化方法，相当于是把这部分的 load/offload 放到 FWD 和 BWD 的计算过程中进行。对于这个优化，我们不妨让
  我们不妨让上式取等，也就是 M(OS_move) 相当于是 activation 的最大值 max(M(A, t))，即：

  • M(OS_move) = max(M(A, t)) [3]

• OS_nvme：一直都放在 nvme 上的 OS；这部分 OS 只能临时放在 CPU 上，访问方法是取固定几个 CPU chunk 作为 buffer，以循环的方式去访问他们。

  注意到：

  M(OS) = M(OS_cpu) + M(OS_move) + M(OS_nvme) = 4 * M(P)

  把 [1], [3] 带入 [4] 有：

  • M(OS_nvme) = 5 * M(P) - M(GPU) - M(CPU)

所以理论上来说，对于一个 100B 的模型来说，M(P) = 200G，在 1T 的内存的场景下，应该不需要有 M(OS_nvme)。然后可以在 warmup 的时候统计 maxM(A, t)，从而规定好 OS_cpu 和 OS_move。

from patrickstar.

暂时想到了这样的实现方案：

在完成 warmup 后，根据 activation 的最大值来计算出 OS_cpu 与 OS_move 对应的 chunk 数，并标记在对应的 optimizer state chunk 上。adam 开始时 OS_cpu 和 OS_move 都放在内存上，OS_nvme 通过共用 buffer 的形式循环访问。在 Adam 结束的时候，异步进行 OS_move 向 nvme 的写入，直到 forward 结束后，再逐渐移动回内存。

这个实现主要有 2 个问题：

1. 这意味着 nvme offload 这个操作目前是绑定在 optimizer 这个事儿上了。不过我认为这相较于 deepspeed 的把参数和 activation 放到 nvme 上实际是一种优化；
2. 这个异步写入可能需要用 C++ 实现，或者起一个新进程...

from patrickstar.