Enable activation offloading (#8255)
# Activation offloading for DeepCompile: plan against a floor the run
has actually reached
## What this does
Enables the activation-offload pass to DeepCompile init_z3. The forward
pass saves tensors the backward needs;
they sit in device memory for the whole step. This pass copies chosen
ones to pinned host memory
during the forward and brings each back before the backward reads it.
Enable with `compile.offload_activation: true`. It replaces the prefetch
and selective-gather passes
and is mutually exclusive with `offload_parameters` and
`offload_opt_states` — each plans against the
whole budget on its own.
## Results
Where a job is too tight to run at all, offloading everything rescues it
but costs 40%+ in step
time. Planning which activations to move recovers most of that.
Qwen3-14B, mb4, 8xH200, ZeRO-3, inductor, bf16 optimizer states,
`expandable_segments`. All arms use
DeepCompile; the question is what the offload pass adds.
| seq | baseline (no offload) | forced-all (blind) | ours (planned) |
ours moved | vs forced-all |
| --- | --- | --- | --- | --- | --- |
| 2048 | 1.76s | 2.50s | **1.85s** | 0GB — declines | run fits; pass
correctly does nothing |
| 3072 | 5.46s | 7.81s | **5.41s** | 0GB — declines | run fits; pass
correctly does nothing |
| 3328 | **OOM** | 8.67s | **7.38s** | 11.7GB of 22.0 | **15% faster** |
| 3584 | **DIED** (watchdog) | 10.80s | **9.76s** | 20.5GB of 23.7 |
**9.6% faster** |
| 4096 | **OOM** | 11.17s | 11.65s | 26.8GB of 27.0 | declines to plan;
safe |
Three properties, each measured:
- **It resolves the out-of-memory case.** At seq3328/3584/4096 the job
does not run without it.
- **It beats moving everything, where there is room to plan.** 15% at
seq3328 and 9.6% at seq3584 --
the two lengths that need offloading at all -- moving roughly half the
bytes. Copy cost is linear
at 0.108s per GB (R²=0.96), which is where the time comes from.
- **It does nothing when nothing is needed** — at seq2048/3072. It keeps
every activation resident
and costs nothing measurable, while blind offloading costs +42%/+43%.
At seq4096 there is only ~2.3GiB of real headroom, so it correctly
declines and matches forced-all
rather than dying.
The place this pass would pay is where recompute is expensive and PCIe
is idle -- long sequences
where attention recompute scales O(s^2) while a copy stays O(s), or
composed with recompute rather
than against it. Neither is measured here.
---------
Signed-off-by: pengdurice <pengduhit@gmail.com>
