VolvoxAI for Python

Use VolvoxAI from Python to run models on CPUs and native GPUs, tokenize text, plan model execution, schedule work, train or adapt models, and perform post-training quantization. ONNX conversion and command-line PTQ are included. Python calls the same C engine and generated protobuf API as native applications and the JavaScript/WASM package.

Generated messages are standard protobuf classes, exposed as volvoxai.pb. Construct nested messages with keyword arguments; use HasField("buffer") to check a message field's presence, CopyFrom to replace it, and extend for repeated fields. SerializeToString() and FromString() read and write the same schema-defined wire format. Python clients use these operations directly with the existing Synurang module host. The wheel declares protobuf==7.35.1; the qualified CPython runtime uses its native upb implementation.

InferenceSession and AsyncInferenceSession load a model, reuse its compiled execution context and accept and return NumPy arrays. quantize applies weight and activation policies, consuming calibration batches when needed; TrainingSession supports training, saving and checkpoint resumption. Native libraries, CPU threads and result lifetime are managed automatically. Failures raise Python exceptions at the call site.

The package includes the native engine library, its shared-library loader, all eight generated synchronous and asynchronous service clients, the protobuf messages, and the pinned Synurang runtime. It also bundles the native inference runner used by model tooling. No source checkout, compiler, CUDA toolkit or Docker installation is needed to use an installed wheel.

Installation and supported environments

The initial binary distribution targets Linux x86_64 with glibc 2.28 or later. The wheel uses the C ABI through ctypes and a CPython stable-ABI DLPack capsule adapter. One cp310-abi3 wheel serves CPython 3.10–3.14. CPU, CUDA, Vulkan and OpenGL backends are compiled in. GPU execution requires the corresponding host device and driver; CPU execution works without them. CUDA kernels use a compute-capability 7.5 PTX baseline. Windows, macOS and Linux ARM64 wheels are not part of this release.

After publishing to PyPI, install the latest published release:

python3 -m venv .venv
. .venv/bin/activate
python -m pip install volvoxai

To install the local build before publishing:

# Run from the repository root.
VOLVOXAI_VERSION="$(python -c 'import json; print(json.load(open("package.json", encoding="utf-8"))["version"])')"
python -m pip install "dist/python/$VOLVOXAI_VERSION/"*.whl
volvoxai --version

NumPy, ONNX and SafeTensors are installed as dependencies. The optional volvoxai[tflite] extra installs the TensorFlow Lite importer dependencies. Hugging Face downloads can use the separately installed huggingface_hub tool. PyTorch is an optional DLPack consumer/producer, not a dependency of the wheel. No separate DLPack library is installed.

Run a model

Convert an ONNX model into a VolvoxAI package. Bind unresolved dimensions with --input-shape or --dimension-bound; volvoxai export --help describes all supported conversion options.

volvoxai export --model model.onnx --out model/model.safetensors

The resulting graph.json and model.safetensors are also usable by the npm and native distributions. Conversion does not change application preprocessing.

For a model with one float32 input of shape [1, 3]:

import numpy as np
import volvoxai as vx

values = np.array([[1.0, 2.0, 3.0]], dtype=np.float32)
with vx.InferenceSession("model") as session:
    outputs = session.run(values)
    for name, array in outputs.items():
        print(name, array.shape, array.dtype, array)

"model" is a directory path. FP32, INT8 and mixed precision are read from the graph and weights; there is no precision argument. A quantized model can still require float32 application inputs. Preserve the input dtype declared by the model instead of casting inputs to INT8 because its weights are quantized.

The session finds graph.json (or a named *.graph.json) and its SafeTensors weights. You can pass the directory, graph file or corresponding weight file. Omit the model path to discover a single package in the current directory or its immediate child directories:

with vx.InferenceSession() as session:
    outputs = session.run(values)

If multiple models exist, such as separate encoder/decoder packages or FP32 and INT8 exports, pass the intended model path. The SDK raises an ambiguity error instead of choosing one. Named graphs prefer their matching weight filename. If a graph directory contains multiple weight files, select its files or shards explicitly with weights=["path/to/shard1.safetensors", "path/to/shard2.safetensors"]. These paths are relative to the working directory. ONNX conversion remains the explicit export step above.

For a local dense Hugging Face Llama, Qwen2 or Qwen3 checkpoint, use vx.InferenceSession.from_pretrained("checkpoint", max_context=128). It builds the graph from configuration and loads the original F32/F16/BF16 SafeTensors shards without creating an F32 weight export. Execution uses F32; tokenization remains with your application. Scaled RoPE uses small temporary position tables that are removed after loading. This method never downloads files. See the decoder workflow for complete inputs, CUDA, incremental rows and W4A16/W4A8 package export.

For multiple inputs, use names from session.inputs, a tuple of immutable Python TensorSpec objects. session.outputs exposes the output contracts. The converter may assign canonical input names such as input0:

with vx.InferenceSession("model") as session:
    print(session.inputs, session.outputs)
    outputs = session.run({session.inputs[0].name: values})
    # Repeated calls reuse the model and execution context.
    selected = session.run(values, output_names=[session.outputs[0].name])

Inspect names, shapes and dtype strings without interpreting protobuf enums:

with vx.InferenceSession("model") as session:
    for spec in session.inputs:
        print(spec.name, spec.shape, spec.dtype)  # e.g. input0 (1, 3) float32
        for symbol, bounds in spec.constraints.items():
            print(symbol, bounds.min, bounds.max, bounds.multiple_of)

A dynamic contract can have shape=(1, "S", 320) with constraints["S"] == vx.Dimension(min=1, max=256, multiple_of=1). These describe the legal domain; each returned array has its actual integer shape. Metadata, its shape tuple and dimension constraints are read-only. Storage formats without a NumPy representation still have dtype names such as bfloat16 or float8_e5m2 in metadata.

run() always returns a dictionary of output names to owned, writable NumPy arrays. Arrays remain valid after later runs and after session close. Dtypes and concrete shapes are preserved; C validates each complete input batch. Strided and non-native-byte-order input arrays are supported. Packed output formats with no supported NumPy dtype raise TypeError; their bytes remain available in an inline Execute packet or through CopyTensors for retained outputs.

Contiguous, aligned little-endian inputs are borrowed until run() returns; do not modify them from another thread during the call. Other layouts are normalized without changing precision or scalar rank. Native input buffers and direct writes into output arrays avoid protobuf tensor-payload copies. C still owns execution storage and output snapshots, so this is not a promise of zero copies throughout the engine. Static output shapes need no extra metadata query; dynamic outputs query the completed result's concrete shapes.

The default backend is CPU. CPU thread count is selected by the C engine; on Linux x86_64 this accounts for CPU affinity and physical cores. VOLVOXAI_THREADS or an explicit positive cpu_threads can override it. None and 0 select automatic threading. GPU execution requires only a backend name, such as "cuda", "vulkan" or "opengl":

with vx.InferenceSession("model", backend="cuda") as session:
    outputs = session.run(values)
    print(session.backend)  # The actual compiled backend.

A requested GPU is required and cannot silently fall back to CPU. session.report retains compilation and route evidence. The session waits for pending execution, reads only requested outputs and releases each native result. Calls on one session are serialized. Close it with a with block or close(); this retires its host, model and execution context. Use generated services for scheduling or incremental decoding.

Keep tensors on the GPU and share them through DLPack

Use run() when you want NumPy outputs. Use run_tensors() when outputs should remain in native memory, particularly between GPU inference calls. It accepts NumPy arrays, VolvoxAI tensors, and compatible DLPack producers. The default backend is CPU. Choose cuda, vulkan, opengl or metal to retain GPU outputs on that backend, subject to the platform/build support below.

For a model package with one float32 input of shape [1, 3] and output logits:

import torch
import volvoxai as vx

x = torch.tensor([[1.0, 2.0, 3.0]], dtype=torch.float32, device="cuda")

with vx.InferenceSession("model", backend="cuda") as session:
    with session.run_tensors(x) as outputs:
        logits = outputs["logits"]
        print(logits.shape, logits.dtype, logits.device)
        # For example: (1, 10), float32, cuda:0

        scores = torch.from_dlpack(logits)  # Shares the CUDA output buffer.
        predicted_class = scores.argmax(dim=-1)
        print(predicted_class.cpu().tolist())  # Only the selected IDs go to CPU.

# The external view still owns its buffer after Tensor and session close.
print(scores.shape)
del scores

Install a PyTorch distribution suitable for your CUDA driver separately when using this integration. Native C owns VolvoxAI tensor storage and device lifetime. DLPack does not automatically translate device allocations between different GPU APIs. See the DLPack Python protocol.

Backend	Pass a VolvoxAI tensor to another session	External sharing
CPU	Yes	NumPy and CPU DLPack
CUDA	Yes, on the same device	CUDA DLPack, including PyTorch
Vulkan	Yes, on the same VkDevice	Exported native VkBuffer; imports currently require a module-owned buffer
OpenGL	Yes, in the same EGL context	Exported native buffer name; imports currently require a module-owned buffer
Metal	Adapter implemented; macOS execution not yet qualified	Same-device MTLBuffer and compatible kDLMetal DLPack; four-byte aligned offsets and byte extents

Vulkan/OpenGL use the native buffer descriptor, so Tensor.__dlpack__() raises BufferError for those backends. Pass retained tensors directly between VolvoxAI models. Graphics-API producers must finish writes before a call. Metal DLPack objects must carry an actual id<MTLBuffer> as required by the DLPack C structure; a framework's GPU support alone does not establish compatibility. The initial Metal adapter uses buffer blits, whose macOS alignment requirements require four-byte offsets and byte lengths.

You can also connect two compatible models without installing a GPU array library. For example, assume an encoder returns memory and a decoder accepts that tensor plus an int32 tokens input:

import numpy as np
import volvoxai as vx

with vx.Runtime() as runtime, \
     runtime.inference_session("encoder", backend="cuda") as encoder, \
     runtime.inference_session("decoder", backend="cuda") as decoder:
    image = np.load("encoder_input.npy")  # Exact encoder input dtype/shape.
    with encoder.run_tensors(image, output_names=["memory"])["memory"] as memory:
        result = decoder.run_tensors({
            "memory": memory,
            "tokens": np.array([[1]], dtype=np.int32),
        }, output_names=["logits"])
        with result["logits"] as logits:
            host_logits = logits.numpy()  # Explicit, independent host copy.
            print(host_logits.shape, host_logits.dtype)

Use one Runtime owner and the same backend for both sessions; the example also works with vulkan or opengl and their drivers. A CPU input uploads once, the intermediate memory stays on the GPU, and numpy() performs the final explicit download.

Runtime(budget=vx.pb.RuntimeBudget(...)) configures the existing engine admission limits for every session sharing that owner. Unset fields retain their engine defaults. Large decoder captures may need an explicit max_unconsumed_result_bytes: admission reserves the selected outputs at their declared maximum shape, rather than just the current sequence length. The default is 64 MiB. Retained output tensors and their access leases remain charged until their final release. See RuntimeBudget for all limits.

Choose the actual input/output names and shapes from each model's inputs and outputs metadata. Each call returns TensorOutputs, a dictionary of Tensor objects. Use with session.run_tensors(...) as outputs: or outputs.close() to release a group in one native call. Individual tensor context managers and tensor.close() also work. Existing DLPack views retain their own leases when the group closes. CPU tensors can also be shared with np.from_dlpack(tensor). tensor.numpy() always makes an independent host copy. Tensor has no implicit NumPy conversion that downloads a GPU buffer.

For repeated execution, keep unchanged inputs and feedback state inside the session. Supply a complete batch for the first call; later calls can combine new inputs, reuse_inputs, and feedback. For example, a decoder whose cache input is past, output is present, and declared entrypoints are initialize and transition can use:

with decoder.run_tensors(initial_inputs, entrypoint="initialize", output_names=["logits"]) as outputs:
    token = int(outputs["logits"].numpy().argmax())

with decoder.run_tensors(
    {"token": np.array([[token]], dtype=np.int32), "position": np.array([1], dtype=np.int32)},
    entrypoint="transition",
    reuse_inputs=["memory"],
    feedback={"past": "present"},
    output_names=["logits"],
) as outputs:
    token = int(outputs["logits"].numpy().argmax())

Use your model's actual names and supply any additional control inputs. Together, the three input sources must bind each input exactly once. Feedback uses the previous internal graph output, including outputs that were not returned. Editing an exported tensor cannot change that internal state. output_names=[] executes without exporting snapshots. Output selection does not skip graph computation. run() and run_tensors() share the same bindings and context state. Call decoder.new_context() before initializing a new sequence; it retains the compiled model and existing exported tensors. Supplying a complete input batch does not reset a context. Stateless models with one full evaluation entrypoint can omit entrypoint; other models require an explicit name or scope=vx.pb.ExecutionScope(...). The model’s session.execution metadata lists the supported entrypoints and ranges.

CUDA, Vulkan and OpenGL reuse unchanged input buffers in their reserved domains. Feedback copies directly from the previous internal output to its next input, eliminating the intermediate exported snapshot. Aliased sources and movable layouts use private snapshots to preserve simultaneous binding semantics. CPU and Metal use this safe snapshot path. Device-produced indices or routing values that require host validation must still be supplied as host arrays. Validation failures preserve prior state. An execution failure after input commit leaves the context failed; call session.new_context() and initialize it with a complete input batch before running again.

The memory and execution contract is explicit:

• Export through DLPack shares output storage. Inputs are copied into the engine's execution arena and outputs are copied into independent snapshots; GPU inputs and snapshots use device-to-device copies. The complete engine execution is therefore not a zero-copy operation.
• Released allocations are reused by the session, bounded to 64 idle buffers and 64 MiB. Session close frees idle storage; live outputs are not reused. CUDA/Vulkan/OpenGL group input and output copies before waiting for completion.
• Calls wait for GPU completion. Input producers establish ordering on the CUDA legacy default stream, including when PyTorch uses other streams. Vulkan, OpenGL and Metal producers complete writes before the call; raw native handles do not encode cross-queue fences or context switching. Output snapshots are ready before export. This first implementation favors a predictable synchronous contract; it does not overlap inference calls.
• Outputs survive subsequent calls and session close. External NumPy/PyTorch views retain their storage after the VolvoxAI wrapper closes. Those DLPack consumers hold exclusive access; release that external lease before reusing the original VolvoxAI tensor. Passing the modified PyTorch object itself allows its producer protocol to establish the required stream ordering.
• DLPack inputs need a supported exact dtype, dense contiguous layout and, for CUDA inputs, the session's CUDA device and primary context. Strided views, foreign devices, managed allocations and unsupported DLPack devices are rejected instead of being implicitly moved or converted. Make an explicit contiguous copy in the producing library when needed.
• Index and route inputs that require host-side value validation must remain NumPy arrays. This includes some embedding/gather indices. Large encoder activations can remain on the GPU while those small control inputs stay on CPU.
• DLPack shares tensor data, not autograd history. Detach PyTorch tensors with gradients before inference. Packed low-bit types are outside this initial dense tensor interface. AsyncInferenceSession continues to provide the asynchronous NumPy workflow; run_tensors() currently belongs to the synchronous InferenceSession.

Shared owners and generated buffer operations

Use vx.Runtime to pass buffer IDs between sessions in one native owner:

with vx.Runtime() as runtime:
    with runtime.inference_session("encoder") as encoder:
        with runtime.inference_session("head") as head:
            with encoder.run_tensors({"x": x}) as encoded:
                with head.run_tensors({"hidden": encoded["hidden"]}) as output:
                    scores = output["scores"].numpy()

Replace paths, names and shapes with the model contract. Standalone sessions have separate owners; sharing between them requires an explicit DLPack import or copy. runtime.tensor(value) produces an independent CPU snapshot. runtime.tensor(value, copy=False) imports a dense CPU/CUDA DLPack owner once. tensor.numpy(copy=False) acquires a direct read-only CPU mapping; writable=True requests an exclusive write lease.

The generated API returns Tensor.buffer (BufferView), containing an ID, offset and logical byte length. Reuse this descriptor in another model input on the same owner. Use VxBufferServiceClient.copy_tensors for independent CPU snapshots, inline readback or copies into BorrowedBuffer host destinations. retain_buffers creates fresh IDs; release_buffers retires them in a batch. Only BeginBufferAccess exposes a pointer/resource under an explicit lease. BorrowedBuffer inputs can be carried in the execution call without a separate import RPC. Names and shapes remain model contracts validated in C.

DLPack exports acquire an exclusive external lease, because consumers may write. Other native reads/inputs/exports of that buffer return BUSY until the consumer releases it. A NumPy numpy(copy=False) read view permits other readers, and its lease survives Tensor/session close. Shared trainer parameters use read-only leases; step/commit/rollback/reset return BUSY while they live. Independent parameter snapshots do not block training. No memory-sharing operation connects VolvoxAI autograd to PyTorch.

For bounded external CPU/CUDA access, with tensor.torch_access() as view: ends permission at scope exit while remaining aliases retain only allocation lifetime. CUDA records PyTorch's current stream on the tensor's device; supply streams=[first, second] when multiple streams access the view. Order their conflicting work yourself and stop using every alias after the scope. Native input/readback/reuse then waits only for that allocation's recorded consumers. Unrelated buffers remain usable. Ordinary torch.from_dlpack(tensor) continues to use automatic lifetime and a conservative CUDA drain on final release.

See buffers and tensors for complete inference, training, generated C and transport workflows, including current GPU and WebGPU restrictions.

Native C uses these same messages without the Python adapter; see the complete C client and before/after benchmark. GPU copies inside the engine are included in those measurements.

See the proto schema for the exhaustive operation contracts and native integration for C dispatch.

Async inference

import asyncio
import numpy as np
import volvoxai as vx

async def main():
    values = np.array([[1, 2, 3]], dtype=np.float32)
    async with vx.AsyncInferenceSession("model") as session:
        outputs = await session.run(values, timeout=30)
        for name, array in outputs.items():
            print(name, array.shape, array.dtype, array)

asyncio.run(main())

Model initialization happens on entering the async context. Outside a context, use await session.open() and await session.close(). Metadata and output selection match the synchronous API. A session belongs to one event loop and serializes its runs. It uses the generated async services; native work can run while the event loop handles other tasks.

Each run snapshots its inputs when it acquires the session. Keep inputs stable until that run returns, including while it waits behind another run. timeout covers that wait and execution and raises asyncio.TimeoutError. Cancellation raises asyncio.CancelledError. If native work has already begun, the adapter waits for it to finish and releases the result before returning the exception. This safe drain can exceed the timeout; it keeps arrays alive until C has stopped using them. Cancellation while queued submits no work. A completed cancellation leaves the session reusable.

Profile and debug a session

InferenceSession.trace() measures the calls inside a with block, like torch.profiler.profile. Leaving the block stops the trace and reads its typed records; export_chrome_trace() writes JSON for Perfetto:

import numpy as np
import volvoxai as vx

with vx.InferenceSession("path/to/model") as session:
    with session.trace(detail="nodes", memory=True) as trace:
        with trace.annotate("frame 1"):
            session.run(np.load("input.npy"))
    trace.export_chrome_trace("trace.json")
    nodes = [e for e in trace.events if e.HasField("node")]

InferenceSession.debug() starts an isolated, steppable execution. By default it compiles a separate model with preserve_node_boundaries=True, so each source node remains observable. Ordinary session.run() keeps using its original optimized model. Pass preserve_node_boundaries=False to inspect that optimized schedule instead.

The debugger pauses before the first node. continue_(break_on_nonfinite=True) stops after the first captured output containing NaN or infinity; tensor() returns a snapshot as a NumPy array in its storage dtype:

with vx.InferenceSession("path/to/model") as session:
    with session.debug(np.load("input.npy", allow_pickle=False)) as debug:
        debug.continue_(break_on_nonfinite=True)
        if debug.stop_reason == vx.pb.DebugStopReason.DEBUG_STOP_REASON_NONFINITE:
            event = debug.events()[-1]
            print(debug.plan.steps[event.step].source_node_ids)
            for snapshot in event.snapshots:
                if snapshot.status == vx.pb.DebugTensorStatus.DEBUG_TENSOR_STATUS_AVAILABLE:
                    print(debug.tensor(snapshot))

A model that fails while debugged sets debug.failure instead of raising. Debugging requires the full library bundled in the wheel. If VOLVOXAI_LIBRARY explicitly selects libvolvoxai-lite, select the full libvolvoxai before creating the session; a live session cannot switch libraries. build_profile reports the loaded profile. Close each debug scope before closing its parent session.

debug.collect() reads the session info, plan, every event and available raw tensor bytes into a JSON-compatible volvoxai-debug/v1 artifact. Collect while the debug scope is open; the returned artifact survives its close. Collection does not advance execution: a paused session yields a partial capture. Use collect(values=False) for metadata only, or constants=False when opening the debugger to leave weights out of the capture. See compare Python captures for backend and FP32/INT8 comparisons, and profiling for trace collection.

Compare two saved captures without running an engine:

report = vx.debug_compare("float.json", "int8.json", mode="quantization", min_sqnr_db=20)
print(report["comparable"], report["activationCoverage"], report["worst"])

The same function accepts capture dictionaries. Use mode="exact" with atol/rtol for the same precision across backends or schedules; mode="auto" chooses from the recorded dtype and quantization. Reports use the same JSON fields and formats as the JavaScript comparison tool.

Handle errors

Catch an exception where your application can handle a failure. Uncaught exceptions include a traceback pointing to the failed call:

try:
    with vx.InferenceSession("model") as session:
        outputs = session.run(values)
except vx.VolvoxAIError as error:
    print(error.operation, error.status, error.code, error.stage)
    print(error.report)  # Original typed diagnostics and backend evidence.
    raise

The same rule applies to all synchronous and asyncio service clients exported by volvoxai, including release methods that return an OperationReport directly. Every non-OK operation status raises VolvoxAIError. The exception retains report, response, operation, status, code, stage, backend and the offending node index when present. Classify failures with proto enums; message text is descriptive and does not determine error identity. A required report missing from a response raises an internal error.

Successful responses keep their reports for backend and execution diagnostics. Execute and Generate return complete selected outputs when the call finishes. Scheduled callers wait for a terminal request with WaitRequest, then transfer its outputs once with TakeRequestResult. Nested plan/revision evidence and status descriptions are data; they are not recursively treated as failed calls. For BUSY, overload or deadlines, inspect error.status and choose recovery according to that operation's contract. The SDK performs no automatic retry.

vx.FfiError carries Synurang call failures and its grpc_code, including RPC deadlines and unavailable methods. vx.PluginClosedError is its subclass for a closed host. These remain distinct from an engine operation's report. Async task cancellation propagates as asyncio.CancelledError. Awaiting an async method raises at the await expression.

The package applies this policy to the generated methods after their single response decode. Import service clients from volvoxai; the flat volvoxai_client module is the raw generated transport binding and returns operation reports without the package's exception policy. The former public check() helper has been removed.

Public API across languages

All three languages use the operations and types in proto/volvoxai.proto. Their call and error conventions follow the host language:

Language	Public call surface	Failure handling
C	Generated request/response structs and codecs, dispatched through Synurang_GetApi	Check the terminal call status, then the operation report's status
Python	InferenceSession, AsyncInferenceSession, quantize, TrainingSession; Vx*ServiceClient / AsyncClient and vx.pb for the complete API	A direct call or await raises VolvoxAIError; call failures raise FfiError
TypeScript	Vx*ServiceClient, pb objects and EngineHost / FullEngineHost	await rejects with VolvoxAIError; call failures raise RpcError

InferenceSession adapts paths and arrays and sequences existing generated CreateRuntime, LoadModel, GetModelInfo, CompileModel, CreateExecutionContext, Execute, CopyTensors and ReleaseBuffers calls. It introduces no engine operation or numerical implementation. Advanced service callers retain direct access to every proto request, response and owner-scoped handle.

Train, save and resume

Training uses the full library automatically. For a classifier whose exported graph accepts one float32 input of shape [1, 3] and produces two or more class logits named logits:

import numpy as np
import volvoxai as vx

# Replace these illustrative batches with your preprocessed training dataset.
batches = [
    (np.array([[1, 2, 3]], dtype=np.float32), np.array([1], dtype=np.int64)),
    (np.array([[3, 1, 2]], dtype=np.float32), np.array([0], dtype=np.int64)),
]
with vx.TrainingSession(
    "classifier",
    loss=vx.CrossEntropyLoss(output="logits"),
    optimizer=vx.AdamW(lr=1e-3),
) as trainer:
    print(trainer.inputs, trainer.trainable_names)
    for inputs, targets in batches:
        result = trainer.step(inputs, targets)
        print(result.loss, result.optimizer_step, result.update_applied)
        for metric in result.metrics:
            print(metric.name, metric.correct, metric.examples)
    trained = trainer.save("trained-classifier")
    trainer.save_checkpoint("checkpoints/step-2.pb")

with vx.InferenceSession(trained) as session:
    print(session.run(batches[0][0]))

Loss, backward propagation and weight updates execute in C. The convenience layer currently exposes cross entropy, AdamW and SGD, the choices supported by the trainer contract. Labels must be integers that fit int32; int64 NumPy labels are accepted after bounds checking. For multiple model inputs, pass a name-to-array mapping. For multiple losses, pass a sequence of uniquely named CrossEntropyLoss objects and a target mapping keyed by those loss names.

All float32 weights are selected by default. Use trainable_names=["classifier.weight", "classifier.bias"] to train a subset using names from your model's SafeTensors files. Unsupported graphs or parameter choices raise the C engine's typed error. optimizer=None keeps the current C configuration: a new trainer defaults to AdamW with learning rate 0.001, beta1 0.9, beta2 0.999, epsilon 1e-8, zero weight decay and no clipping. vx.SGD(lr=0.01) selects SGD. Both optimizer adapters expose weight_decay and max_gradient_norm; AdamW also exposes beta1, beta2 and epsilon.

step() waits for completion and changes only the trainer's private weights. save() exports those weights with the exact loaded graph for inference. Its ModelPackage result carries path, graph_path and ordered weight_paths and can be passed directly to an inference session, including sharded models. commit() publishes a model revision inside the session and makes the current weights, optimizer and RNG the new rollback baseline. rollback() returns to the last commit, the imported checkpoint, or the initial model state. An explicitly supplied optimizer is applied again on each subsequent step.

Resume with the same graph, loss configuration and trainable selection:

with vx.TrainingSession(
    "classifier", loss=vx.CrossEntropyLoss(output="logits"),
    checkpoint="checkpoints/step-2.pb",
) as trainer:
    result = trainer.step(batches[0][0], batches[0][1])
    print(result.optimizer_step)  # 3

Omitting optimizer preserves the checkpoint's exact optimizer configuration and state. Checkpoints use the generated protobuf format and include weights, optimizer slots, step count, RNG seed and optional application metadata. Pass bytes with save_checkpoint(path, metadata=b"..."); omitting metadata preserves imported metadata. C validates a checkpoint against the loaded model. Loss selection, dataset position and trainable selection remain application settings; store the latter two in application metadata when needed.

step(..., statistics=True) fills result.gradients with the global gradient norm and, for each parameter, its gradient, parameter and update norms; a step that clips reports the norm and clip_scale without the option. A NaN or infinity raises VolvoxAIError with code TRAINING_LOSS_NONFINITE or TRAINING_GRADIENT_NONFINITE, restoring the trainer's private weights:

try:
    trainer.step(inputs, targets, locate_nonfinite=True)
except vx.VolvoxAIError as error:
    print(error.response.gradients.first_nonfinite_parameter)
    print(error.report.offending_node)  # CPU trainers: first node that produced it

For gradient accumulation, set a positive normalizer on each loss for the whole window and pass accumulation_steps=N to each step. flush=True applies a partial window; reset=True discards unfinished accumulation before the supplied batch. Export and checkpoint operations reject unfinished windows. Saved model files and checkpoint paths are never overwritten; use a fresh destination for each export. Close the session to release all native resources.

Library selection and advanced services

Python ships one engine library, and discovery is automatic. An explicit file override (VOLVOXAI_LIBRARY) is authoritative. Otherwise discovery checks VOLVOXAI_LIBRARY_DIR, the bundled library, then the source build. A missing explicit file or a selected library that cannot load raises an error.

open_library() reaches every service: CPU/GPU inference, tokenization, graph planning, scheduling, buffers, training, optimizer/checkpoint operations, adapters and PTQ. The repository also builds a smaller inference-only libvolvoxai-lite for the native and browser releases; it is not installed by the wheel, and the Training and Quantization clients do not work against it.

For advanced training control, load a model through VxInferenceServiceClient, then use VxTrainingServiceClient on the same host to create a trainer, submit training steps, inspect state and export or commit the result. RollbackTrainer restores the pinned baseline. The training integration example demonstrates an SGD update and rollback checked against an independent oracle.

Text processing uses VxTextServiceClient: create a tokenizer, encode text, decode tokens, then release it. VxPlanningServiceClient inspects and edits graph plans. VxSchedulerServiceClient manages execution requests and batch queues. These clients share the selected host and its handle ownership.

Every service also has an AsyncClient for asyncio:

import asyncio
import volvoxai as vx

async def main():
    async with vx.AsyncModuleHost(vx.open_library()) as host:
        platform = vx.VxPlatformServiceAsyncClient(host)
        info = await platform.get_platform_info(vx.pb.Empty())
        print(info.library_version, info.compiled_backends)

asyncio.run(main())

The proto schema is the authoritative API contract and is also bundled under volvoxai/schema/. For live discovery, call VxPlatformServiceClient.describe_api with pb.DescribeApiRequest(include_types=True). It returns the methods and message contracts actually present in the loaded library.

Post-training quantization

Use one workflow for grouped weight storage and static affine regions. Weight storage is independent of packed execution arithmetic; see the quantization guide for the complete contract.

W2 and W3 are explicit CPU/WASM options alongside W4/W8. algorithm="importance" uses representative inputs, and ordered QuantizationRule values select mixed weight precision. See the low-bit tuning workflow. TrainingSession(qat=pb.QatOptions(...)) separately enables weight-only QAT on selected original FP32 nodes; it does not run automatically during PTQ.

import volvoxai as vx

packed = vx.quantize("model", output="model-i4", weight_bits=4,
                     group_size=128, algorithm="mse-clip")
with vx.InferenceSession(packed, packed_arithmetic="f16") as session:
    outputs = session.run(values)

static = vx.quantize("model", output="model-static", group_size=0,
    static_activations=vx.pb.StaticActivations(dtype=vx.pb.DataType.DATA_TYPE_I8),
    calibration_data=calibration_batches())

For a Conv-only weight policy, weight_bits=4, group_size=0, algorithm="mse-clip" with the same StaticActivations fits one scale per output channel and writes real nibble-packed Conv4 weights. Activations remain I8/U8 and bias remains I32; CPU and WASM implement this route. Retain other weight-bearing operators with rules, or pack their Linear weights in a separate grouped plan. Calibration must observe the same activation computation used in deployment.

algorithm is rtn, mse-clip or importance. The engine reports calibration requirements after resolving rules. Static affines and importance fitting use representative batches; RTN/MSE grouped weights need no data. A batch is a NumPy array for a single input or a name-to-array mapping. Input dtype and preprocessing must match the model. Data is consumed once, incrementally. Failed calibration publishes no package. samples_per_batch counts represented examples and cpu_threads fixes calibration workers.

On native, set calibration_backend="cuda" (CLI: --calibration-backend cuda) to run required calibration observations on CUDA. CPU remains the default. CUDA computes activation ranges and input-energy statistics on the device; final weight fitting and package writing remain on the host. result.calibration_backend records the observation backend and is empty for data-free plans. Unsupported CUDA graph operations raise an error without CPU fallback.

Weight-only QAT separately supports TrainingSession(..., backend="cuda", qat=qat_options). CUDA fits and reconstructs the selected W2/W3/W4 weights and applies their straight-through gradient masks, while FP32 masters and optimizer state retain the usual checkpoint and rollback behavior. This does not simulate activation quantization or add GPU inference support to exported W2/W3 or low-bit Conv packages; evaluate those with their supported CPU/WASM arithmetic.

Use generated QuantizationRule messages for mixed precision or retained nodes:

rules = [vx.pb.QuantizationRule(node_id="output", keep_float=vx.pb.Empty())]
packed = vx.quantize("model", output="model-mixed", weight_bits=4, rules=rules)

Rules select node_id, operator_name or node_glob, and choose weights or keep_float. The first matching rule wins. Session packed_arithmetic defaults to f32; alternatives are dynamic-i8, f16 and explicit f32-fast. dynamic-i4 is accepted as a reserved option for future implementation, but compilation currently raises VolvoxAIError with BACKEND_UNSUPPORTED on every backend, including CUDA. It never runs another arithmetic under an A4 label. The resolved enum is available as session.packed_arithmetic, including on AsyncInferenceSession. No named WnAn presets or legacy aliases are accepted.

volvoxai ptq --graph model/graph.json --weights model/model.safetensors \
  --weight-bits 4 --group-size 128 --algorithm mse-clip --out model-i4

volvoxai ptq --graph model/graph.json --weights model/model.safetensors \
  --group-size 0 --static-activations i8 --calibration calibration --out model-static

volvoxai ptq --graph conv-model/graph.json --weights conv-model/model.safetensors \
  --weight-bits 4 --group-size 0 --algorithm mse-clip --static-activations i8 \
  --calibration calibration --out model-conv4

--rules reads an array of protobuf JSON rule objects. Calibration .npz keys are model input names. The full C library performs all numerical work; outputs are graph.json and model.safetensors, and existing files are never overwritten.

Measure PTQ sensitivity

Use vx.ptq_sensitivity to decide which supported nodes should remain in float. It measures model outputs against the original FP32 model with each node quantized alone (isolated) and with that node retained in float while the other selected nodes are quantized (leave-one-float). The second mode reports SQNR recovered over the fully quantized baseline.

Opt into modes=["leave-one-higher"] with higher_weight_bits=3, 4 or 8 to measure raising individual low-bit nodes instead of retaining them in F32. The result adds higherRecoveryDb and bestToKeepHigher; the application chooses its final rules and separately validates task accuracy. Existing default sweeps and default quantization precision are unchanged.

Prepare a separate evaluation/ directory of held-out .npz input batches using the same names, dtypes and preprocessing as calibration/. Then run:

volvoxai ptq sensitivity --graph model/graph.json --weights model/model.safetensors \
  --group-size 0 --static-activations i8 \
  --calibration calibration --evaluation evaluation --out sensitivity.json

Both modes run by default. Repeat --mode isolated and --mode leave-one-float to select modes, and repeat --output <name> to select F32 graph outputs. Use the same --weight-bits, --group-size, --algorithm and --packed-arithmetic as the package you intend to produce. Reports record these settings in configuration. For grouped weights, supply only evaluation data:

volvoxai ptq sensitivity --graph model/graph.json --weights model/model.safetensors \
  --weight-bits 4 --group-size 128 --packed-arithmetic f16 \
  --evaluation evaluation --out packed-sensitivity.json

Every candidate uses the same quantization plan lifecycle. Grouped RTN and MSE clipping plans are ready immediately and have calibrationBatches=0. Supplying calibration data for these configurations is an error. Static activations and importance fitting require calibration. Omitting --output selects all graph outputs, which must be F32 for this measurement. --out is a new JSON report file; the command creates no model packages. The same quantization and float-retention flags apply, and --cpu-threads can limit worker threads.

The Python API accepts finite iterables of NumPy batches directly:

import json
from pathlib import Path
import numpy as np
import volvoxai as vx

def batches(directory):
    for path in sorted(Path(directory).glob("*.npz")):
        with np.load(path, allow_pickle=False) as batch:
            yield {name: batch[name] for name in batch.files}

static = vx.pb.StaticActivations(dtype=vx.pb.DataType.DATA_TYPE_I8)
report = vx.ptq_sensitivity("model", calibration_data=batches("calibration"),
    group_size=0, static_activations=static, evaluation_data=batches("evaluation"))
Path("sensitivity.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
print(report["baseline"], report["bestToKeepFloat"])

# Review the measured rows and choose the nodes your accuracy budget needs.
chosen = report["bestToKeepFloat"][:1]
result = vx.quantize("model", calibration_data=batches("calibration"),
    group_size=0, static_activations=static, output="mixed-precision",
    rules=[vx.pb.QuantizationRule(node_id=node, keep_float=vx.pb.Empty()) for node in chosen])

Choose node IDs from the measured report and express them as keep_float rules (or protobuf JSON rules with --rules). Evaluate the resulting package on your application's task metric before deploying it.

The sweep runs on CPU through the full C library. It copies each finite calibration/evaluation iterable once and reuses those batches for every candidate, so keep the sets small enough to retain in memory. The cost grows with quantizable node count and selected modes. baseline.sqnrDb is the worst selected output over all evaluation batches; rows include isolated, leaveOneFloat and, when measurable, recoveryDb. mostSensitive ranks isolated error, and bestToKeepFloat ranks recovered SQNR. Unsupported selections carry a note; a null SQNR can mean zero signal, zero error or a nonfinite measurement and is not a quality guarantee. The JSON report format is volvoxai-ptq-sensitivity/v1, shared with the JavaScript tool.

Build and publish

From the repository root:

make build_wheel

This builds volvoxai-wheel-build:<version> from python/Dockerfile.wheel, using package.json.version. Its pinned manylinux base supplies the Linux compatibility baseline; native CPU/CUDA/Vulkan/OpenGL inference and full libraries are compiled from the current checkout. The source checkout is mounted read-only. Temporary native outputs do not replace the existing native/ or npm release artifacts. WHEEL_BUILD_JOBS defaults to two. Native build intermediates are cached under build/python-wheel/, invalidated when the builder image changes. Every qualification run creates fresh Python environments. Failed candidates remain in that ignored build directory for diagnosis and are not published into dist/python/.

The build checks native profile boundaries and CUDA PTX, repairs the wheel with auditwheel, checks PyPI metadata with Twine, and installs it into clean CPython 3.10–3.14 environments. It exercises generated API coverage, sync/async calls, CPU inference, tokenization, planning, scheduling, training/rollback, ONNX conversion/PTQ and automatic exceptions across every sync/async service method. Session qualification includes model/library discovery, repeated inference, dynamic multi-input models, output selection, dtype preservation and cleanup on errors. Both FP32 exports and native PTQ packages run through the NumPy API. Workflow qualification also covers immutable metadata, host-buffer transport, async cancellation/deadline cleanup, streaming PTQ, an independent SGD oracle, AdamW checkpoint continuation, commit/rollback and fresh-file publication on every supported Python version. Tensor tests cover DLPack ownership, legacy/versioned capsules, CPU sharing, dynamic shapes and late external releases. For source-tree DLPack tests, first run python python/build_dlpack_bridge.py with setuptools and Python development headers installed.

wheel-validation.json records the CPU scope and test files for that exact wheel; a rebuild does not inherit GPU results from an older artifact. Physical GPU execution is qualified separately on a GPU host using VOLVOXAI_WHEEL_TEST_BACKEND=cuda, vulkan or opengl with python/tests/test_installed_wheel.py.

For CUDA tensor interoperability, install a compatible PyTorch distribution on the GPU host. To require all three Linux GPU backends, run:

VOLVOXAI_TEST_CUDA=1 VOLVOXAI_TEST_NATIVE_BUFFERS=cuda,vulkan,opengl \
  python -I python/tests/test_tensor_interop.py

This additionally checks nondefault streams, exact shared pointers, results that survive session close, cross-session reuse, INT32 values, a native INT8 PTQ package, interpreter shutdown and CUDA profiler transfer events. VOLVOXAI_TENSOR_TEST_EVIDENCE=path.json saves the profiler's copy event names and framework versions. The copy test uses a small warmed float32 model; it is not a receipt-model throughput benchmark.

The profiling tests import PyTorch before creating CUDA contexts. With the qualified PyTorch 2.5.1/CUDA 12.1 and driver 535.309.01 combination, calling cuInit before importing PyTorch can make CUPTI crash at process exit. A driver-only reproducer confirms this without loading VolvoxAI. The import order is a profiling-environment requirement; ordinary engine execution does not depend on PyTorch or CUPTI.

Successful builds publish these local files:

dist/python/<version>/volvoxai-<version>-cp310-abi3-manylinux_2_28_x86_64.whl
dist/python/<version>/SHA256SUMS
dist/python/<version>/wheel-validation.json

The engine uses its generated C ABI. The DLPack capsule adapter uses CPython's stable ABI with a Python 3.10 minimum; it has no CUDA/framework link dependency. The package is not a pure-Python any wheel. The version comes from package.json. Build provenance records the source digest and whether the checkout was dirty.

To build, validate and upload to PyPI in one command:

make publish_pypi

This runs build_wheel, then uses the same Docker image's pinned Twine to check the wheel metadata and upload the current package.json version. A build or validation failure stops the upload. No host Twine installation is needed.

To check the existing wheel without rebuilding or uploading:

make publish_pypi DRY_RUN=1

This verifies SHA256SUMS and runs twine check --strict in the wheel image, without reading registry credentials. make build_wheel also runs this dry run as its final step, so a packaging failure fails the build. The dry run does not check PyPI upload permissions or whether the version is already published.

If ~/.pypirc exists, the target mounts it read-only in the upload container. Twine automatically uses the token saved in its [pypi] section. Set TWINE_REPOSITORY to select another configured index, such as TestPyPI.

Environment credentials (TWINE_USERNAME and TWINE_PASSWORD) override the file. Without saved credentials, Twine prompts in a terminal. Set TWINE_NON_INTERACTIVE=1 for automation. The target also forwards TWINE_REPOSITORY_URL when using another package index.

make build_wheel only creates and validates local artifacts; make publish_pypi also uploads them. Attach the same wheel and checksum to GitHub Releases. JS/WASM retain their existing dist/<version>/ names and npm packaging.