D1: Activation Policies for Transformer Training

A workstation-scale measurement study on 2× RTX A6000.

D1 asks one fixed-scope question: if the pipeline partition and schedule are held fixed, do Transformer submodules sit in cost regimes different enough that a per-submodule save | recompute | offload policy can beat coarse-grained checkpointing?

The current answer is narrow:

• Selective recompute is the practical positive result. It shares the Pareto frontier with full-layer recompute at the measured OOM-pressure point.
• Naive saved-tensor CPU offload is the negative result. PyTorch hook-based offload survives memory pressure but is about 17× slower than recompute.
• The result fits Fold3D's design lesson. The offload failure comes from per-tensor request/launch overhead, the structural problem Fold3D [1] engineers around with batched offload and known scheduling windows.

D1 does not propose a new distributed schedule. It isolates activation lifecycle policy inside a fixed pipeline partition/schedule.

See docs/RUNBOOK.md for the test commands, CPU smoke checks, and the ordered GPU artifact refresh queue.

planned_selective_recompute and full_layer_recompute both lie on the Pareto frontier — neither dominates. The trade-off is 9.8% more memory for 2.1% less step time. Uniform CPU offload survives the memory pressure but is about 17× slower than recompute, and save_all OOMs at this point.

A four-stage harness on PyTorch: profile → plan → run → compare.

• Profile records per-submodule activation sizes, H2D / D2D transfer latency across tensor sizes, and a 2-stage sequential pipeline bubble (I measured the forward bubble at 13.26 ms for pp=2).
• Plan runs a greedy planner over save | recompute | offload decisions with a memory budget, at profile-point granularity (QKV / attention-core / attention-projection / FFN units).
• Run executes the planner's selective-recompute decisions through PyTorch checkpointing, alongside matched baselines: full-layer recompute, attention-only, attention+FFN, save-all, uniform CPU offload through saved_tensors_hooks, and a dedicated-stream offload variant.
• Compare produces three-repeat measurements with stderr and scaling sweeps along sequence, layer, and hidden dimensions.

To reproduce or refresh the headline artifacts, use the ordered GPU command queue in docs/RUNBOOK.md.

I read both papers before writing any code, and tried to let what I read shape the questions I asked, not the answers I would accept.

a) I started from the wrong end on purpose. Fold3D [1] section 5.2 makes it clear that per-tensor offload kernel launches dominate at scale, and that AIAO scheduling only hides DP.sync inside computation because the computation window is known in advance. Knowing this, I deliberately started from the opposite end — naive per-tensor saved_tensors_hooks, no batching, no scheduling — to see empirically where it would break.

It broke exactly where the paper said it should. I measured 52.58 GB/step of saved-tensor traffic moving at ~1.5 GB/s effective host throughput, an order of magnitude below PCIe limits. Kernel-launch and request overhead were dominating, not bandwidth. Adding a dedicated CUDA stream gave only 4.8% improvement, which made me confident the conclusion was robust: naive overlap is not enough without a structured lifecycle.

b) I held scheduling fixed on purpose. From NASPipe [2] I took the lesson that scheduling and execution are separable concerns. I held partition and pipeline schedule fixed and varied only activation policy, so that anything I found would be orthogonal to scheduling work and could later compose with CSP-style designs.

c) I scaled up until the question had stakes. At seq=512, every recompute and offload policy looked worse than save_all — the regime simply wasn't memory-pressured enough for the trade-off to matter. So I ran sweeps along sequence / layer / hidden until I found the OOM-pressure point where save_all OOMs but recompute and offload policies still run. That is the regime the headline result lives in, and it is where the question "does per-submodule policy beat coarse policy" actually has consequences.

This is a single-workstation harness, not a replacement for Fold3D-scale scheduling results. The strongest supported claim is about activation lifecycle policy under a fixed partition/schedule at the measured OOM-pressure point.

• A verified nsys overlap proof for offload copies. Current traces hit importer errors; the empirical arithmetic above is the fallback.
• A planner-selected per-submodule offload runtime. Currently only uniform offload is wired end-to-end; the planner can choose per-unit but the exec path does not honor it.
• Integration of cuda/batched_activation_copy.cu (a standalone batched pack/unpack primitive) into the autograd offload path. The Fold3D-style batching it would enable is exactly what should close the kernel-launch overhead gap I measured.
• Statistical strength of the selective-recompute advantage. The 2.1% effect over ±0.2% stderr across three repeats is significant but small; more repeats would harden the claim.

The natural extension is whether a segment-level lifecycle, designed along Fold3D's lines, would close the bubble/transfer gap at workstation scale, or whether the gap is structural to single-host hardware. That question needs multi-host measurements I do not have access to here.

[1] F. Li, S. Zhao, Y. Qing, X. Chen, X. Guan, S. Wang, G. Zhang, and H. Cui. "Fold3D: Rethinking and Parallelizing Computational and Communicational Tasks in the Training of Large DNN Models." IEEE Transactions on Parallel and Distributed Systems, 34(5):1432-1449, 2023. DOI: 10.1109/TPDS.2023.3247883.

[2] S. Zhao, F. Li, X. Chen, T. Shen, L. Chen, S. Wang, N. Zhang, C. Li, and H. Cui. "NASPipe: High Performance and Reproducible Pipeline Parallel Supernet Training via Causal Synchronous Parallelism." In Proceedings of the 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS 2022), pp. 374-387, 2022. DOI: 10.1145/3503222.3507735.