PaperlayerView exact revision on arXiv

Source-linked AI summary

SlackSwap: Bandwidth-Aware Activation Swapping and Recomputation for Memory-Efficient DNN Training

Hui Feng, Xun Zhou, Yuxin Wang

arXiv:2610.13003v1cs.DC

TL;DR

GPU compute has outgrown host-link bandwidth, limiting how many saved activations can be swapped during computation. SlackSwap allocates the link budget to costly-to-recompute tensors and recomputes the rest; in simulation, it retains near-ideal throughput at half activation memory while outperforming the best baseline.

  • Problem

    GPU compute outpaced host-link bandwidth, while existing swap policies ignore link-time budgets and block checkpointing leaves the link idle.

    Source: §1 · p. 1

  • Method

    From one profiled iteration, SlackSwap spends the forward-pass link budget on tensors most expensive to recompute, then recomputes the cheapest eligible tensors to cover memory excess.

    Source: §1 · p. 1, §3.3 · p. 3

  • Results

    At 50% of peak activation memory, SlackSwap retains 96.3%, 97.6%, and 99.6% of ideal throughput for BERT-Large, GPT-2 XL, and Llama-7B, respectively, and achieves 1.10–1.11× the best baseline.

    Source: §4.2 · p. 4

  • Takeaways & Limitations

    In simulation, SlackSwap reaches activation-memory budgets that block checkpointing cannot, and shifts toward swapping with coherent links of hundreds of GB/s or long sequences.

    Source: §6 · p. 6

    All results are simulated; the simulator omits allocator fragmentation, HBM contention from copies, and GPUs sharing PCIe switches or host-memory bandwidth.

    Source: §4.5 · p. 5

Abstract

from arXiv · show

Swapping saved activations to host memory once enabled GPUs to train models that exceeded their device-memory capacity when compute throughput and host-link bandwidth were roughly balanced. For example, on a 2015 Titan X, PCIe bandwidth was sufficient to offload all VGG-16 activations during the forward pass. This balance no longer holds. From the Titan X to the H100, peak training throughput increased by 141x, whereas host-link bandwidth improved by only 3.9x. On an A100, the host link can hide transfers for only 11-38% of the activations in the transformer models we study. Existing swapping approaches, such as vDNN, SuperNeurons, and SwapAdvisor, offload tensors according to memory requirements without explicitly respecting the available link-time budget, resulting in transfer-induced stalls; block checkpointing, in contrast, leaves the host link underutilized. We present SlackSwap, which treats host-link time as an explicit resource budget. Using a single profiled iteration, SlackSwap allocates the forward-pass link budget to tensors that are most expensive to recompute, resolves the remaining memory deficit by recomputing the least expensive tensors, and schedules FIFO offloads and earliest-deadline-first prefetches without per-layer synchronization barriers. In trace-driven simulations of BERT-Large, GPT-2 XL, and a Llama-7B-style model on an A100, SlackSwap retains 96-100% of unlimited-memory GPU throughput while using half the activation memory, achieving 1.10-1.11x the throughput of the best baseline. Its plans are within 0.5% of a genetic search over 6,600-14,000 candidate plans, while requiring less than 0.4 seconds to generate.

1 Introduction

The host link has not kept pace with GPU compute, so activation swapping can no longer assume that transfers fit within forward-pass computation. SlackSwap budgets link time explicitly, combining selective swapping with recomputation to manage activation memory.

1 Introduction

  • From the Titan X to the H100, peak training FLOP/s grew 141× while host-link bandwidth grew 3.9×, reducing the share of activations hideable during the forward pass to 9–27%.On the Titan X, PCIe bandwidth could move every VGG-16 feature map during the forward pass.Source: §1 · p. 1
  • At half activation memory, vDNN, SuperNeurons, and SwapAdvisor achieve 36%, 66%, and 43% of unlimited-memory GPT-2 XL throughput, while block checkpointing leaves the link idle.The swap policies' copies sit on the critical path; block checkpointing re-runs 13% of the step.Source: §1 · p. 1
  • On an A100, recomputation costs 1–11 ps per byte while an exposed PCIe 4.0 copy costs 40 ps per byte, favoring swaps for the most expensive-to-recompute bytes.The link budget should go to those bytes; bytes beyond it should be recomputed.Source: §1 · p. 1
  • SlackSwap uses one profiled iteration to allocate the forward-pass link budget and recompute remaining memory excess; at half activation memory it retains 96–100% of ideal throughput and reaches 1.10–1.11× the best baseline.Its two-pass planner and barrier-free runtime are evaluated against four baselines.Source: §1 · p. 1

2 Background and Motivation

Saved activations can be kept, swapped to host memory, or recomputed, but the link can hide only a limited share of them. The section motivates allocating that link capacity selectively rather than treating memory relief as the only constraint.

2.1 Two ways to free a saved activation

  • Saved activations are retained for backward use, swapped to host memory and prefetched before use, or dropped and rebuilt by rerunning their producing operator.For GPT-style layers using FlashAttention, the paper reports 34 bytes saved per token per hidden unit in BF16.Source: §2.1 · p. 2
  • Existing systems choose swap targets by layer type or reuse distance, while block checkpointing reruns whole blocks and does not use the link.vDNN, SuperNeurons, and SwapAdvisor use different offload and prefetch policies.Source: §2.1 · p. 2, §2.3 · p. 2

2.2 How much can the link hide?

  • Before the activation-memory peak at the end of the forward pass, D2H transfers can free at most W T_F bytes, so the hideable fraction is bounded by forward-pass link capacity.The passage defines W as link bandwidth and T_F as forward time.Source: §2.2 · p. 2
  • Tensor cores raised the machine ratio β to 5,900–8,800 on the A100, while halving activation precision only doubled workload ratio I.I is forward FLOPs per saved byte; β is sustained FLOP/s per byte/s of host link.Source: §2.2 · p. 2
  • GPT-2 XL has I=1,200, so an A100 can hide 15% of its 22 GB of activations, and D2H—not H2D—binds because the backward pass lasts about 2T_F.Even a perfect swap-only policy stalls below 85% of peak memory.Source: §2.2 · p. 2

2.3 What should happen to the rest?

  • On an A100, rebuilding each GPT-2 XL tensor costs 1.9–87× less than an exposed copy, so swapping is valuable within the link budget and should favor costly or impossible rebuilds.Block checkpointing can spend 4–11× more recompute time on elementwise bytes than rebuilding those bytes alone.Source: §2.3 · p. 2

3 SlackSwap

SlackSwap plans activation handling from a profiled iteration under a memory budget. It first selects feasible swaps within the link budget, then recomputes eligible tensors to resolve remaining excess, with validation and runtime mechanisms to repair and execute the plan.

3.1 Problem

  • Given a profiled operator sequence, tensor sizes, production and use times, recompute costs, and memory budget, SlackSwap assigns each tensor to keep, swap, or recompute.The objective is to keep memory within budget while minimizing iteration time.Source: §3.1 · p. 2

3.2 Pass 1: spend the link budget

  • Pass 1 considers tensors in decreasing recompute cost per byte and admits swaps only when FIFO offloads and deadline-ordered prefetches meet their timing and memory constraints.Offloads must finish before the guarded end of the forward pass, and prefetches before each tensor's backward use.Source: §3.2 · pp. 2–3

3.3 Pass 2: recompute the cheapest bytes

  • Pass 2 recomputes the eligible resident tensor with the lowest cost per byte of memory excess, falling back to full-block recomputation when no candidate remains.This can go below the memory floor of checkpointing alone because swapped block inputs also free memory.Source: §3.3 · p. 3

3.4 Validation, repair and runtime

  • The plan is replayed in an event model and repaired when observed prefetch waits, lagging offloads, or memory-fit failures warrant changes.At runtime, two copy streams use production-order D2H and deadline-order H2D without per-layer barriers; operators wait only for a needed input or memory.Source: §3.4 · p. 3

4 Evaluation

The evaluation uses A100 trace-driven simulations to compare SlackSwap with swapping and checkpointing baselines across workloads, activation budgets, and hardware settings. SlackSwap retains near-ideal throughput at half the activation memory, while ablations and sensitivity tests show the importance of its link-budget and runtime choices.

4.1 Methodology

  • The A100 simulator models one GPU iteration for BERT-Large, GPT-2 XL, a Llama-2-7B-style model, and VGG-16, comparing four activation-management baselines.It models PCIe 4.0 at 25 GB/s per direction and assumes a fragmentation-free allocator; baselines are vDNN, SuperNeurons, SwapAdvisor, and block checkpointing.Source: §4.1 · pp. 3–4

4.3 Where the time goes

  • At 40% of peak memory, SlackSwap recomputes 5.3% of the GPT-2 XL step and 0.8% of the Llama-7B step, with stalls below 0.1%.For GPT-2 XL at 50%, SlackSwap takes 409 ms versus 399 ms ideal, while vDNN’s barriers stretch the step to 1,105 ms.Source: Table 2 · p. 4, §4.3 · p. 4

4.4 Ablation and optimality

  • Ablations lose 0.1–7.9 points when swap candidates are ranked by production time, 3.0–10.7 points without the link-budget pass, and 32–37 points with per-layer barriers.SlackSwap plans are within 0.45% of a seeded genetic search over 6,600–14,000 plans and take at most 376 ms to generate.Source: Table 3 · p. 4, §4.4 · p. 4

4.2 Throughput across memory budgets

  • At 50% of peak activation memory, SlackSwap retains 96.3%, 97.6% and 99.6% of ideal throughput for BERT-Large, GPT-2 XL and Llama-7B, achieving 1.10–1.11× the best baseline.At 30% of peak, it retains 89.7–97.0% of ideal throughput and achieves 1.07–1.16× the best baseline.Source: §4.2 · p. 4, Figure 4 · p. 5

4.5 Sensitivity

  • For Llama-7B, the hideable fraction rises from 36% at 2K tokens to 60% at 32K, while SlackSwap throughput approaches ideal from 96.9% to 99.8%.At 30% of peak, SlackSwap’s lead over the best baseline ranges from 1.07× to 1.16× across the V100, A100, and H100 results reported.Source: §4.5 · p. 5
  • All reported results are simulated for one GPU with resident model states; allocator fragmentation, copy-related HBM contention, and shared PCIe or host-memory bandwidth are not modeled.The paper identifies combining SlackSwap with ZeRO-Offload’s gradient traffic as future work.Source: §4.5 · p. 5
  • At 30% of peak, increasing link bandwidth from 12 to 200 GB/s shifts GPT-2 XL from 1.6 GB swapped and 14.6 GB recomputed to 16.7 GB swapped and no recomputation.SlackSwap stays 0–12 points above the best baseline; swap-centric policies need 100–450 GB/s per direction to beat checkpointing.Source: Figure 6 · p. 5, §4.5 · p. 5

5 Related Work

Prior work swaps activations, recomputes them, or combines both, but SlackSwap sizes swaps by link time rather than memory pressure.

5 Related Work

  • Earlier swapping systems move activations to host memory, while SlackSwap sizes swaps by link time rather than memory pressure.The cited systems include vDNN, TFLMS, SuperNeurons, SwapAdvisor, FlashNeuron, and Sentinel.Source: §5 · p. 5
  • Recomputation methods range from checkpointing every √n layers and integer-programming schedules to greedy runtime eviction and selected-operator recomputation.These approaches are attributed to Chen et al. [5], Checkmate [6], DTR [7], and Megatron-LM [12], respectively.Source: §5 · pp. 5–6
  • Hybrid methods choose or optimize swapping and recomputation using runtime access patterns, dynamic programming, or mixed-integer programming.SlackSwap instead is described as a closed-form, millisecond planner that assigns the explicit link budget to tensors most expensive to rebuild.Source: §5 · p. 6

6 Conclusion

SlackSwap allocates a limited host-link budget to expensive-to-recompute tensors and recomputes the rest. In simulation, it maintains near-ideal throughput at half activation memory, with the strategy shifting toward swapping when links are coherent or sequences are long.

6 Conclusion

  • At half activation memory, SlackSwap holds 96–100% of ideal throughput in simulation and reaches memory budgets that block checkpointing cannot.Its plan spends the link budget on tensors most expensive to recompute and recomputes the rest.Source: §6 · p. 6
  • SlackSwap allocates link time to tensors most expensive to recompute, recomputes the rest, and keeps copies off the critical path.Source: §6 · p. 6
  • With coherent links of hundreds of GB/s or long sequences, SlackSwap shifts from recomputation toward swapping.Source: §6 · p. 6
Loading 2610.13003v1…