MuNet

A greenfield, experimental C++17 graph compiler for Vulkan training and inference, with a small PyTorch-style Python interface.

MuNet is the building blocks for creating neural networks: tensors, layers, automatic differentiation, optimizers, graph compilation, device execution and model interchange. The installed library has NumPy as its only Python runtime dependency. Model architectures, datasets and training recipes belong in examples or applications.

Version 0.3 adds the operations needed to build and train the RT-DETR example. That model is an acceptance test and reference application under examples/rtdetr/; it is not packaged in the library. Full COCO convergence and physical-device performance are not yet established. The runtime retains the munet-nn PyPI project and munet_nn import name; import munet is also supported. This is a breaking rewrite of the 0.1 API.

The C++ core owns graph validation, shape inference, symbolic reverse-mode differentiation, fusion, temporary-buffer planning, CPU reference execution, GLSL generation, and Vulkan execution. Python constructs models, captures a static graph once, invokes the shader compiler on cache misses, and exposes model conversion and serialization.

What runs today

• Float32 tensors with positive static dimensions, rank up to eight.
• Convolution and pooling with backward, batched matrix multiplication, normalization, attention, differentiable bilinear sampling, indexing, TopK and broadcast arithmetic.
• PyTorch-style modules, persistent train/eval buffers, SGD, AdamW parameter groups, gradient clipping and EMA.
• A compiled forward/loss/backward/optimizer/EMA step and native assignment operation. Parameters remain in Vulkan memory between calls; .item(), .numpy(), saving, and explicit interop perform readbacks.
• Single-consumer scalar-operation fusion, including view expressions feeding matrix multiplication; temporary buffers reused after their final consumers.
• GLSL → SPIR-V compilation on demand, a persistent SPIR-V cache, Vulkan pipeline creation, and recorded command-buffer replay.
• Data-only .mnet save/load, including training state; native graph → ONNX → native graph conversion for the supported subset.
• Import of supported eval-mode PyTorch modules, including the pinned RT-DETR reference, through the modern ONNX exporter. Matching native layers accept PyTorch state dictionaries.
• An explicit CPU reference backend. Selecting Vulkan never silently runs native operators on the CPU.
• A disconnect-tolerant training swarm: one durable owner, C++ worker binaries, capability-based leases, cached offline computation, result retries, and sample-weighted MSE/SGD rounds.

Read the RT-DETR example, installation and releases, the architecture and decisions, training swarm guide, RT-DETR acceptance plan, and validation report.

Install

After 0.3.1 is published, the existing PyPI project installs the library and both swarm commands:

python -m pip install --upgrade 'munet-nn>=0.3.1'
munet-node --version
munet-server --version

Release workflows build Linux x86-64/aarch64 wheels, standalone node/server archives, and a CMake SDK. Preview artifacts are available from the PR's Wheels run. See installation and release details.

Build and run

For local development and testing, use the Makefile (Linux, Python 3.10+):

sudo apt-get install build-essential python3-dev python3-venv libvulkan-dev glslang-tools vulkan-validationlayers libcurl4-openssl-dev libssl-dev
make setup
make test
make test-vulkan

Use make setup VULKAN=0 and make test VULKAN=0 for a CPU-only build. make build rebuilds native code after edits; make smoke runs a short training example. Commands use .venv and the source tree automatically, without shell activation. make install also installs the library and node/server commands into .venv for use outside the checkout. See make help and the local development instructions for prerequisites, targeted tests and configuration overrides.

On Ubuntu with a working Vulkan driver:

sudo apt-get install build-essential libvulkan-dev glslang-tools libcurl4-openssl-dev libssl-dev
python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install '.[interop,test]'
python examples/train_mlp.py --device vulkan --steps 200

pip builds the C++ extension using CMake and pybind11. The native runtime loads the host Vulkan loader on first GPU use; ONNX and PyTorch are conversion/development dependencies, not execution engines. Install the GPU vendor's Vulkan driver separately. Linux software testing can use mesa-vulkan-drivers.

For an explicit CPU-only build:

python -m pip install '.[interop,test]' -Ccmake.define.MUNET_VULKAN=OFF
python examples/train_mlp.py --device cpu

For manual source development without rebuilding a wheel:

python -m pip install cmake pybind11 numpy pytest onnx onnxruntime scipy torch onnxscript
python tools/build.py                  # add --cpu-only when needed
PYTHONPATH=python python -m pytest -q
MUNET_TEST_VULKAN=1 PYTHONPATH=python python -m pytest -q
# With vulkan-validationlayers installed:
PYTHONPATH=python python tools/test.py --vulkan-validation

The opted-in Vulkan tests fail when the driver/compiler is unavailable. They do not skip or select the CPU instead. tools/build.py --cmake-arg=-DVulkan_INCLUDE_DIR=... accepts custom SDK paths.

RT-DETR example

The model, matching/loss recipe, image and COCO handling, training loop, and acceptance tests live together under examples/rtdetr/. Set it up separately:

make setup-rtdetr
make test-rtdetr
PYTHONPATH=python .venv/bin/python -m examples.rtdetr.train --help

Use from examples.rtdetr import RTDETR from the repository root. There is no munet.models or munet_nn.models API. See the example guide for training, checkpoint/resume, inference and evaluation commands. make setup builds the library without downloading PyTorch or the RT-DETR reference; make test installs the optional numerical-reference tools and runs library tests.

Python training

import numpy as np
import munet as mu

rng = np.random.default_rng(7)
model = mu.nn.Sequential(
    mu.nn.Linear(4, 16, rng=rng),
    mu.nn.ReLU(),
    mu.nn.Linear(16, 2, rng=rng),
)
optimizer = mu.optim.SGD(model.parameters(), lr=0.03)

@mu.compile(device="vulkan:0")
def train_step(x, target):
    optimizer.zero_grad()
    loss = mu.nn.functional.mse_loss(model(x), target)
    loss.backward()
    optimizer.step()
    return loss

x = rng.normal(size=(32, 4)).astype(np.float32)
y = rng.normal(size=(32, 2)).astype(np.float32)
for step in range(100):
    loss = train_step(x, y)
    if step % 20 == 0:
        print(loss.item())             # Explicit scalar readback.
print(train_step.stats())

The first call captures Python control flow and compiles the whole step. Subsequent calls replay the C++ execution plan. Python side effects inside the function happen only while tracing. Tensor-dependent Python branches are rejected. Input shapes and float32 dtype are guarded. Model structure and closure constants are static: create a new compiled function after changing them. Train/eval changes recapture automatically; AdamW group settings remain live between replays.

This initial API is tracing-only; it does not yet provide general eager tensor execution. Each compiled object has one specialization and one in-flight execution. A returned result borrows output storage: call .numpy() before the next invocation if you need to retain a copy. Stale results raise an error. Compiled calls are serialized; concurrent mutation of parameters shared across compiled objects is unsupported.

Conversion

from munet.interop import to_onnx, from_onnx, from_torch

inference = mu.compile(model, device="vulkan")
prediction = inference(x).numpy()
mu.save(inference, "model.mnet")

restored = mu.load("model.mnet", device="vulkan")
to_onnx(restored, "model.onnx")
imported = from_onnx("model.onnx", device="vulkan")
mu.save(imported, "roundtrip.mnet")

# Optional: pip install '.[torch]'
# imported = from_torch(torch_model.eval(), (example_torch_tensor,), device="vulkan")

ONNX support is default-domain opsets 13–21, float32 computation (integer tensors such as indices and quantized values are carried as exact integer-valued floats; integer/bool model inputs are converted on the way in, up to 2^24, and integer outputs are returned as numpy arrays of their declared type; from_onnx(..., precision="fp32") selects the exact path), and the operators listed in munet.interop.SUPPORTED_ONNX; docs/op_coverage.md is the machine-checked coverage table (all 191 operators of opsets 13-21; string/bytes operators and data-dependent output sizes from constant inputs, sequences, optionals and control flow resolved at import, recurrent layers unrolled). NonMaxSuppression imports when its result is a model output: it runs on the device with a static bound (max_output_boxes_per_class rows per batch and class) and the call returns that output as an int64 [num_selected, 3] numpy array with the padding trimmed; graphs that consume the result inside the model are rejected (use nn.functional.nms/batched_nms for fixed-size results). Before import, from_onnx merges layers at the weight level (optimize=True; program.optimizations lists the rewrites; all exact up to float32 rounding): BatchNorm folds into adjacent convolutions, sums of parallel convolutions on one input (plus an identity or BatchNorm(identity) term) re-parameterize into one convolution, chains of constant linear layers without a nonlinearity collapse when that does not add multiplies, and sibling MatMul/Gemm/Conv layers reading the same input become one wider layer (with their bias adds) whose outputs are slices. Measured (mixed, RTX 3060): a RepVGG-style network in training form 4.28 -> 3.02 ms, ResNet-50 exported with unfolded BatchNorm 9.64 -> 9.42 ms; exports that already fold BatchNorm and re-parameterize (current torch.onnx, deploy-mode RT-DETR) are unchanged within noise. QDQ and QLinear quantized graphs import and run on int8 kernels (docs/quantization.md). Inputs with symbolic dimensions return a ShapeSpecializedProgram that compiles one static specialization per concrete input shape (max_specializations=8 by default, least recently used evicted); Shape/Size fold to each specialization's shape. Local functions are inlined. Unsupported nodes/modes, external tensor data, sparse tensors and ONNX training metadata raise explicit errors.

Export produces a standard inference graph and preserves supported computation, not original Python classes or exact graph topology. Native training programs must be exported by compiling their trained model separately. A supported inference graph does not reconstruct the original training behavior of an arbitrary PyTorch model.

mu.save(train_step, "step.mnet") and mu.load(...) preserve compiled computation and current device state. munet.checkpoint.save_state/load_state supplies data-only state archives for application training loops; the RT-DETR example's DetectorTrainer uses it for model, optimizer, EMA, scheduler and RNG resume. This is a new versioned archive format; no reader for the previous MuNet format is included.

Rust and C

Each release ships the same runtime three ways, for Linux x86_64/aarch64 (macOS arm64 and Windows x86_64 are being brought up):

• Python: pip install munet-nn (wheels; other platforms build from the source distribution).
• Rust: cargo add munet-nn (imported as munet). The crate builds the C++ core with CMake from the bundled sources (C++17 compiler and CMake >= 3.20; Vulkan headers are vendored), or links a prebuilt C API from MUNET_LIB_DIR. See rust/README.md.
• C: munet-c-<version>-<os>-<arch> archives on the GitHub release (header munet/munet_c.h, shared library and static archives).

Models are exported once from Python and loaded natively:

program = from_onnx("detector.onnx", input_shapes={"images": (1, 3, 640, 640)})   # static shapes
mu.save(program, "detector.mnet")   # ONNX imports record input/output names and declared dtypes

let mut program = munet::Program::load("detector.mnet", &munet::Options::default())?;
program.run(&[&image, &sizes])?;            // f32 slices in program.inputs() order (integers as exact floats)
let scores = program.read_by_name("scores")?;

Running on a GPU needs a Vulkan driver and glslangValidator (kernels are compiled once and cached in MUNET_CACHE_DIR, shared with the Python package, as are tile configurations autotuned from Python with autotune=True). Options { external_memory: true, .. } exports the program's device memory (opaque fd) so another Vulkan or CUDA context can write inputs and read outputs in place.

Training swarm

Prepare a job from a native model and finite float32 arrays:

from munet.swarm import create_job

create_job(model, x, y, "training-job", global_batch_size=32,
           micro_batch_size=4, epochs=10, lr=0.03, include_vulkan=True)

Build and launch native workers:

sudo apt-get install libcurl4-openssl-dev libssl-dev
cmake -S . -B build-node -DMUNET_PYTHON=OFF -DMUNET_SWARM_NODE=ON -DCMAKE_BUILD_TYPE=Release
cmake --build build-node -j

# Supply the same MUNET_SWARM_TOKEN (enrolment secret, at least 32 characters) to each process;
# each node exchanges it for its own revocable node token on registration.
python -m munet.swarm training-job --port 8765
# In another terminal:
./build-node/munet-node --owner http://127.0.0.1:8765 --state node-a --device vulkan:0

The owner assigns chunks from the same checkpoint, combines gradients by sample count, and applies one SGD update per global batch. Nodes can disappear, reconnect, or restart using their state directories. Lost work is leased again; duplicate results never apply twice. Each node needs a full model replica. The current loss is MSE, and models must be deterministic and independent across samples.

The worker needs no Python or shader compiler. SPIR-V is prepared on the owner, and the driver compiles pipelines on each device. This Linux/POSIX prototype uses JSON transport and explicit gradient/checkpoint transfers; it has not established large-model or WAN performance. Longer offline local-SGD and fast multi-GPU collectives remain separate future modes. See the swarm guide for TLS, deployment, restart rules, constraints and the runnable MLP example.

Reliability, observability and scaling

• GPU loss. Vulkan failures raise mu.VulkanError (a RuntimeError carrying the VkResult); VK_ERROR_DEVICE_LOST raises mu.DeviceLostError, after which the program fails fast. mu.compile(..., recover="rebuild") instead rebuilds the executor once and retries, restarting parameters from their host mirror (the last .numpy() read).
• Threads. One compiled program may be called from many threads: native calls on a plan are serialized by a mutex and run without the GIL, and separate programs run in parallel.
• Logging and metrics. Modules log to munet.* loggers (the library installs only a NullHandler). mu.metrics is a stdlib registry (counters, gauges, histograms, mu.metrics.span()) with a Prometheus text exporter mu.metrics.prometheus(); the swarm owner serves it at /metrics.
• Several GPUs in one machine. mu.DataParallel(loss_fn, params, optimizer, ["vulkan:0", "vulkan:1"]) splits each batch on dim 0, computes gradients on every device in parallel, averages them on the host (weighted by shard size, deterministic for a fixed device list) and applies one optimizer step; dp.step(x, y) returns the batch loss.
• Memory safety. make test-asan / make test-tsan run the tests against sanitizer builds (-DMUNET_SANITIZE=...), and make fuzz-smoke runs the libFuzzer/Atheris targets in fuzz/ (swarm program JSON, ONNX import, checkpoints).

Runtime limits

Kernels are float32 except matmul. Matmuls of any rank (batch dimensions are folded into rows when the right operand is unbatched, otherwise computed per batch with broadcasting) use a GEMV kernel for M<=32, a tiled register kernel (16x16 threads of 4x4 outputs, register-prefetched K slices) for precision="fp32", and by default an fp16-input/fp32-accumulate cooperative-matrix kernel (see precision below). On an RTX 3060 the mixed kernel reaches about 19-24 TFLOP/s from 1024^3 to 4096^3 (GPU timestamps, resident weights); convolutions use implicit-GEMM kernels on the same path. See docs/perf for measured tables.

One device-local buffer holds the entire planned arena, with a host staging mirror. Graph/weight metadata also consumes host memory. arena_bytes reports only the planned device arena, not total process memory. Tensor-sized descriptors allow the total arena to exceed maxStorageBufferRange; individual tensors must fit it. Two-dimensional dispatches respect device workgroup limits. Host input arrays are copied straight from the NumPy buffers into one of two compact input staging buffers (no intermediate copy), so the host can stage call N+1 while the GPU executes call N; each replay waits only for the one that used the same staging buffer. No CPU wait occurs between individual kernels. On a 128 MB input-bound workload this took a call from 31.0 ms to 15.1 ms (13.6 ms when calls are pipelined; host memory bandwidth limits further overlap). Independent graph branches are not yet spread over multiple queues.

Remaining work includes mixed precision beyond matmul (convolutions, attention), operator fusion, autotuning, general eager execution, asynchronous staging, fast multi-GPU collectives and physical platform validation. The detector has a bounded shape-specialization cache; switching cached programs synchronizes shared state through the host. The swarm still implements its narrow MSE/SGD contract, and does not yet train RT-DETR across workers. See the detector guide and coverage inventory.

Configuration

Setting	Meaning
device="vulkan:N"	Select a physical Vulkan device by index from mu.devices()
device="cpu"	Explicit C++ reference execution
fuse=False	Disable expression fusion for numerical/debug comparisons
precision="mixed" (default)	Matmuls (N,K>=16) and convolutions (GEMM dimensions >=16) take fp16 inputs (round-to-nearest, saturating at +-65504, NaN preserved) and accumulate in fp32 on cooperative matrices; convolutions run as implicit GEMMs, and their backward passes rescale the gradient operand by a power of two from max|dy| so tiny gradients do not underflow. Error is input rounding only: about 3e-4 relative (Frobenius) for one unit-scale matmul. It compounds through deep networks: on the RT-DETR R50 training test, losses stay within 6e-4 of fp32 and parameter gradients have cosine >=0.991 (median relative error 1.6%). fp16 copies of constants and parameters the step never updates stay resident and refresh after Parameter.assign; the fp32 originals of such weights are then kept only in host-visible memory read by the refresh (stats()["host_weight_bytes"]), so device memory holds each inference weight once, as fp16 (or int8/int4 for weight-only quantized layers, see docs/quantization.md). Needs VK_KHR_cooperative_matrix; otherwise fp32
nn.functional.scaled_dot_product_attention	softmax(QK^T*scale + additive mask)V over [B,H,T,D], causal and grouped-query heads; MultiheadAttention uses it when weights/dropout are not requested. Mixed Vulkan inference runs a fused flash-style kernel (head sizes multiples of 16 up to 128; S and P stay on chip); training graphs share the softmax probabilities with the backward; other plans use matmul + softmax ops
profile=True	Serialize kernels and record per-kernel GPU timestamps; read them with compiled.describe() (stats()["gpu_ms"] is always available)
precision="mixed_safe"	As mixed, plus per-tensor power-of-two rescaling of the gradient operands of fp16 matmuls before conversion (exact), so cotangents far below fp16's normal range (tested at 1e-7) keep their relative accuracy. Costs one reduction per gradient tensor (transformer block train step +10%), so it is opt-in; convolution backward is rescaled in every mixed mode
precision="fp32"	Full float32 matmuls and convolutions (exact opt-out; relative error about 1e-6). Convolutions run as fp32 implicit GEMMs on the same register-tiled kernel as matmuls; on an RTX 3060 this path is 1.0-1.5x PyTorch fp32 on GEMM/CNN workloads and up to 2.3x on transformer forwards (docs/perf/fp32.md)
MUNET_VULKAN_LIBRARY	Optional explicit host Vulkan loader path
MUNET_GLSLANG	Absolute path to glslangValidator
MUNET_CACHE_DIR	Persistent SPIR-V and driver pipeline cache directory (corrupt entries are rebuilt; a read-only directory disables caching). stats() splits cold start into shader_compile_ms and device_setup_ms
autotune=True / MUNET_AUTOTUNE=1	Time a bounded set of tile configurations for each GEMM/convolution shape on first use (GPU timestamps on single-op micro-programs), keep the fastest in a per-device/driver cache, and fall back to the heuristics without an entry or once MUNET_AUTOTUNE_BUDGET_S (default 30 s) is spent; stats() reports autotune_ms. Measured: RT-DETR 20.4 -> 18.7 ms, skinny/tall GEMMs 4-8% faster; first-time tuning of RT-DETR takes minutes, so it is opt-in
MUNET_DISABLE_COOPMAT=1	Force the portable kernels (no cooperative matrices) under precision="mixed", e.g. to test the fallback on a GPU that has them
MUNET_TEST_VULKAN=1	Require Vulkan execution in the test matrix
MUNET_SWARM_TOKEN	Enrolment secret: nodes present it once to /v1/register and receive a node-scoped token (revoke with munet-server JOB --revoke WORKER)
recover="rebuild"	On DeviceLostError (GPU crash/reset), rebuild the executor once and retry from the parameters' host mirror
MUNET_TEST_INJECT_DEVICE_LOST=N	Test hook: the N-th Vulkan queue submission reports VK_ERROR_DEVICE_LOST
MUNET_SWARM_NODE	Path to a built worker executable for integration tests
MUNET_SWARM_NODE_WRAPPER	Command prefix for the worker in integration tests (e.g. valgrind ...)
MUNET_TEST_SWARM=1	Require native swarm integration tests; fail if the worker binary is missing

The SPIR-V cache contains baseline Vulkan 1.1 modules keyed by generated source and shader compiler version. Device pipelines are built and retained per compiled object; their driver-specific cache is not persisted yet.