Metadata-Version: 2.4
Name: ckptplan
Version: 0.1.0rc1
Summary: Profiles checkpointable PyTorch model blocks and selects a gradient-checkpointing plan for an activation-memory budget.
Author: Ryan Abraham
Maintainer: Ryan Abraham
License-Expression: MIT
Project-URL: Homepage, https://github.com/ryamabra/ckptplan
Project-URL: Repository, https://github.com/ryamabra/ckptplan
Project-URL: Issues, https://github.com/ryamabra/ckptplan/issues
Project-URL: Changelog, https://github.com/ryamabra/ckptplan/blob/main/CHANGELOG.md
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Operating System :: OS Independent
Requires-Python: <3.13,>=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch<2.14.0,>=2.5.0
Provides-Extra: dev
Requires-Dist: pytest>=7.4; extra == "dev"
Requires-Dist: build>=1.0; extra == "dev"
Requires-Dist: mypy>=1.10; extra == "dev"
Dynamic: license-file

<p align="center">
  <img src="assets/ckptplan-wordmark.svg" alt="ckptplan logo" width="420">
</p>

# ckptplan

Profiles checkpointable PyTorch model blocks and selects a gradient-checkpointing
plan for a user-specified activation-memory budget, minimizing estimated
recomputation overhead.

Gradient checkpointing trades compute for memory: a checkpointed block discards
its activations during the forward pass and recomputes them during the backward
pass. Checkpointing *everything* is the usual default and is rarely what you
want — it pays the maximum recompute cost for memory you may not need to save.
`ckptplan` measures each block, then picks the subset to checkpoint that meets
your memory budget at the lowest estimated recompute cost.

**Status: release candidate (v0.1.0rc1), tagged locally and not yet published.** The
full v0.1 API works, is covered by a 163-test CPU suite, and its four
checkpointing planners have passed genuine gradient-correctness checks on real
A10G hardware — see [Verified evidence](#verified-evidence). No memory-saving
or throughput percentage is claimed as a gate; see
[Measured, and not claimed](#measured-and-not-claimed).
[`MVP_SPEC.md`](./MVP_SPEC.md) (Revision 3.4) is the accepted design;
[`STATE.md`](./STATE.md) tracks progress and open blockers;
[`CHANGELOG.md`](./CHANGELOG.md) lists what this release candidate includes and
the metadata decisions still awaiting the maintainer's confirmation.

## Built with

- **Python** — public library, planning algorithms, profiling, and benchmarks
- **PyTorch** — tensor execution, autograd, and gradient checkpointing
- **CUDA** — GPU activation-memory measurement and A10G validation
- **Modal** — reproducible cloud GPU benchmark workflows
- **pytest** — CPU and correctness test suite
- **GitHub Actions** — Python/PyTorch compatibility CI

GitHub reports the repository as Python-only because CUDA is accessed through
PyTorch rather than through standalone `.cu` source files.

## Install

```bash
pip install -e ".[dev]"
```

Or, from a built wheel:

```bash
python -m build
pip install dist/ckptplan-*.whl
```

Python 3.10–3.12, PyTorch >=2.5.0,<2.14.0. Licensed under MIT (see
[`LICENSE`](./LICENSE)).

## The pipeline

Five calls, in order. A runnable version of exactly this is in
[`examples/end_to_end.py`](./examples/end_to_end.py).

```python
import torch
from ckptplan import (
    declare_blocks, profile_blocks, plan_checkpoints, apply_plan, run_benchmark,
)

device = "cuda"
model = torch.nn.ModuleList([
    torch.nn.TransformerEncoderLayer(
        d_model=64, nhead=4, dim_feedforward=256,
        dropout=0.0,          # a stochastic block cannot be checkpointed
        batch_first=True, device=device, dtype=torch.float32,
    )
    for _ in range(4)
])
example_inputs = (torch.randn(2, 32, 64, device=device),)

# 1. Declare the checkpointable blocks, in execution order.
blocks = declare_blocks(model, [(f"layer{i}", layer) for i, layer in enumerate(model)])

# 2. Measure each block: isolated activation bytes and recomputation time.
profiles = profile_blocks(blocks, example_inputs, device=device, dtype=torch.float32)

# 3. Choose which blocks to checkpoint, under an activation budget.
activation_total = sum(p.activation_bytes_estimate or 0 for p in profiles)
plan = plan_checkpoints(
    profiles, blocks,
    target_kind="activation_budget_bytes",
    target_value=activation_total // 2,     # keep at most half the activations
    planner="dynamic_programming",
)

# 4. Build a runnable module that checkpoints exactly the selected blocks.
container = apply_plan(blocks, plan, example_inputs, None)
output = container(*example_inputs)

# 5. Optionally, measure it — including a gradient correctness check against
#    an equivalent no-checkpoint plan.
result = run_benchmark(
    blocks, plan, example_inputs, None, lambda out: out.float().square().mean(),
    device=device, dtype=torch.float32, check_correctness=True,
)
print(result.peak_allocated_bytes, result.step_latency_ms_mean, result.correctness_passed)
```

### What each step does

| Call | Returns | Notes |
|---|---|---|
| `declare_blocks` | `tuple[CheckpointableBlock, ...]` | Blocks must be disjoint module subtrees with unique ids. Never calls `forward()`. |
| `profile_blocks` | `tuple[BlockProfile, ...]` | Measures activation bytes (CUDA only) and genuine full-recomputation timing. Restores all caller-owned module state, including on error. |
| `plan_checkpoints` | `CheckpointPlan` | Deterministic: identical inputs give a bit-identical plan. Planners: `dynamic_programming`, `greedy`, `uniform`, `checkpoint_all`, `no_checkpoint`. |
| `apply_plan` | `CheckpointedSequential` | Reuses the original module instances and preserves parameter identity, so existing optimizers keep working. |
| `run_benchmark` | `BenchmarkResult` | Latency, peak allocated/reserved memory, and an optional correctness check against a no-checkpoint reference. |

Use `validate_plan` to check a serialized plan against a model before applying
it; it re-derives every block's execution signature and verifies the model
fingerprint.

## CPU is timing-only

**Activation-based planning requires CUDA.** On CPU, PyTorch exposes no
allocator counters equivalent to `torch.cuda.max_memory_allocated`, so
`profile_blocks` cannot measure activation bytes. Rather than invent a number,
it reports:

```python
profile.timing_only               # True
profile.activation_bytes_estimate # None
profile.activation_bytes_method   # None
```

and `plan_checkpoints` refuses those profiles outright:

```
TimingOnlyProfileError: activation-based planning requires real activation-byte
profiles; CPU timing_only profiles are not valid planner inputs
```

This is a deliberate guard: the planner optimizes recompute cost subject to a
*memory* constraint, and it will not pretend to satisfy a budget it cannot
measure. On CPU you can still use `declare_blocks` and `profile_blocks` for
timing, and `apply_plan`/`run_benchmark` work with any plan you already have.

`examples/end_to_end.py` runs on either device: on CUDA it completes all five
steps; on CPU it stops at step 3 and prints the reason.

```
$ python examples/end_to_end.py
device: cpu  (torch 2.13.0)

1. declare_blocks -> 4 blocks: ['layer0', 'layer1', 'layer2', 'layer3']

2. profile_blocks:
     timing_only               = True
     activation_bytes_estimate = None
     activation_bytes_method   = None
     forward_time_ms_mean      = 0.3812
     eligible_for_checkpoint   = True

3. plan_checkpoints -> TimingOnlyProfileError (expected on CPU)
     activation-based planning requires real activation-byte profiles; CPU
     timing_only profiles are not valid planner inputs
```

## Blocks that cannot be checkpointed

`profile_blocks` marks a block ineligible rather than silently producing wrong
gradients. A block is excluded when it is stochastic (dropout and friends —
recomputation would not reproduce the forward pass), stateful in a way
recomputation would re-apply, or has no differentiable output. The reason is
recorded in `profile.exclusion_reason` and carried into the plan.

## Verified evidence

- **159 CPU tests pass** (163 including packaging/metadata checks added for
  this release), run via `.venv/bin/python -m pytest -q`, across Python
  3.10/3.12 and PyTorch 2.5.0/2.13.0 in CI.
- **All four checkpointing planners passed genuine A10G gradient-correctness
  checks.** `benchmarks/matrix_a10g_result.json` (24-layer, 1.2B-parameter
  transformer, seq_len 2048, batch 1, `rtol=1e-3, atol=1e-5`), re-run after
  fixing two defects in the original correctness harness (see `STATE.md`'s
  "Correctness Evidence — CORRECTED" section for the full defect history):

  | planner | correctness_passed | max_grad_diff |
  |---|---|---|
  | checkpoint_all | true | 7.105e-15 |
  | uniform | true | 6.217e-15 |
  | greedy | true | 5.329e-15 |
  | dynamic_programming | true | 5.329e-15 |

  Every value above is a real, non-null result from the normal completion
  path (`oom: false`) — not a null placeholder and not a value produced by the
  OOM-fallback path, which under the fixed code can only ever leave
  `correctness_passed` as `None`. A second, independent boundary run at
  seq_len 512 also passed with exact `max_grad_diff: 0.0`
  (`benchmarks/boundary_correctness_result.json`).
- `no_checkpoint`'s correctness fields (`correctness_passed`, `max_grad_diff`)
  are `null` **by design, not by gap**: it is the reference plan itself, so
  `run_benchmark` skips the correctness check for it rather than comparing it
  against itself.
- **Known, honest limitation — not glossed over:** at seq_len 4096 / batch 4,
  `no_checkpoint` itself OOMs on a single A10G
  (`benchmarks/oom_boundary_a10g.json`). `checkpoint_all` completes at that
  configuration, but **no correctness comparison can exist there**, because
  there is no baseline run to compare its gradients against — the reference
  itself cannot execute, isolated or otherwise. This is a hardware/harness
  ceiling at this model scale, not a defect in the checkpointing logic.

## Measured, and not claimed

Per MVP_SPEC.md §12.5, no release gate asserts a percentage of memory saved or a
bound on throughput overhead. Two things are worth stating plainly:

- **Reported, not gated:** memory reduction, step-time overhead, and the
  prediction gap. The profiler's additive per-block isolated activation estimate
  legitimately exceeds the measured whole-model peak reduction, because
  parameters, gradients, and allocator reuse dominate the end-to-end peak. That
  gap is reported, not corrected.
- **Gradient correctness is now verified for the two A10G configurations that
  can run at all** (seq_len 512 and seq_len 2048, see
  [Verified evidence](#verified-evidence) above) — the two defects that
  previously made every `correctness_passed` value in this repository
  meaningless (a shared-parameter self-comparison, and an indentation bug that
  routed the real comparison through the OOM handler) are both fixed, tested,
  and re-verified against real A10G runs. What remains unverified is
  correctness at configurations where no reference can execute at all — see
  the seq_len 4096 / batch 4 limitation above. See STATE.md's "Correctness
  Evidence" section for the full defect history.

`benchmarks/report.py` re-reports any saved benchmark JSON locally, for free:

```bash
python benchmarks/report.py benchmarks/matrix_a10g_result.json
```

## Development

```bash
pip install -e ".[dev]"
pytest -q
```

CI runs the suite on CPU across Python 3.10/3.12 and PyTorch 2.5/2.13
([`.github/workflows/ci.yml`](./.github/workflows/ci.yml)). The GPU benchmarks
under `benchmarks/` require Modal and an A10G and are run separately, never in
CI.

## Documents

- [`MVP_SPEC.md`](./MVP_SPEC.md) — accepted v0.1 design (Revision 3.4).
- [`ARCHITECTURE.md`](./ARCHITECTURE.md) — component structure.
- [`STATE.md`](./STATE.md) — progress, decisions, and what is proven vs pending.
