Selective activation checkpointing (SAC) decides which activations to save for the backward pass and which to recompute: each choice is a point in the plane of peak memory and recomputation time, and the goal is that plane's Pareto frontier. Current practice falls short of this goal. Hand-tuned policies fix a single point that does not adapt when the model, batch or kernels change; PyTorch's SAC reruns each block's forward pass in program order, recomputing values the backward pass never reads; and the memory budget of torch.compile treats recomputation as free or forbidden and selects saved tensors by bytes alone, an assumption that fails for the custom kernels that produce many of a large model's activations. We present TapeSAC, which automates SAC in eager PyTorch without editing the model's forward code. TapeSAC records each block's forward pass at the dispatcher and, during backward, replays only the producers of the values autograd consumes, so saving a tensor eliminates every producer needed only for it. A planner selects the saved tensors from measured recomputation costs under a single limit on step-peak memory. On Qwen3.8-27B and Qwen3.6-35B-A3B, replay alone reduces step time by 3.0% and 1.7% versus full checkpointing at no greater memory cost; memory-aware planning raises these gains to 6.3% and 13.9% under a 90%-of-device-memory limit. The evaluated hand-tuned policies fall short of this memory-recomputation frontier and, on Qwen3.6-35B-A3B, run out of memory at larger batches. Over 500 steps on real data, TapeSAC reduces training time by 8-11%; on Qwen3.5-397B, it is 2.7% faster than a model-specific hand-tuned policy while using 14% less memory.
Zejian Wang, Yifan Zhao, Kaihui Ren, Xiaoling Zhou, Yubin Wang, Chenhao Guan, Lin Zhang, Yufeng Li 等 14 位作者
Fully Sharded Data Parallelism (FSDP2) maintains parameters and gradients as per-parameter tensors, whereas its collective path operates on flat, rank-major buffers. Converting between these layouts copies all-gather outputs into parameter storage and packs gradients into full-group fp32 buffers. Layout conversion alone accounts for 6.4% of step time in Llama-3.1 70B training at 8K tokens per GPU. Communication kernels impose a second cost: a default NCCL reduce-scatter can slow a concurrent GEMM that would otherwise complete in a single wave by 69%. We present TIDE, a set of zero-copy collectives built on public NCCL 2.31+ APIs. TIDE preserves FSDP2's per-parameter storage and compute kernels while folding layout placement into collective addressing. TIDE-AG delivers shards directly to parameter storage using NVSwitch multicast and paced recursive-doubling RDMA, with no SMs used for transport or placement. TIDE-RS consumes per-parameter bf16 gradients directly and reduces them into fp32 shards (a mixed-precision mode that NCCL does not provide) using copy engines and short-lived thread blocks (CTAs) for node-local accumulation and NVSwitch for cross-node reduction. A two-phase pipeline separates gradient reads from network progress; Slot-Reuse optionally stages gradients in vacated all-gather slots. On 64 GPUs, TIDE reduces step time by up to 27.1% relative to Stock FSDP2 and 6.0% relative to FSDP2+SymmMem across Llama-3.1 70B, Wan2.2, and Qwen3-30B-A3B. It reduces reduce-scatter SM time by factors of 12.5-14.5 and cuts interference with concurrent compute by up to an order of magnitude. Its reduce-scatter buffer remains nearly constant from 16 to 64 GPUs at a fixed per-GPU output size, and its peak memory is up to 4.7 GiB below FSDP2+SymmMem's and, with Slot-Reuse on Wan2.2, below Stock FSDP2's.