Training Large Models Without Python: A Two‑Year Review of the GoMLX Go ML Framework

Architecture Refactor and Independent Compute Engine

Since version v0.28 the low‑level graph definition and JIT compilation logic have been moved to a separate compute repository (https://github.com/gomlx/compute), establishing a clear "high‑level API + pluggable compute engine" separation.

Clear responsibilities : the compute repo defines and executes multi‑dimensional computation graphs; the main GoMLX repo provides ergonomic APIs for autodiff, layers, training loops, datasets and checkpoint management.

Backend migration : the original Go backend now lives in github.com/gomlx/compute/gobackend; the XLA backend resides in github.com/gomlx/go-xla/compute/xla.

Third‑party backend support : developers can implement a custom backend by satisfying the notimplemented interface and running backendtest.RunAll(t, myBackend) to pass the standard test suite.

Four Core Abstractions

Backend ( compute.Backend ) – bridges Go processes to hardware (CPU, GPU, TPU), handles JIT compilation and memory movement. A process typically creates one backend and reuses it globally. Automatic selection follows the order CUDA‑GPU → Metal → CPU; explicit selection is possible via the GOMLX_BACKEND environment variable or compute.NewWithConfig("go").

Graph – a pure‑Go representation of a computation using *graph.Node objects. Nodes represent future values (shape and dtype only) and cannot be read until the graph is executed with .Call(). Errors such as shape mismatches are raised as panics during graph construction and captured by the Exec wrapper, which converts them to standard Go errors on .Call(). This mirrors JAX @jax.jit and TensorFlow @tf.function semantics.

Tensor – the concrete multi‑dimensional array ( shapes.Shape + dtypes.DType) that exists in both host and device memory. Transfers are lazy; explicit FinalizeAll() releases device memory for large allocations.

Store – manages trainable parameters and scopes. model.Store holds the actual tensor data; model.Scope is a lightweight pointer to a hierarchical namespace inside the store, enabling layered variable naming such as scope.In("layer1").

Multi‑Backend Coverage

xla – based on OpenXLA/PJRT, JIT‑compiles to CPU, Nvidia GPU (and likely AMD/Intel), and TPU; supports only static shapes. Ideal for training large models where peak performance matters.

go – pure Go implementation without CGO, portable to WASM, includes SIMD‑accelerated matrix multiplication (AVX2/AVX512) and basic operator fusion. Suited for embedded devices, browser inference, or any environment where a single binary is required.

go‑darwinml (experimental) – CoreML bindings for macOS/iOS with Metal acceleration, targeting native Apple‑ecosystem inference.

Backend selection can be automatic or forced, e.g. export GOMLX_BACKEND=xla:cuda to use XLA with Nvidia CUDA.

Ecosystem Integration

go‑huggingface – native Go tokenizers, safetensors/GGUF parsing, and direct model download from HuggingFace.

onnx‑gomlx – converts ONNX models to GoMLX graphs, enabling inference and fine‑tuning without onnxruntime.

numpy interop – github.com/gomlx/gomlx/core/tensors/numpy reads NumPy arrays for data exchange.

Docker + JupyterLab – a pre‑built image bundles GoMLX, JupyterLab and GoNB, with optional CUDA support for quick experimentation.

Practical Example: Training an MLP in <60 Lines

import (
    "github.com/gomlx/compute"
    "github.com/gomlx/gomlx/backends/default"
)

// 1. Select backend (auto‑selects the best available)
backend, _ := compute.New()

// 2. Create a store for parameters
store := model.NewStore()

// 3. Build an in‑memory dataset of (x, y) → (r, g, b) mappings
inputs := make([][]float32, 0, width*height)
labels := make([][]float32, 0, width*height)
for y := 0; y < height; y++ {
    for x := 0; x < width; x++ {
        nx := float32(x)/float32(width)*2.0 - 1.0
        ny := float32(y)/float32(height)*2.0 - 1.0
        inputs = append(inputs, []float32{nx, ny})
        r, g, b, _ := img.At(x, y).RGBA()
        labels = append(labels, []float32{float32(r)/65535.0, float32(g)/65535.0, float32(b)/65535.0})
    }
}

ds, _ := dataset.InMemoryFromData(backend, "image_pixels", []any{inputs}, []any{labels})
ds.BatchSize(512, false).Shuffle().Infinite(true)

// 4. Define a three‑layer MLP model function
modelFn := func(scope *model.Scope, spec any, inputs []*Node) []*Node {
    x := inputs[0]
    h := denseLayer(scope.In("layer1"), x, 64)
    h = activation.Relu(h)
    h = denseLayer(scope.In("layer2"), h, 64)
    h = activation.Relu(h)
    h = denseLayer(scope.In("layer3"), h, 64)
    h = activation.Relu(h)
    y := Sigmoid(denseLayer(scope.In("output"), h, 3))
    return []*Node{y}
}

// 5. Configure trainer (Adam optimizer, MSE loss)
trainer := train.NewTrainer(
    backend, store, modelFn,
    loss.MeanSquaredError,
    optimizer.Adam().LearningRate(0.003).Done(),
    nil, nil,
)

// 6. Run training loop for 5000 steps, logging every 1000 steps
loop := train.NewLoop(trainer)
train.EveryNSteps(loop, 1000, "log_metrics", 0, func(l *train.Loop, metrics []*tensors.Tensor) error {
    fmt.Printf("Step %5d: MSE Loss = %.6f
", l.LoopStep, metrics[0].Value())
    return nil
})
_, _ = loop.RunSteps(ds, 5000)

The example wires the four abstractions together: backend compiles and runs the graph, store holds parameters, graph (via modelFn) describes the forward pass, and trainer/loop orchestrates optimization. Each component can be swapped independently (e.g., change optimizer, switch backend via GOMLX_BACKEND, or modify the model architecture).

Latest Features

Built‑in KAN (Kolmogorov‑Arnold Networks) and VNN (Vector Neural Networks) layers, including variants such as GR‑KAN and SO(3)‑equivariant VNN.

Gradient checkpointing API that trades recomputation for reduced device memory, useful for large‑scale models.

Experimental distributed execution using XLA Shardy (GSPMD) for multi‑GPU/TPU training; the feature is marked as actively improving.

Command‑line tool gomlx_checkpoints visualizes loss curves and compares multiple runs using Plotly.

Comparison with the Python Ecosystem

Performance: XLA‑based backends achieve comparable speeds to JAX, TensorFlow and PyTorch/XLA because they share the same compilation engine.

Portability: The pure Go backend can be compiled to WebAssembly, enabling inference in browsers or any platform that supports Go without requiring a Python runtime.

Trade‑offs: The Go ecosystem is smaller, APIs are more verbose, and the model zoo is limited compared to Python. However, GoMLX offers a single‑binary deployment model and transparent, “no‑magic” abstractions that align with Go’s design philosophy.

Roadmap and Current State

GitHub stars >1.5k, forks >70, >3500 commits; latest release v0.27.3.

Dedicated documentation site (https://gomlx.github.io/docs/overview/) covering tutorials, core concepts, training, layers and debugging.

Active community on Slack (#gomlx), Google Groups and GitHub Discussions.

Long‑term goals: make Go a first‑class citizen for model training, support large‑scale distributed training, and enable compilation of models to C libraries or WebAssembly for consumption from any language.

Key Takeaways

GoMLX has evolved from a proof‑of‑concept into a full‑stack ML platform with a layered architecture, multiple backends, and an expanding ecosystem. While it still trails Python in terms of model zoo size and community resources, it provides a compelling option for teams that need Go‑native training and inference, especially when single‑binary deployment or WASM compatibility is required.