Add OutputDiscardCheckpoint

[AkshitaB]

Summary

Add OutputDiscardCheckpoint, an activation-recompute primitive for cases where the output of a checkpointed region dominates memory (rather than its intermediates).

Vanilla torch.utils.checkpoint discards intermediates inside the wrapped function but can't free the output -- downstream consumers and their saved-for-backward references hold it live. OutputDiscardCheckpoint extends that pattern: forward runs under no_grad, the output's storage can be freed after downstream forward consumes it, and a backward hook recomputes the forward and rebinds the freed storage in place (via a C++ share_storage extension, with a Python fallback). The tensor object survives so existing autograd saved-tensor references stay valid; only its underlying bytes are recycled.

Useful for fat-output ops where the output is wider than the input -- precision casts (bf16 -> fp32 doubling), FFN up-projections, attention outputs before SDPA fuses them.

Contents

• src/olmo_core/nn/output_discard_checkpoint.py -- the primitive (287 lines), with a code-block usage example and :param:/:returns: docstrings.
• src/olmo_core/nn/__init__.py -- exports OutputDiscardCheckpoint.
• src/test/nn/output_discard_checkpoint_test.py -- four tests: storage discard/restore, grad parity vs a non-checkpointed Sequential, Python-fallback path forced via monkeypatch, and a 3D-Linear regression test (see below).

Bug found and fixed during integration testing

Writing the FFN integration script surfaced a real bug in the Python fallback. Tensor.set_(new_storage, ...) swaps dst's StorageImpl for a new one. Any autograd-saved view of dst that still referenced the original StorageImpl -- e.g. the 2D-reshape view that MmBackward saves when Linear is called with a 3D input -- would see the original storage, which was resize_(0)'d and never refilled. Backward then hit "tensor has non-zero numel but data is not allocated".

The C++ path didn't have this bug: it mutates dst's existing StorageImpl in place via set_data_ptr, so saved views see the new data through the same StorageImpl. The fix makes the Python fallback equivalent: resize dst's existing storage and copy src's bytes into it. Costs an extra allocation + copy during recompute on machines without a C++ toolchain, but is correct for tensors with saved views (which is most real workloads).

The 3D-Linear regression test ensures this doesn't sneak back in.

Standalone scripts (not run in CI)

Two scripts under src/scripts/ -- intentionally outside src/test/ so pytest never picks them up. Useful for human-driven benchmarking and verification:

• src/scripts/benchmark_odc.py -- compares baseline / torch.utils.checkpoint / ODC (C++) / ODC (Python fallback) on a fat-output Linear -> activation -> Linear workload. Reports peak GPU memory and forward/backward wall time. Supports --scenarios grid for a shape sweep.
• src/scripts/odc_ffn_integration_check.py -- wraps the real OLMo FeedForward (SwiGLU) with ODC around the fat activation(w1(x)) * w3(x) intermediate that w2 saves for backward. Asserts output + gradient parity vs the baseline; exits non-zero on failure.

Test plan

• pytest -v src/test/nn/output_discard_checkpoint_test.py -- 4 tests pass (incl. forced Python fallback and 3D-Linear regression).
• make checks -- isort / black / ruff / mypy all clean.
• python src/scripts/odc_ffn_integration_check.py --device cpu -- PASS (output and grad parity at 0 diff in fp32).
• On GPU: python src/scripts/benchmark_odc.py -- visually verify ODC peak memory < baseline.
• On GPU: python src/scripts/odc_ffn_integration_check.py --dtype bf16 -- verify parity holds in bf16.

🤖 Generated with Claude Code

[AkshitaB]

Benchmarking

========== n_layers=8  (d_model=2048, d_ff=8192, batch=2, seq=2048, dtype=torch.float32) ==========

  [fp32_cast] x.float() under ODC; consumer is an fp32 Linear (router pattern)
  config                          peak (MB)   fwd (ms)   bwd (ms)   total (ms)   mem saved   time delta
  ---------------------------- ------------ ---------- ---------- ------------  ----------  -----------
  baseline                            544.3       5.76      10.88        16.65  (baseline)   (baseline)
  torch.utils.checkpoint              544.3       5.79      10.89        16.68       +0.0%        +0.2%
  ODC (C++ if available)              544.3       5.73      10.88        16.60       +0.0%        -0.3%
  ODC (python fallback)               544.3       5.76      10.88        16.64       +0.0%        -0.1%

  [up_proj] fat Linear up-projection, no activation inside the discarded region
  config                          peak (MB)   fwd (ms)   bwd (ms)   total (ms)   mem saved   time delta
  ---------------------------- ------------ ---------- ---------- ------------  ----------  -----------
  baseline                           2560.3      44.07      84.83       128.90  (baseline)   (baseline)
  torch.utils.checkpoint             2560.3      44.42      84.93       129.35       +0.0%        +0.3%
  ODC (C++ if available)             2400.3      44.12     106.80       150.92       +6.2%       +17.1%
  ODC (python fallback)              2464.3      44.06     107.34       151.40       +3.7%       +17.4%

  [silu_up] silu(up(x)) inside; activation adds a saved intermediate to recompute
  config                          peak (MB)   fwd (ms)   bwd (ms)   total (ms)   mem saved   time delta
  ---------------------------- ------------ ---------- ---------- ------------  ----------  -----------
  baseline                           3584.3      44.49      85.95       130.44  (baseline)   (baseline)
  torch.utils.checkpoint             2656.3      44.88     107.87       152.75      +25.9%       +17.1%
  ODC (C++ if available)             2560.3      44.48     108.53       153.02      +28.6%       +17.3%
  ODC (python fallback)              2592.3      44.56     109.31       153.86      +27.7%       +18.0%

  [swiglu] OLMo SwiGLU FFN; three fat intermediates saved during recompute
  config                          peak (MB)   fwd (ms)   bwd (ms)   total (ms)   mem saved   time delta
  ---------------------------- ------------ ---------- ---------- ------------  ----------  -----------
  baseline                           6240.3      67.26     130.38       197.63  (baseline)   (baseline)
  torch.utils.checkpoint             3776.3      67.60     174.82       242.42      +39.5%       +22.7%
  ODC (C++ if available)             3904.3      67.31     175.86       243.17      +37.4%       +23.0%
  ODC (python fallback)              3904.3      67.32     176.54       243.86      +37.4%       +23.4%

  [rms_norm] RMSNorm + Linear + residual; cheap recompute, small fat-output savings
  config                          peak (MB)   fwd (ms)   bwd (ms)   total (ms)   mem saved   time delta
  ---------------------------- ------------ ---------- ---------- ------------  ----------  -----------
  baseline                            800.4       6.18      11.78        17.96  (baseline)   (baseline)
  torch.utils.checkpoint              800.3       6.34      12.00        18.33       +0.0%        +2.1%
  ODC (C++ if available)              560.3       6.19      11.99        18.19      +30.0%        +1.3%
  ODC (python fallback)               560.3       6.20      12.18        18.38      +30.0%        +2.3%

  [fp32_softmax] softmax in fp32 (attention/routing pattern); 2x size if base dtype is bf16/fp16
  config                          peak (MB)   fwd (ms)   bwd (ms)   total (ms)   mem saved   time delta
  ---------------------------- ------------ ---------- ---------- ------------  ----------  -----------
  baseline                            576.3       5.86      11.45        17.30  (baseline)   (baseline)
  torch.utils.checkpoint              800.3       6.10      11.69        17.79      -38.9%        +2.8%
  ODC (C++ if available)              560.3       5.93      11.69        17.61       +2.8%        +1.8%
  ODC (python fallback)               560.3       5.92      11.87        17.79       +2.8%        +2.8%

[TianhuaTao]

Looks ok for now as an alpha feature.
There might be concerns regarding RNG state, autocast context, and cases where requires_grad == False, but for now we don't have them in our active paths, so it should be fine.