Ibbyml/lyra-small-1.6B

Model

Introduction

Lyra Small is a 1.62B parameter language model trained from scratch on ~500M tokens. This checkpoint is primarily meant to demonstrate correctness and perf for the Lyra LLM pretraining stack.

This is an early base-model checkpoint. It has not been instruction-tuned, and its text can still be repetitive or inaccurate.

Summary

Architecture	Mixture-of-Experts (MoE)
Total Parameters	1.62B
Active Parameters	1.0B
Number of Layers	16
Number of Dense Layers	12
Attention-Layer Composition	8 Global + 8 Local
Attention Mechanism	Gated GQA
Number of Attention Heads	16
MoE Hidden Dimension (per Expert)	2048
Dense Hidden Dimension	3072
Number of Routed Experts	15
Number of Shared Experts	1
Experts per Token	4 (3 Routed + 1 Shared)
Vocabulary Size	201088
Context Length	4096
Activation Function	SwiGLU

Hidden dimensions are the widths after SwiGLU gating.

Results

The run reached a validation loss (CE) of 3.06 and perplexity of 21.3 after ~500M tokens. Training took ~4h 40m on a single TPU v6e, with training updates running at ~34K tokens/s and compiled peak memory of 28.8 GiB.

Training

Optimizer	Muon + Adam
Total Training Tokens	~500M
Tokens Per Update	65536
LR Schedule	Cosine (Warmup + 50% plateau)
Peak LR	Muon 0.0015, Adam 0.0002
Precision	FP32 Parameters, BF16 Compute
Hardware	1x TPU v6e

Usage

Download the Model

On a TPU Machine:

git clone https://github.com/ibbyml/lyra.git
cd lyra
uv sync --extra tpu

# Download the model (~6GB)
uv run hf download Ibbyml/lyra-small-1.6B --local-dir ./lyra-small

Run with JAX

uv run scripts/serve.py \
  --variant lyra-small \
  --weights ./lyra-small/model \
  --device tpu \
  --max-new-tokens 128

The checkpoint contains model weights in Orbax format.

Citation

If you find Lyra useful in your research, please cite:

@misc{Lyra,
  author = {Ibrahim Khan},
  title = {Lyra: A performant LLM pretraining stack in pure JAX.},
  year = {2026},
  publisher = {GitHub},
  url = {https://github.com/ibbyml/lyra}
}

License

MIT