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:

pip install 'zkdv[jax]'

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. The explicit transaction API remains demonstrated by examples/jax/basic_training.py.

Torch 2.13 quick start

Install only the optional Torch frontend:

pip install 'zkdv[torch]'

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.

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:

@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:

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:

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

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:

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:

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.