Skip to content

Your first training run

In this tutorial you will connect a Jax or Torch training job to DVBench to attest what datasets you used during training.

1. Install DVBench

pip install "dvbench[torch]"

pip install "dvbench[jax]"

2. Connect a Client

Begin by connecting to DVBench.

from dvbench import Client  # import before Jax/Torch
client = Client(API_URL)

What's API_URL?

As we are currently in (probably very broken) alpha testing, you will receive API_URL from someone, presumably whomever you found out about this from... Sorry for the wet paint.

3. Wire up a Neural Network

We can't help you with this part, but we trust you've got it :)

import torch
from torch import nn
from torch.nn import functional as F

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = nn.Linear(10, 2).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

import jax
import jax.numpy as jnp
import optax
from flax import linen as nn

class TinyClassifier(nn.Module):
    @nn.compact
    def __call__(self, tokens):
        return nn.Dense(2)(tokens.astype(jnp.float32))

model = TinyClassifier()
params = model.init(
    jax.random.key(0),
    jnp.ones((1, 4), dtype=jnp.uint32),
)["params"]
optimizer = optax.adam(1e-2)
opt_state = optimizer.init(params)

Using Flax TrainState?

This tutorial follows ZKDV's basic explicit transaction API. A separate how-to guide will show how zkdv.jax.contrib.TrainState declares the same optimizer transition through Flax's familiar apply_gradients() workflow.

3. Configure your proof tape

This is where you configure the proof tape, and you'll have to make some decisions.

  • name: run names must be kebab-case, with a dash like/this-example.
  • rolling_window: every rolling_window number of tokens is considered a sample.
# register an experiment
experiment = client.experiment("your-name/your-run")

# initialize your zkdv tape
proof = experiment.proof(backend="torch")
proof.start(rolling_window=32)

# prepare your Torch models for tracing
model, optimizer = proof.prepare(model, optimizer)

# register an experiment
experiment = client.experiment("your-name/jax-example")

# initialize your zkdv tape
proof = experiment.proof(backend="jax", config=config)
proof.start(rolling_window=4, probe_ratio=0.01)
Tap for more on rolling_window

A rolling_window of 4 means that every 4 tokens is considered a "sample." For instance, the quick brown fox and brown fox the quick are considered two distinct samples when rolling_window=4, but considered a shuffling of the same sample when rolling_window=2. A longer window results in a shorter tape, but offers less flexibility in shuffling.

3. Declare the replayable update

We need a deterministic description of how you got from one parameter to the next.

@proof.attest
def attested_step(model, optimizer, batch):  # you must keep this signature
    logits, loss = proof.forward(model, **batch)  # instead of model(**batch)
    proof.backward(loss)  # instead of loss.backward()
    optimizer.step()
    optimizer.zero_grad()

@proof.attest
def attested_step(params, opt_state, batch):  # you must keep this signature

    def loss_fn(current_params):
        logits, loss = model.apply({"params": current_params}, batch)
        return loss
    gradients = jax.grad(loss_fn)(params)

    # manually compute parameter updates and next state
    deltas, next_state = optimizer.update(gradients, opt_state, params)

    # return a tuple of PyTrees: (updates, next optimization state)
    return deltas, next_state

This function has to be deterministically serializable. Common culprits that prevent this include custom kernels. Rest assured, we don't call this function that many times (usually <1% of training), so you can definitely wire up a "slow-path" here without loosing MFUs.

4. Annotate your training

Your training loop remains largely the same, and can contain any custom logic:

for _ in range(...):
    batch = ...

    with proof.transaction(batch=batch):
        logits, loss = model(**batch)
        proof.backward(loss)  # instead of loss.backward()
        optimizer.step()
        optimizer.zero_grad()

@proof.jit(donate_argnames="params")  # instead of jax.jit, same signature
def train_step(params, opt_state, batch):

    # open a transaction
    transaction = zkdv.jax.Transaction(
        batch=batch,  # declare what data is used
        params=params,  # current parameters
        opt_state=opt_state,  # current optimizer state
    )

    # train normally
    deltas, next_opt_state = compute_update(params, opt_state, batch)
    next_params = optax.apply_updates(params, deltas)

    # update transaction with the information produced
    transaction.update(
        deltas=deltas,
        params=next_params, 
        opt_state=next_opt_state,
    )
    return next_params, next_opt_state

# train loop stays the same
for _ in range(...):
    batch = ...
    params, opt_state = train_step(params, opt_state, batch)

5. Finish the proof

Call finish() before evaluation.

proof.finish()
client.close()

save_your(model.state_dict())  # do so after proof is finished

proof.finish()
client.close()

save_your(params)  # do so after proof is finished

A complete working example is in examples/jax/proof_jax.py and examples/torch/proof.py.

Next: evaluate the checkpoint.