论文

Replay Contracts 与选择性保留改进激活重算内存效率

State-Aware Activation Recomputation with Replay Contracts and Selective Preservation

精选理由

训练时省显存的激活重算常因状态不一致悄悄算错,这篇论文把状态变成显式契约,H100 上步时快了 18%,写训练框架的值得看。

arXiv 论文 2611.00590 提出用版本化 replay contract 将激活重算所需状态显式化,解决 mutable buffer、别名字节和随机数生成器状态在重放时改变导致梯度计算错误的问题。联合规划器按 segment 在保留激活、整 buffer 快照和选择性区域保留之间权衡内存与恢复成本。在冻结 sparse-write H100 负载、每项 50 个新进程的测量中,联合规划在 44 GiB 预算下比正确计价的 eager/whole 规划器减少 18.16% 的配对步时;在 32 GiB 下比固定区域方案快 4.98%。CIFAR-100 dense-queue 应用在 6 个训练种子下得到字节级一致的测试分数。

原文 · arXiv cs.AI

State-Aware Activation Recomputation with Replay Contracts and Selective Preservation

Activation recomputation saves accelerator memory by replaying forward segments during the backward pass, but replay is only correct if it sees the state the original forward pass saw. Mutable buffers, aliased bytes and random-generator state can change in between, so re-executing unchanged code can differentiate a different computation. We make this state explicit. A versioned replay contract specifies what replay must restore, and a joint planner chooses, per segment, between activation retention, whole-buffer snapshots and selective-region preservation, charging each option its measured memory and recovery cost. Independent derivative and byte-level checks validate the contract, and a Torch/JAX cohort retains all 5,900 execution records, including matches, rejections and divergences. On frozen sparse-write H100 workloads, each measured over fifty fresh processes, joint planning reduces paired step time by 18.16% against a correctly costed eager/whole planner at a 44 GiB budget (approximate 97.5% ratio interval $[.8172,.8195]$). At 32 GiB, where the eager/whole planner has no feasible static policy, joint planning is 4.98% faster than fixed regions at 29.94 rather than 26.20 GiB. A complete CIFAR-100 dense-queue application runs seven methods for 100 encoder and 25 probe epochs at six training seeds, with byte-identical test scores within each seed. There, joint planning uses 22.41% more allocated memory than whole snapshots, with a process-time geometric-mean ratio of 1.0080 (seed range $[.9845,1.0191]$). Together, the results map where selective preservation pays off: sparse writes with large activations benefit, while in the dense-queue workload whole snapshots use less allocated memory.