Zlynx Documentation
CAUTION
Zlynx is currently an experimental library. APIs are subject to change without notice. We are in the early stages of development and welcome feedback.
An experimental, research-oriented deep learning library built on JAX and Flax NNX. Zlynx explores providing fine-grained control over model architectures, training loops, and distributed setups without the weight of larger frameworks.
🚀 Getting Started
If you are exploring Zlynx, these guides offer an early look at setting up and running initial experiments.
- Installation — Basic setup for Zlynx and its dependencies.
- Quick Start — A brief overview of current experimental workflows.
- Create a Model — Defining architectures using the experimental
Zbase class. - Training — An introduction to the current
Trainerimplementation. - Checkpointing — Early support for saving and loading models.
💡 Concepts & Explorations
Deep dives into the experimental interfaces and ideas currently in Zlynx.
- Sharding & Parallelism — Initial concepts for distributed training.
- PEFT (LoRA, etc.) — Experimental parameter-efficient fine-tuning utilities.
- GaLore — Exploratory memory-efficient full fine-tuning.
- Logging Backends — Current state of logging integrations.
📚 Examples
See how Zlynx can be used in its current state.
- MNIST Tutorial — A basic training example using Zlynx.
🛠️ API Reference (WIP)
Documentation for the available parts of the Zlynx project.
- Core API — Base classes and initial utilities.
- Model API — Built-in architectures (Under Development).
- Module API — Current building blocks: Attention, MLP, RoPE, and PEFT.
- Trainer API — The experimental training loop and configuration.
Project Goals
- Exploratory — Built for research and experimentation with JAX/NNX.
- Modularity — Exploring reusable blocks for custom architectures.
- Scalability — Aims to simplify sharding and parallelism in the long term.
- Integration — Initial efforts to support industry standards like SafeTensors and HuggingFace Hub.