Metadata-Version: 2.4
Name: zkdv
Version: 0.0.0
Requires-Dist: blake3>=1.0.0
Requires-Dist: numpy>=2.0
Requires-Dist: flax>=0.12.7 ; extra == 'jax'
Requires-Dist: jax>=0.10.2 ; extra == 'jax'
Requires-Dist: optax>=0.2.6 ; extra == 'jax'
Requires-Dist: torch>=2.13,<2.14 ; extra == 'torch'
Provides-Extra: jax
Provides-Extra: torch
Summary: Add your description here
Requires-Python: >=3.12
Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM

# ZKDV

ZKDV produces a verification transcript for updates made by JAX or Torch
training loops. Shared configuration lives at the package root; the optional
frontends are loaded lazily as `zkdv.jax` and `zkdv.torch`. Root records name
the selected prover `whitebox-jax` or `whitebox-torch`.

## Quick start

Install the optional JAX frontend:

```console
pip install 'zkdv[jax]'
```

```python
import jax
import jax.numpy as jnp
import optax
import zkdv


# Initialize ZKDV once, before compiling the training step.
config = zkdv.ZKDVConfig(
    sampling_policy=zkdv.SamplingPolicy(
        state_check_probability=0.01,
        update_check_probability=0.01,
    ),
)
proof = zkdv.jax.ZKDV("proof", config)
proof.start(rolling_window=2, probe_ratio=0.01)


# Create the state as usual. ZKDV adds only the proof owner and an optional
# exportable replacement for apply_fn.
def apply_fn(variables, batch):
    return variables["params"] * batch


state = zkdv.jax.contrib.TrainState.create(
    proof=proof,
    apply_fn=apply_fn,
    params=jnp.arange(8, dtype=jnp.float32),
    tx=optax.adam(1e-3),
)


# Replace @jax.jit with @proof.jit. The body remains ordinary Flax/JAX code,
# including arbitrary loss construction and any number of apply_fn calls.
@proof.jit(donate_argnums=(0,))
def train_step(state, batch):
    def loss(params):
        predictions = state.apply_fn({"params": params}, batch)
        return jnp.mean(predictions)

    return state.apply_gradients(grads=jax.grad(loss)(state.params))

for _ in range(10):
    state = train_step(state, jnp.ones((8,), dtype=jnp.float32))

# Close once, after the training loop.
proof.close()
```

The example independently samples state and update checks with probability
`0.01`; choose probabilities appropriate for the run. `proof/records.cborl` is
the resulting transcript. `ZKDV.close()` waits for all queued commitments and
sampled checks, flushes the tape, and raises a semantic `ZKDVPoisonedError` if
asynchronous verification failed.

A runnable version is available at
[`examples/jax/train_state.py`](examples/jax/train_state.py). The explicit
transaction API remains demonstrated by
[`examples/jax/basic_training.py`](examples/jax/basic_training.py).

## Torch 2.13 quick start

Install only the optional Torch frontend:

```console
pip install 'zkdv[torch]'
```

```python
import torch
from torch import nn
from torch.nn import functional

import zkdv

proof = zkdv.torch.ZKDV("proof")
proof.start(rolling_window=32, probe_ratio=0.01)
model = nn.Linear(10, 2).cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
model, optimizer = proof.prepare(model, optimizer)


@proof.attest
def attested_step(model, optimizer, batch):
    inputs, targets = batch
    loss = functional.cross_entropy(proof.forward(model, inputs), targets)
    proof.backward(loss)
    optimizer.step()
    optimizer.zero_grad()


data = [
    (
        torch.randn(8, 10, device="cuda"),
        torch.randint(2, (8,), device="cuda"),
    )
]
for inputs, targets in data:
    with proof.transaction(batch=(inputs, targets)):
        loss = functional.cross_entropy(model(inputs), targets)
        proof.backward(loss)
        optimizer.step()
        optimizer.zero_grad()

proof.close()
```

`prepare()` accepts any number of values. It prepares recognized stateful
objects, returns unrecognized values unchanged, returns its sole value for one
argument, and returns a tuple for multiple arguments. A prepared optimizer
retains native eager behavior outside attestation; inside `@proof.attest`, its
`step()` lowers through the matching functional `torch.optim` implementation.
All standard optimizers with a standalone functional API are supported. ZKDV
rejects `LBFGS` and unknown optimizer subclasses with instructions to call
`proof.register_functional_optimizer(...)`.

Prepare the stateful root model and the optimizer whose transition is
attested. Additional stateful modules captured by the attested function must
be composed under that root; otherwise lowering fails with a `prepare()`
instruction. Stateless modules, tensors, batches, configuration values, and
ordinary application objects do not require preparation and pass through
unchanged.

Open `proof.transaction(...)` before the ordinary model forward so mutable
module buffers are included. During attestation, use `proof.forward()` and
`proof.backward()` to expose the module call and gradient boundary to AOT
lowering. Generate randomness outside the attested transition and include it
in the committed batch. Implicit random operators are rejected because replay
must recompute the same deterministic transition.

The first transaction seals an AOT Inductor artifact for replay and commits
its exact bytes. Parameter, buffer, optimizer, and batch commitments remain
rank-local under DDP and composable FSDP. Periodic replay snapshots use a
bounded pinned-host pool, flow-controlled CUDA copies, and explicit mutation
barriers. Unchecked policies retire production deltas; checked flow donates
dead reconstruction buffers before replay.

A fully commented runnable example is available at
[`examples/torch/basic_training.py`](examples/torch/basic_training.py).

## Flax `TrainState`

`zkdv.jax.contrib.TrainState` inherits `flax.training.TrainState`, so applications
may inherit it again to add ordinary Linen state fields. Its
`apply_gradients()` implementation performs the same Optax transition and also
registers the resulting additive update with the active `proof.jit` trace.

Create it through Flax's standard `TrainState.create(...)` API and place the
result with ordinary JAX transforms such as `jax.device_put` or `jax.jit` with
`out_shardings`. The application supplies `params` and `tx`; TrainState creates
and carries all optimizer state on the application's behalf. The same state
behaves as an ordinary Flax state under `jax.jit` and registers its transition
automatically when the function is decorated with `proof.jit`.

The sampled program does not guess where the loss lives. ZKDV replays the
staged training function through its registered optimizer transition with
challenged parameters, committed optimizer state, and the invocation's dynamic
non-state inputs. It stops at that symbolic boundary and projects the update
and optimizer coordinates selected by the verifier; the unused user return
path and parameter application are not exported. Static JIT arguments are
fixed by the first program commitment. If the production `apply_fn` cannot
pass `jax.export`, supply `attest_apply_fn`; it replaces `state.apply_fn` only
in sampled replay.

Donating the public state argument aliases only its parameter buffers in the
lowered executable. Optimizer history remains available to a possible sampled
check, preserving the same memory and protocol boundary as the explicit API.

## The two compiled functions

### `ZKDV.jit`

`ZKDV.jit` accepts the ordinary `jax.jit` options and can be used as either a
decorator or a callable:

```python
@proof.jit(in_shardings=..., out_shardings=..., donate_argnums=(0,))
def train_step(...):
    ...

# Equivalent:
train_step = proof.jit(train_step, donate_argnums=(0,))
```

The decorated function may have an arbitrary input ABI and return any PyTree.
A contrib state declares the transaction through `apply_gradients()`.
Lower-level integrations instead declare exactly one transaction explicitly:

```python
transaction = zkdv.jax.Transaction(
    batch=batch_before,
    params=params_before,
    opt_state=opt_state_before,
    overlap=overlap_policy,
)

# Compute the real update normally.

transaction.update(
    deltas=additive_parameter_update,
    params=params_after,
    opt_state=opt_state_after,
)
```

These calls are symbolic markers while JAX traces the function. They do not
open fictitious native transactions during tracing or compilation. The outer
wrapper opens exactly one transcript transaction for each actual invocation.
The overlap policy is host-side transaction metadata: changing it does not
retrace the function, add an XLA conditional, or copy a decision from a device.

Choose the policy separately for every invocation:

- `zkdv.overlap.ALWAYS` releases the training admission slot as soon as a
  sampled check has acquired its device inputs. Training and verification may
  then overlap, using two training generations plus verification workspace.
  The trusted core queues the eventual `CheckResult`; a failed result poisons
  subsequent operations and is always observed by `close()`.
- `zkdv.overlap.UNCHECKED` pipelines ordinary commitments but makes a sampled
  check a fence. At most one already-admitted successor is retained while the
  check completes. This is the default.
- `zkdv.overlap.NEVER` admits no successor until the current commitment and any
  sampled check retire. This is the lowest-residency, fully sequential mode.

The policies bound live buffer generations, not exact allocator bytes. XLA
temporaries and the attested program determine the remaining workspace.

Parameter inputs may be donated. Pre-update optimizer state cannot be donated,
because a sampled verifier may still need it; the explicit API rejects that
donation with a semantic error. Donating a contrib `TrainState` lowers to
parameter-only donation automatically.

### `ZKDV.attest`

The attested function is the hermetic program used to recompute one sampled
parameter update and one sampled post-update optimizer-state value. Its input
ABI is fixed:

```python
def attested_step(
    params_before,
    opt_state_before,
    batch,
    parameter_index,
    optimizer_state_index,
):
    ...
    return selected_delta, selected_post_optimizer_state
```

Both outputs must be scalar floating-point JAX values. The program must be
exportable by JAX and reproduce the same update represented by the transaction.
It may be a slower, more portable implementation than the production training
function. Both indices are scalar `uint64` values, so the ABI covers the full
index range committed by the protocol without enabling x64 for ordinary
training computations. `zkdv.jax.index` selects from arbitrary PyTrees without
materializing a concatenated parameter- or optimizer-sized buffer.

`ZKDV.attest` supports the `jax.jit` input, static-argument, donation, backend,
device, inline, and compiler options. Its output layout is fixed by the
protocol, so an explicitly supplied `out_shardings` is ignored with a warning.

## Lifecycle and configuration

```python
config = zkdv.ZKDVConfig(
    name="pretraining",
    sampling_policy=zkdv.SamplingPolicy(
        state_check_probability=0.01,
        update_check_probability=0.01,
    ),
)
proof = zkdv.jax.ZKDV(
    "proof",
    config,
    max_in_flight=2,
    replay_snapshot_interval=12,
)
proof.start(rolling_window=32, probe_ratio=0.01)
```

- `rolling_window` controls rolling hashes within each committed batch row; it
  must not exceed that row's packed width. It is unrelated to the number of
  training steps.
- `probe_ratio` is the fraction of each parameter leaf's last dimension retained
  by its orthonormal projection (default `0.01`) and remains fixed for the run.
- `max_in_flight` is the hard ceiling for unresolved Python-side pipeline
  transactions. Per-transaction overlap policies may tighten that ceiling; the
  default of two provides the intended `ALWAYS` and `UNCHECKED` behavior.
- `replay_snapshot_interval` bounds checked replay. The default, `12`, launches
  an asynchronous full pre-state copy every 12 steps and replays at most 11
  preceding transitions. A value `N` launches that copy every `N` steps;
  choose `1` to copy every step and check exactly one transition. Copies target
  JAX compiler-addressable host memory (pinned when available) and preserve
  donation. A background dispatcher launches one precompiled whole-tree D2H
  before the due update. The training thread waits only until JAX has registered
  the source read; JAX then prevents donated storage from being reused before
  that read completes. Unchecked flow never reads the host result; only a
  sampled check does. This ordering is portable and exact, but it does not by
  itself prove that a backend overlaps transfer and training. For check
  probability `p`, `1 / p` is the expected number of
  steps between checks; choose `N` directly as the maximum acceptable replay
  length rather than relying on that expectation. ZKDV prefers pinned-host
  placement and otherwise uses the compiler-addressable host memory kind
  exposed by the active JAX backend. A bounded JAX host pool is primed before
  training and donated into periodic copy outputs, avoiding cold pinning
  without retaining an additional device-state tree.
- `SamplingPolicy.DEBUGCheckEvery(n)` is available for deterministic testing
  and profiling; use the probabilistic policy for normal operation.
- Call `join()` to wait without closing, or `close()` once at the end of the
  run. Training code does not need to join individual transactions.

The parameter and optimizer PyTrees, input representation, and device/sharding
layout must remain compatible with the first compiled invocation. For
multi-device `NamedSharding`, ZKDV preserves every parameter, delta, and
optimizer leaf's existing layout. Session probe maps and small control inputs
are replicated over that same mesh; JAX supplies the required portable
cross-shard reductions. DDP replicas therefore commit one logical update, while
tensor-parallel shards contribute their disjoint coordinates exactly once.
The protocol commits a rectangular `uint32` batch. The contrib lowering
losslessly encodes every dynamic non-state input into that representation. The
explicit `Transaction` API continues to accept an application-provided packed
batch directly.

## Documentation

Preview the complete site, including rustdoc:

```console
scripts/serve-docs
```

For live Markdown/CSS reloads without the linked rustdoc pages, use
`uv run --group docs zensical serve` instead.

Build the complete static site:

```console
scripts/build-docs
```

The build treats documentation warnings as errors, renders the Python API from
the package source, and embeds rustdoc for the public protocol and hashing
crates.

