Metadata-Version: 2.2
Name: munet-nn
Version: 0.4.0
Summary: Experimental C++ graph compiler for portable Vulkan training and inference
Author: MuNet Contributors
License: Proprietary
Project-URL: Homepage, https://github.com/bryan-pakulski/MuNet
Project-URL: Repository, https://github.com/bryan-pakulski/MuNet
Requires-Python: >=3.10
Requires-Dist: numpy>=1.24
Provides-Extra: interop
Requires-Dist: onnx>=1.16; extra == "interop"
Provides-Extra: vk
Requires-Dist: onnx>=1.16; extra == "vk"
Requires-Dist: onnxruntime>=1.18; extra == "vk"
Provides-Extra: test
Requires-Dist: pytest>=8; extra == "test"
Requires-Dist: onnx>=1.16; extra == "test"
Requires-Dist: onnxruntime>=1.18; extra == "test"
Requires-Dist: scipy>=1.10; extra == "test"
Provides-Extra: torch
Requires-Dist: torch>=2.6; extra == "torch"
Requires-Dist: onnx>=1.16; extra == "torch"
Requires-Dist: onnxscript>=0.3; extra == "torch"
Description-Content-Type: text/markdown

# 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](examples/rtdetr/README.md). 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](examples/rtdetr/README.md), [installation and releases](docs/install.md), [the architecture and decisions](docs/architecture.md), [training swarm guide](docs/swarm.md), [RT-DETR acceptance plan](docs/rtdetr-acceptance.md), and [validation report](docs/validation.md).

## Install

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

```bash
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](docs/install.md).

## Build and run

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

```bash
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](docs/install.md#local-development-with-make)
for prerequisites, targeted tests and configuration overrides.

On Ubuntu with a working Vulkan driver:

```bash
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:

```bash
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:

```bash
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:

```bash
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](docs/rtdetr.md)
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

```python
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

```python
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](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](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](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:

```python
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
```

```rust
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:

```python
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:

```bash
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](docs/swarm.md) 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](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](docs/rtdetr.md) and [coverage inventory](docs/rtdetr-coverage.json).

## 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](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](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.
