README
¶
GoMLX -- Accelerated ML Libraries for Go
GoMLX is a fast and easy-to-use set of ML libraries built on top of XLA, a just-in-time compiler of numeric computations to CPU, GPU and TPU.
Quick Start: see our tutorial, or a guided example for Kaggle Dogs Vs Cats using GoMLX.
This under development, and should be considered experimental for now.
It was developed primarily as a platform to easily experiment with ML ideas, and to allow one to use Go for ML. But hopefully it can grow beyond that -- see Long-Term Goals below.
It was developed with a strong focus on being simple to read and reason about, leading the user to correct and transparent mental model of what is going on (no surprises). Documentation is kept up-to-date (if it is not well documented, it is as if the code is not there) and error messages are useful and try to make it easy to solve issues. Similar approach to the one described in Section 2 of Yggdrasil Decision Forests paper.
Long-term Goals
- Building and training models in Go -- as opposed to Python (or some other language) -- with focus on:
- Being simple to read and reason about, leading the user to a correct and transparent mental model of what is going on. Even if that means being more verbose when writing.
- Clean, separable APIs: individual APIs should be self-contained and decoupled where possible.
- Composability: Any component should be replaceable, so they can be customized and experimented. That means sometimes more coding (there is not one magic train object that does everything), but it makes it clear what is happening, and it's easy to replace parts with a third party versions or something custom.
- Up-to-date documentation: if the documentation is not there or if it's badly written, it's as if the code was not there either.
- Clear and actuable error reporting:
- To be a robust and reliable platform for production. Some sub-goals:
- Support modern accelerator hardware like TPUs and GPUs.
- Save models to industry tools like TensorFlow Serving.
- Import pre-trained models from Hugging Face Hub and TensorFlow Hub.
- Compile models to binary as in C-libraries and/or WebAssembly, to be linked and consumed (inference) anywhere (any language).
- To be a productive research and educational platform to learn and try new ML ideas:
- As composable and decoupled as possible, to allow anything to be tweaked and replaced without much hassle.
- Simple to read and well documented for anyone wanting to see how things are done.
- Support mirrored training on multiple devices and various forms of distributed training (model and/or data parallelism) in particular to support for large language models and similarly large model training.
What is already done
A working proof-of-concept, with many important components in place, from the bottom to the top of the stack. But only still a very narrow slice of what a major ML library/framework should provide (like TensorFlow, Jax or PyTorch).
- XLA integration for model training and evaluation -- including GPU (and presumably TPU, but never tested so likely not working). Unfortunately there is no lean inference-only version yet: so to run predictions one has to also import the whole engine and XLA machinery.
- Autograd: automatic differentiation -- only gradients for now, no Jacobians.
- Context: automatic variable management for ML models.
- Small ML layers library with some of the most popular machine learning "layers": dense (simple FFN layer),
activation functions, Layer Normalization, Batch Normalization, Convolutions, Pooling, Dropout, Multi-Head-Attention (for transformer layers), PiecewiseLinear (for calibration and normalization). - Training library, with some pretty-printing. Including plots for Jupyter notebook, using Bash Kernel.
- Also, various debugging tools: collecting values for particular nodes for plotting, simply logging the value of nodes during training, stack-trace of the code where nodes are created (TODO: automatic printing stack-trace when a first NaN appears during training).
- SGD and Adam optimizers.
- Various losses and metrics.
- Examples: Synthetic linear model; Adult/Census model; Cifar-10 demo; Dogs & Cats classifier demo; IMDB Movie Review demo.
Installation
For now Linux only. It does work well also in WSL (Windows Subsystem for Linux) in Windows.
And likely it would work in Macs --> contributions are very welcome, I don't have a Mac.
Linux packages dependencies.
The library depends on the following libraries:
libunwind8: usually available in most Linux systems.liblzma5: compression library, also usually available.- TC Malloc, usually packaged as
libgoogle-perftools4: fast malloc version, and memory debugging tools.
Typically, this can be installed with:
$ sudo apt-get install libunwind8 libgoogle-perftools4 liblzma5
gomlx_xla.tar.gz installation
This is the C-library with the bindings used to link GoMLX to XLA. The library is pretty large, ~500Mb (with GPU and TPU, it statically links most of what it needs) -- for Just-In-Time (JIT) compilation it includes the whole LLVM compiler.
The gomlx_xla.tar.gz package provided in the Release area includes a libgomlx_xla.so file and a few .h files
needed for the compilation of GoMLX. They are separated on two top level directories /lib and /include, and for
now the recommended way is to just untar them in /usr/local, which is usually in the default
path for inclusion and dynamic library loading. For instance:
$ cd /usr/local
$ tar xzvf .../path/to/gomlx_xla.tar.gz
This should be enough for most installations. If CGO is not finding the library, you may need to configure some environment variables.
After that, just import it as with any Go library.
Tutorial
See the tutorial here. It covers a bit of everything.
After that look at the demos in the examples/ directory.
The library itself is well documented (pls open issues if something is missing), and the code is not too hard to read (except the bindings to C/XLA, which were done very adhoc).
Finally, feel free to ask questions: time allowing (when not in work) I'm always happy to help. See TBD.
How is it implemented ?
At the lowest level it includes a Go dialect to construct computation graphs, and a C library integrating the TensorFlow's XLA just-in-time compiler -- soon to be migrated to OpenXLA.
In the future ideally it may also support other compilers, like IREE, Apache's TVM, and broaden the support for different accelerators.
To support it GoMLX implements autograd (automatic differentiation) and a minimal tensor library with "local" (on the host cpu) or "device" (on the accelerator device, e.g. a GPU) storage.
At a high level the library provides support for a Context object that hold model variables, a convenient trainer and a library of ML layers and tools: dense layers, convolutions, batch and layer normalization, transformer like attention, ADAM optimizer (and Adamw, Adamax), etc.
Collaborating
The project is looking forward contributions for anyone interested. Many parts are not yet set in stone, so there is plenty of space for improvements and re-designs for those interested and with good experience in Go, Machine Learning and APIs in general.
No governance guidelines have been established yet, this also needs work.
Advanced Topics
Documentation
¶
There is no documentation for this package.
Source Files
Directories
¶
| Path | Synopsis |
|---|---|
|
cmd
|
|
|
constraint_generator
command
|
|
|
gomlx_checkpoints
module
|
|
|
examples
|
|
|
adult/demo
command
Linear generates random synthetic data, based on some linear mode + noise.
|
Linear generates random synthetic data, based on some linear mode + noise. |
|
cifar
Package cifar provides a library of tools to download and manipulate Cifar-10 dataset.
|
Package cifar provides a library of tools to download and manipulate Cifar-10 dataset. |
|
cifar/demo
command
Demo for cifar library: it implements 2 models, a FNN and a CNN.
|
Demo for cifar library: it implements 2 models, a FNN and a CNN. |
|
dogsvscats/demo
command
demo for Dogs vs Cats library: you can run this program in 3 different ways:
|
demo for Dogs vs Cats library: you can run this program in 3 different ways: |
|
imdb
Package imdb contains code to download and prepare datasets with IMDB Dataset of 50k Movie Reviews.
|
Package imdb contains code to download and prepare datasets with IMDB Dataset of 50k Movie Reviews. |
|
imdb/demo
command
IMDB Movie Review library (imdb) demo: you can run this program in 4 different ways:
|
IMDB Movie Review library (imdb) demo: you can run this program in 4 different ways: |
|
linear
command
Linear generates random synthetic data, based on some linear mode + noise.
|
Linear generates random synthetic data, based on some linear mode + noise. |
|
notebook
Package notebook allows one to check if running within a notebook.
|
Package notebook allows one to check if running within a notebook. |
|
notebook/bashkernel
Package bashkernel implements tools to output rich content to a Jupyter notebook running the bash_kernel (https://github.com/takluyver/bash_kernel).
|
Package bashkernel implements tools to output rich content to a Jupyter notebook running the bash_kernel (https://github.com/takluyver/bash_kernel). |
|
notebook/gonb/margaid
Package margaid implements automatic plotting of all metrics registered in a trainer, using the Margaid library (https://github.com/erkkah/margaid/) to draw SVG, and GoNB (https://github.com/janpfeifer/gonb/) to display it in a Jupyter Notebook.
|
Package margaid implements automatic plotting of all metrics registered in a trainer, using the Margaid library (https://github.com/erkkah/margaid/) to draw SVG, and GoNB (https://github.com/janpfeifer/gonb/) to display it in a Jupyter Notebook. |
|
experimental
|
|
|
collector
Package collector implements DataCollector, which attaches itself to a computation graph executor and collects values of any selected computation graph node.
|
Package collector implements DataCollector, which attaches itself to a computation graph executor and collects values of any selected computation graph node. |
|
nest
Package nest implements the Nest generic container, used for generic inputs and outputs in GoMLX.
|
Package nest implements the Nest generic container, used for generic inputs and outputs in GoMLX. |
|
Package graph is the core package for GoMLX.
|
Package graph is the core package for GoMLX. |
|
graphtest
Package graphtest holds test utilities for packages that depend on the graph package.
|
Package graphtest holds test utilities for packages that depend on the graph package. |
|
ml
|
|
|
context
Package context defines the Context and Variable types: Context organizes variablesMap and variablesMap manages the storage of values typically used as variablesMap.
|
Package context defines the Context and Variable types: Context organizes variablesMap and variablesMap manages the storage of values typically used as variablesMap. |
|
context/checkpoints
Package checkpoints implements checkpoint management: saving and loading of checkpoints.
|
Package checkpoints implements checkpoint management: saving and loading of checkpoints. |
|
context/ctxtest
Package ctxtest holds test utilities for packages that depend on context package.
|
Package ctxtest holds test utilities for packages that depend on context package. |
|
context/initializers
Package initializers include several weight initializers, to be used with context.
|
Package initializers include several weight initializers, to be used with context. |
|
data
Package data is a collection of tools that facilitate data loading and preprocessing.
|
Package data is a collection of tools that facilitate data loading and preprocessing. |
|
layers
Package layers holds a collection of common modeling layers.
|
Package layers holds a collection of common modeling layers. |
|
train
Package train holds tools to help run a training loop.
|
Package train holds tools to help run a training loop. |
|
train/commandline
Package commandline contains convenience UI training tools for the command line.
|
Package commandline contains convenience UI training tools for the command line. |
|
train/losses
Package losses have several standard losses that implement train.LossFn interface.
|
Package losses have several standard losses that implement train.LossFn interface. |
|
train/metrics
Package metrics holds a library of metrics and defines
|
Package metrics holds a library of metrics and defines |
|
train/optimizers
Package optimizers implements a collection of ML optimizers, that can be used by train.Trainer, or by themselves.
|
Package optimizers implements a collection of ML optimizers, that can be used by train.Trainer, or by themselves. |
|
types
|
|
|
shapes
Package shapes defines Shape and DType and associated tools.
|
Package shapes defines Shape and DType and associated tools. |
|
slices
Package slices are
|
Package slices are |
|
tensor
Package tensor provides a `Tensor` interface with 2 different implementations: `Local` and `Device`.
|
Package tensor provides a `Tensor` interface with 2 different implementations: `Local` and `Device`. |
|
Package xla wraps the XLA functionality, plus extra dependencies.
|
Package xla wraps the XLA functionality, plus extra dependencies. |