The tensor library#

Train anywhere. Get the same weights.

vision.tensor has tensors on the GPU, reverse-mode autodiff, transformer layers, optimisers, int8 inference and tools to audit a training run. It trains a GPT from scratch.

import numpy as np
import vision.tensor as vt
from vision.tensor.nn import GPT
from vision.tensor.optim import AdamW

dev = vt.device("intel")                     # "nvidia", "cpu" (llvmpipe), "numpy", "numpy64"
model = GPT(dev, vocab=65, block_size=128, dim=192, layers=4, heads=6)
opt = AdamW(model.parameters(), lr=2e-3)

loss = model.loss(x, y)                      # x, y: numpy integer arrays [batch, 128]
loss.backward()
opt.step()
print(vt.fingerprint(model.parameters()))    # the same signature on any GPU

Tensors and autodiff#

x = vt.tensor(dev, np.ones((2, 3), np.float32), requires_grad=True)
y = (x @ x.T).sum()
y.backward()
print(x.grad.numpy())

Tensors support +, -, *, /, @, reshape, permute, transpose, .T, sum and mean, plus the functions softmax, cross_entropy, rms_norm, gelu, embedding and matmul. Data goes in with vt.tensor(dev, array) and comes out with t.numpy().

Every gradient is computed in a fixed order, without atomics, so training is reproducible bit by bit.

Modes#

Mode

Guarantee

Runs on

Mini-GPT training (RTX 4050)

auditable (default)

bit-identical on every device, including the CPU and numpy

everything

69 k tokens/s

gpu-exact

bit-identical between GPUs that fuse multiply-add (tested: NVIDIA = Intel)

GPUs; refuses devices without fused FMA, such as llvmpipe

98 k tokens/s

fast

deterministic only on the same kind of GPU

everything

104 k tokens/s

 
dev = vt.device("nvidia", mode="gpu-exact")

If in doubt, use the default. gpu-exact refuses a device that cannot give the same bits, instead of silently giving different ones.

Number formats#

Every matrix product has a number format, chosen per call, per layer or per device:

y = vt.matmul(x, w, format="mx16")
model = GPT(dev, vocab=65, block_size=128, format="mx8")   # every Linear layer
dev.format = "mx16"                                         # the default for 2-D products on dev

The MX formats quantise both operands and compute the product exactly, rounding once at the end, so the result is the same bits on every GPU and on the numpy device, whatever hardware computes it.

Format

Mantissa

Max error (normal data, relative to the largest output)

Product, RTX 4050

Mini-GPT, 600 steps

float32 (auditable)

24 bits

5.5e-7

1.2 TFLOPS

val loss 1.7856, 69 k tokens/s

mx16

14 bits

1.0e-4

2.2–2.8 TOPS

val loss 1.7855, 84 k tokens/s

mx8

7 bits

1.2e-2

4.2–6.1 TOPS

val loss 1.7920, 113 k tokens/s

mx16 trains like float32 here; mx8 is faster and slightly less accurate. mx16r is a variant of mx16, a little less accurate on uneven rows (Qwen2.5: max logit error 0.32 against 0.09, the same text). On GPUs with tensor cores that Vision can use (e.g. the RTX 4050), mx16r and mx8 run on them, with the same bits as on any other device: a prompt of 512 tokens goes from 3028 to 5122 tokens/s. Training the mini-GPT for 300 steps in mx16r gives the same validation loss as mx16 (2.0999) at 110 k tokens/s against 93 k, and the weights trained on the RTX 4050’s tensor cores have the same SHA-256 as those trained on the Intel UHD. The MX formats take 2-D operands (the Linear layers); batched products, such as those of attention, stay in float32 when the format comes from the device, and raise an error when it is asked for explicitly. With mx16, K is at most 16 320 per product.

Language models#

vision.tensor.llm runs pretrained models of the Llama family (Qwen2 / Qwen2.5, Llama, Mistral, SmolLM2; Decoder.load(folder, dev)) from their .safetensors files or from a .gguf file of llama.cpp (read with numpy only; bf16 becomes float32 exactly), with rotary embedding, grouped-query attention, a key/value cache and greedy decoding. In the auditable mode the logits have the same bits on every device, so the generated text is the same.

from vision.tensor.llm import Qwen2
model = Qwen2.load("examples/data/qwen2.5-0.5b", vt.device("intel"), format="mx16")
tokens = model.generate(ids, 64)

On Qwen2.5-0.5B-Instruct (494 M parameters):

Result

Logits against PyTorch in float64

max error 6.3e-5 in float32 (PyTorch float32: 4.0e-5); the same next token at every position

The same logits (SHA-256)

float32 on NVIDIA, Intel and llvmpipe; mx16 on NVIDIA, Intel and Apple (Metal); the same text everywhere, with and without the cache

SmolLM2-360M (Llama architecture)

float32 within 1.0e-4 of PyTorch float64, the same next token everywhere, the same logits on NVIDIA and Intel

Generation, RTX 4050

about 105 tokens/s in mx16 and 155 in mx8 (the laptop GPU’s clock varies by ±15%)

Generation, Intel UHD

41 tokens/s in mx16, 58 in mx8

Constant weights are quantised once (vt.quantized(w, "mx16")). The generation is optimised for speed without changing a bit of the logits: they are the same with or without the cache, and the same as the model computed one product at a time.

GGUF files are read with numpy only (vision.tensor.gguf): F32, F16, BF16, Q8_0, Q4_0, Q4_1, Q5_0, Q5_1, Q4_K, Q5_K and Q6_K are dequantised to float32 with the same bits as llama.cpp’s own gguf-py (checked for all but Q5_K), the Q/K permutation of Llama conversions is undone, and the configuration comes from the metadata. The official Q8_0 and Q4_K_M files of Qwen2.5-0.5B and the Q8_0 of SmolLM2-360M give the same logits on NVIDIA and Intel and the same text as their safetensors.

generate_batch(prompts, steps) serves several sequences at once: every step their new tokens go through the linear layers together, each sequence with its own cache, and in the MX formats every sequence gets exactly the tokens it gets alone. On the RTX 4050, 16 sequences of Qwen2.5-0.5B give 535 tokens/s in total in mx8 (133 for one) and 407 in mx16 (109 for one); on the Intel UHD, 93 in mx8 (55 for one).

Against llama.cpp (build b11887, llama-bench -p 512 -n 128), the same Qwen2.5-0.5B Q8_0 file, tokens/s:

RTX 4050: prompt 512

RTX 4050: generation 128

Intel UHD: prompt 512

Intel UHD: generation 128

llama.cpp CUDA

18 795

266

—

—

llama.cpp Vulkan

13 277

251

1 211

56.2

Vision mx8

7 145

139

885

56.4

Vision mx16r / mx16

5 992

111

388

40.0

Vision is at about half of llama.cpp on the NVIDIA GPU and level with it on the Intel one for generation; it is the one of the three whose logits and text are the same on every GPU.

Besides greedy decoding, generate(..., temperature=0.8, top_p=0.9, seed=1) draws the tokens from a seed, also the same on every machine. On Qwen2.5 the same seed gives the same poem on the RTX 4050 and on the Intel UHD, with or without the cache.

The LLM example answers a prompt and writes a receipt binding the model (SHA-256 of its files), the input, the format and the code to the tokens and to the signature of the logits; --verify runs it again on another device:

python examples/llm.py "What is the capital of France?" --device nvidia --receipt answer.json
python examples/llm.py --verify answer.json --device intel      # MATCH: the same tokens and logits

int8 inference#

The int8 matrix product is exact, so it gives the same result on every vendor. It runs at 3.3–4.7 TOPS on the RTX 4050.

from vision.tensor.quant import quantize_model
quantize_model(model)            # replaces every Linear with an int8 QuantLinear, in place

On the mini-GPT, validation loss is 1.5216 in int8 against 1.5209 in float32, with the same logits signature on NVIDIA, Intel and llvmpipe.

Auditing a training run#

  • Signature on the GPU. SHA-256 runs on the GPU; only the digests go back to the host. GPU and CPU (hashlib) give the same hash.

  • Chained log. Every step records the hash of the batch, the hash of the weights and the loss (exact bits). Every record carries the hash of the previous one.

  • The batch of step t is a pure function of (seed, t), so any step can be redone without running the previous ones, on any number of GPUs.

  • Checkpoints with the hash of the full state, and an auditor that redoes chosen or random stretches of the run on another device and compares bit by bit.

On the mini-GPT, the log with a signature per step costs 11%.

Heterogeneous training#

DataParallel splits micro-batches between devices of different vendors. Neither the number nor the vendor of the GPUs changes the weights: NVIDIA and Intel training together give the same SHA-256 as NVIDIA alone. cross_check recomputes one micro-batch on two devices and compares, which detects a faulty GPU or silent data corruption.

Performance#

Same RTX 4050, same model, batch 32 × 128 tokens:

Training tokens/s

FP32 matmul

Vision, auditable

69 k

0.7–1.4 TFLOPS

Vision, gpu-exact

98 k

1.8–2.1 TFLOPS

Vision, fast

104 k

0.9–2.8 TFLOPS

PyTorch + CUDA (FP32, eager)

185 k

cuBLAS 5.7–6.0 TFLOPS

PyTorch + CUDA (TF32)

262 k

cuBLAS 6.3–9.9 TFLOPS

PyTorch FP32 is still 1.9× faster than gpu-exact. Tensor cores, fused attention and mixed precision are not implemented yet.