A Machine Learning framework from scratch in pure Mojo 🔥
Basalt is a stand-alone machine learning framework that leverages the power of Mojo.
As discussed by Modular, Mojo is a language for the future of AI development. Built on top of MLIR technology, rather than existing GCC and LLVM approaches, Mojo looks and feels like Python code, yet performs much closer to languages like Rust or C++. Parametric functions and compile time parameters allow for the graph to statically compiled. Having the static graph allows for much harder performance optimizations.
Basalt, while still in its infancy, is able to achieve speeds comparable to well established frameworks like Pytorch. Below a snapshot of the current benchmarks. But keep posted, there is much more room for improvement and we are upgrading the project on a daily basis.
Try out the benchmarks yourself:
mojo -I . examples/housing.mojo
mojo -I . examples/sin_estimate.mojo
mojo -I . examples/mnist.mojo
Compare to the alternative PyTorch implementation:
Make sure to install the requirements in python-requirements.txt in your python environment.
python examples/housing.py
python examples/sin_estimate.py
python examples/mnist.py
- Improve matrix multiplication and convolution kernels
- Switch to custom Tensor and TensorShape implementations
- Improve benchmarks and overall model execution performance
- Add profiling and additional performance tests
- Add additional operators: Slice, (Un)Squeeze, Concat, Clip, Gather, Split, FMA ...
- Better layer support and more activation functions
- Graph submodules & graph concatenation
- Computer vision benchmark.
- Better parallelization
- GPU support
- Reworked Dataloader
- Autotuning and related features
- Graph compilation optimizations
- Operator fusion
- ONNX / Max compatibility
Basalt is built by community efforts and relies on your expertise and enthousiasm!
Small fixes and improvements are much appreciated. If you are considering larger contributions, feel free to contact us for a smoother communication channel on Discord. If you find a bug or have an idea for a feature, please use our issue tracker. Before creating a new issue, please:
- Check if the issue already exists. If an issue is already reported, you can contribute by commenting on the existing issue.
- If not, create a new issue and include all the necessary details to understand/recreate the problem or feature request.
- Fork the Project
- Create your Feature Branch
- Commit your Changes
- Push to the Branch
- Open a Pull Request
Once your changes are pushed, navigate to your fork on GitHub. And create a pull request against the original basalt-org/basalt repository.
- Before creating a PR make sure it doesn't break any of the unit-tests. (e.g.
mojo run -I . test/test_ops.mojo)- Introducing new big features requires a new test!
- In the pull request, provide a detailed description of the changes and why they're needed. Link any relevant issues.
- If there are any specific instructions for testing or validating your changes, include those as well.
Distributed under the Apache 2.0 License with LLVM Exceptions. See LICENSE and the LLVM License for more information.