SolenaV2

A from-scratch 336M-parameter language model built with JAX and Equinox

SolenaV2 is an experimental decoder-only Transformer with a complete pipeline for corpus preparation, BPE tokenization, pretraining, supervised fine-tuning, multi-host TPU execution, checkpointing, and autoregressive generation.

Important

SolenaV2 is a language-model project and research prototype. It is not a wrapper around an existing API and is not based on pretrained model weights. Trained checkpoints and compiled token arrays are not currently distributed with this repository.

Overview

SolenaV2 implements a generative language model from the ground up. The principal configuration contains 336,128,000 trainable parameters and was pretrained with a 1,024-token context window on a custom 891M-token corpus. The model, optimizer, data loader, distributed execution, tokenizer pipeline, and generation loop are all implemented in this repository.

The project is designed around three goals:

1. Build and train a complete language model without starting from a pretrained base.
2. Run the same code across small CPU tests and multi-host TPU configurations.
3. Keep the training system observable enough to diagnose optimization and scaling behavior.

Model at a glance

Component	Principal configuration
Architecture	Decoder-only Transformer
Parameters	336,128,000
Vocabulary	32,000 SentencePiece tokens
Context length	1,024 tokens
Transformer blocks	24
Model width	1,024
Attention heads	8
Head width	128
Feed-forward width	4,096
Positional representation	Learned absolute embeddings
Normalization	Pre-LayerNorm
Activation	GELU
Dropout	0.1 during pretraining
Embedding/output weights	Tied
Parameter precision	bfloat16
LayerNorm statistics and logits	float32
Pretraining objective	Next-token cross-entropy
Principal optimizer	Factored Adafactor
Principal hardware	16 TPU v6e chips

At context length 768, the shorter positional table reduces the model to 335,865,856 parameters.

Architecture

Each Transformer block contains:

• custom pre-normalization in float32
• multi-head causal self-attention through jax.nn.dot_product_attention
• a two-layer GELU feed-forward network
• residual connections around attention and the feed-forward network
• dropout on the attention output and feed-forward hidden activations

Token embeddings are added to learned positional embeddings before the first block. A final LayerNorm is followed by a tied projection through the token embedding matrix.

The large profiles enable activation rematerialization with eqx.filter_checkpoint. Language-model logits are evaluated in sequence chunks and cast to float32, preventing the complete batch-by-sequence-by-vocabulary tensor from occupying device memory at once.

Implemented systems

• SentencePiece BPE training with a configurable vocabulary
• Mixed-corpus streaming, filtering, chunking, and exact deduplication
• Deterministic document-level train/validation splitting
• Memory-mapped NumPy token arrays
• Random training windows and deterministic evenly spaced evaluation windows
• CPU, single-host data parallel, and multi-host JAX mesh execution
• Host-side batch prefetching
• Replicated model and optimizer state with sharded batches
• AdamW and factored Adafactor optimizer paths
• Optional Adafactor parameter-block scaling
• Masked weight decay for attention and feed-forward matrices
• Global update RMS, parameter RMS, and update-to-parameter logging
• Best-validation model checkpoints and metadata
• Optional Google Cloud Storage synchronization
• UltraChat supervised fine-tuning with assistant-only loss masks
• Interactive and one-shot text generation
• Top-k and nucleus sampling, repetition penalties, no-repeat n-grams, and stop tokens

Training corpus

The completed pretraining dataset contains:

Split	Tokens
Training	846,817,500
Validation	44,450,403
Total	891,267,903

The configured character-budget mixture is:

Category	Target share	Primary source
Educational web	45%	FineWeb-Edu-dedup
Encyclopedic text	25%	WikiText-103
Synthetic textbooks	15%	Cosmopedia v2
Stories	10%	TinyStories
Source code	5%	The Stack subsets

Some categories include fallback datasets. The exact source order is defined in training/prepdata.py.

Preprocessing performs the following operations:

1. Stream examples from each configured source.
2. Normalize whitespace and discard documents shorter than 200 characters.
3. Split long documents into bounded character chunks.
4. Remove exact duplicate chunks using SHA-1.
5. Write the configured source mixture to the raw corpus.
6. Assign complete documents to train or validation with a deterministic BLAKE2b hash.
7. Encode each document with SentencePiece and append an end-of-sequence token.

The tokenizer is a 32,000-entry SentencePiece BPE model with reserved padding, unknown, beginning-of-sequence, end-of-sequence, user, assistant, and turn-end tokens.

Repository layout

SolenaV2/
├── SolenaV2.py                     # generation CLI
├── config.py                       # model, runtime, and experiment profiles
├── models/
│   └── SolenaV2.py                 # Transformer implementation
├── training/
│   ├── prepdata.py                 # mixed pretraining corpus builder
│   ├── train_bpe.py                # SentencePiece tokenizer training
│   ├── encodedata.py               # deterministic split and token encoding
│   ├── prepdata_ultrachat_sft.py   # UltraChat SFT corpus builder
│   └── train.py                    # training and validation pipeline
└── utils/
    ├── dataset.py                  # training/evaluation window sampling
    ├── distributed.py              # JAX mesh and sharding utilities
    ├── gcs_cache.py                # remote artifact synchronization
    ├── preprocessing.py            # preprocessing helpers
    └── tokenizer.py                # SentencePiece wrapper

Installation

Python 3.10 or newer is required.

git clone https://github.com/LIXIRI0/SolenaV2.git
cd SolenaV2

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -r requirements.txt

On Windows, activate the environment with .venv\Scripts\activate.

JAX installation depends on the target accelerator. The command above is suitable for CPU use. GPU and TPU users should follow the official JAX installation guide or use the environment provided by the TPU VM.

Select a runtime profile

All entry points import config.py. Set the model profile and training stage before launching a script:

export SOLENA_PROFILE=cpu_dev
export SOLENA_TRAIN_STAGE=pretrain

Available profiles include:

Profile	Intended use
cpu_dev	Small local correctness and pipeline tests
cpu_full	Larger CPU-only experiments
collab_tpu	Smaller eight-device TPU environment
kaggle_tpu_8	Eight-device large-model profile
trc_tpu_16	Principal 16-chip TPU v6e profile
trc_tpu_64	64-chip mesh profile
tpu_train	General eight-device TPU training profile

The checked-in defaults are trc_tpu_16 and sft. Select explicit values for local execution.

Storage configuration

CPU profiles use repository-local data and checkpoint directories. TPU Research Cloud profiles use /tmp for fast local storage and can synchronize artifacts with Google Cloud Storage.

TPU profiles currently default to the original project's GCS prefix. Set a different bucket before running your own instance:

export SOLENA_GCS_ROOT=gs://YOUR_BUCKET/solena

To disable remote synchronization and use only local files:

export SOLENA_GCS_ROOT=""

GCS synchronization requires an installed and authenticated gcloud CLI. For Hugging Face sources that require authentication, export FACE_TOKEN or place it in a root-level .env file.

Pretraining

The complete pretraining pipeline has four stages:

export SOLENA_PROFILE=trc_tpu_16
export SOLENA_TRAIN_STAGE=pretrain
export SOLENA_GCS_ROOT=gs://YOUR_BUCKET/solena

# 1. Build the mixed raw-text corpus.
python training/prepdata.py

# 2. Train the SentencePiece BPE tokenizer.
python training/train_bpe.py

# 3. Create deterministic train and validation token arrays.
python training/encodedata.py

# 4. Train the language model.
python training/train.py

Warning

The configured pretraining target is approximately one billion tokens. It is not a quick-start workload. Reduce PRETRAIN_TARGET_TOKENS in config.py before a local end-to-end test.

For a fresh training process, disable checkpoint loading explicitly:

export SOLENA_RESUME=0
python training/train.py

Supervised fine-tuning

The SFT pipeline converts UltraChat conversations into the repository's role format:

<|user|>
...
<|assistant|>
...
<|end|>

It creates token-level masks so that loss is applied to assistant responses rather than user prompts. SFT requires the pretraining tokenizer and a base model checkpoint.

export SOLENA_PROFILE=trc_tpu_16
export SOLENA_TRAIN_STAGE=sft
export SOLENA_GCS_ROOT=gs://YOUR_BUCKET/solena

python training/prepdata_ultrachat_sft.py
python training/encodedata.py
python training/train.py

By default, SFT loads SolenaV2.eqx and writes SolenaV2-sft.eqx. Set SOLENA_SFT_RESUME_FROM_SFT=1 to continue from an existing SFT checkpoint.

Generation

Generation requires a tokenizer and a checkpoint matching the active profile. The script supports plain continuation for the base model and role-formatted prompts for the SFT model.

One-shot generation:

export SOLENA_PROFILE=trc_tpu_16
export SOLENA_TRAIN_STAGE=sft

python SolenaV2.py --prompt "Explain why the sky appears blue."

Interactive generation:

python SolenaV2.py

The generation loop streams decoded output and supports temperature, top-k, top-p, repetition penalties, a repetition window, no-repeat n-grams, minimum and maximum generation lengths, and optional sentence-boundary stopping.

Common environment variables

Variable	Purpose
SOLENA_PROFILE	Selects the model and hardware profile
SOLENA_TRAIN_STAGE	Selects pretrain or sft
SOLENA_DATA_DIR	Overrides the local data directory
SOLENA_CHECKPOINT_DIR	Overrides the checkpoint directory
SOLENA_GCS_ROOT	Sets the remote artifact prefix; empty disables GCS
SOLENA_RESUME	Enables checkpoint loading when nonzero
SOLENA_SAVE_CHECKPOINTS	Enables model checkpoint writes when nonzero
SOLENA_MAX_BATCHES	Caps optimizer steps per nominal epoch
SOLENA_EPOCHS_PER_RUN	Sets the number of epochs in an invocation
SOLENA_VAL_BATCHES	Sets deterministic validation batches
SOLENA_STEP_LOG_INTERVAL	Logs loss and RMS diagnostics every N steps
SOLENA_STEP_LOG_FIRST_N	Logs the first N optimizer steps
SOLENA_ADAFACTOR_PARAM_SCALE	Enables Adafactor parameter scaling
SOLENA_WEIGHT_DECAY	Sets the nominal weight-decay coefficient
SOLENA_SFT_LR	Overrides the supervised fine-tuning rate
SOLENA_LOGIT_CHUNK_SIZE	Controls output-logit memory chunking
SOLENA_PREFETCH_BATCHES	Controls host-side prefetch depth
SOLENA_GEN_MAX_NEW_TOKENS	Sets the generation limit
SOLENA_GEN_TEMPERATURE	Sets sampling temperature
SOLENA_GEN_TOP_P	Sets nucleus-sampling probability mass
SOLENA_GEN_TOP_K	Sets the top-k sampling limit

Checkpoint behavior

Equinox checkpoints contain model leaves and JSON metadata. They do not contain Optax optimizer state. Starting a new training process also reconstructs the data and dropout RNG streams from their configured seeds. A resumed process is therefore a model-weight continuation, not an exact continuation of every piece of training state.

The repository saves a new checkpoint only when the selected validation metric improves. Signal handlers attempt to synchronize the last completed checkpoint and logs before a managed shutdown.

Related research

SolenaV2 training runs were used as the experimental system for a separate optimizer case study. The paper analyzes Adafactor step-size semantics; it does not define the identity or purpose of the SolenaV2 language model.

Nisa Hanay. When Parameter Scaling Stalls Training: Diagnosing Adafactor Step-Size Semantics in a 336M-Parameter Transformer. Zenodo, 2026. https://doi.org/10.5281/zenodo.21296326

Reproducibility notes

• Training windows are sampled with replacement. An epoch in the logs denotes an approximately corpus-sized sampled-token budget.
• Evaluation windows are selected deterministically across the validation array.
• datasets, huggingface-hub, and pyarrow are pinned, while JAX, Equinox, Optax, and SentencePiece versions are currently unpinned.
• The original corpus-building logs did not preserve which fallback source was selected for each category.
• Model checkpoints and compiled training data are not included in the repository.

Citation

To cite the software repository:

@software{hanay2026solenav2,
  author  = {Hanay, Nisa},
  title   = {SolenaV2: A From-Scratch Transformer Language Model in JAX},
  year    = {2026},
  url     = {https://github.com/LIXIRI0/SolenaV2}
}

If you are specifically referencing the Adafactor findings, cite the Zenodo paper linked in the related-research section instead.

License

MIT License