A hands on research laboratory for rebuilding modern language model components, changing their assumptions, and measuring whether the changes actually improve training or inference.
The first complete project in this repository is ForgeLM, a compact decoder only language model written in PyTorch. PyTorch is used for tensors and automatic differentiation, but the tokenizer, attention mechanism, Transformer blocks, optimizer, mixture of experts layer, sampling logic, quantization utilities, and training loop are implemented directly in this repository.
This is not intended to be a production foundation model. It is designed to make serious machine learning ideas small enough to inspect, alter, benchmark, and explain.
- Rebuild the important parts of an LLM rather than hiding them behind a high level model library.
- Add original experimental changes to standard components.
- Keep a clean baseline for every custom idea.
- Make each experiment measurable with tests, probes, and ablations.
- Separate demonstrated behavior from unverified research hypotheses.
ForgeLM contains three primary custom mechanisms and more than ten supporting experiments.
Standard approach
Byte Pair Encoding repeatedly merges the most frequent adjacent token pair. Frequency is useful, but it can prioritize pairs that occur many times inside one repeated phrase.
ForgeLM change
The EntropyAwareBPETokenizer scores a pair using both occurrence frequency and the diversity of its surrounding contexts. A pair that appears next to many different left and right neighbors receives a higher contextual entropy score.
Conceptually, the merge score is:
merge score = pair frequency × contextual diversity adjustment
Why it may matter
A pair that appears in varied contexts may represent a reusable linguistic unit rather than a memorized fragment. This could improve vocabulary efficiency and reduce the number of tokens required for unseen text.
What must be tested
Compare it with ordinary frequency only BPE using the same corpus and vocabulary size. Measure token count, average tokens per word, unknown domain compression, training throughput, and validation perplexity.
Implementation
forgelm/tokenizer.py
Standard approach
Full causal attention allows each token to attend to all earlier tokens. It captures long range relationships but requires quadratic attention work. Sliding window attention is cheaper but cannot directly retrieve information outside its window.
ForgeLM change
Each attention head learns a gate between two attention distributions:
head output = gate × global attention + (1 − gate) × local attention
The gate is learned independently for every head. A head can become mostly local, mostly global, or remain mixed.
Why it may matter
Different heads perform different functions. Some may primarily model nearby syntax while others retrieve distant information. Learning the split avoids manually assigning certain heads to local or global attention.
Important limitation
The current implementation computes both distributions, so it does not yet reduce training cost. It first studies specialization. A later sparse version could skip the unused path once gates become decisive.
What must be tested
Track gate values by layer and head, attention entropy, long context retrieval accuracy, perplexity, and whether gates converge consistently across random seeds.
Implementation
forgelm/attention.py
Standard approach
AdamW keeps exponential averages of gradients and squared gradients, then applies decoupled weight decay.
ForgeLM change
AdamFlux adds two controls to an AdamW style update.
- Gradient surprise measures how strongly the current gradient disagrees with its recent exponential history. Large disagreement reduces the effective step.
- Parameter trust limits update magnitude relative to the norm of the parameter being updated.
Why it may matter
Language model optimization can encounter sudden gradient changes and layers with very different parameter scales. Surprise control may reduce unstable steps, while trust scaling may stop a small parameter tensor from receiving an update that is disproportionately large.
Important limitation
These controls can also slow useful adaptation. AdamFlux is an experimental optimizer, not a claimed replacement for AdamW.
What must be tested
Compare AdamFlux with AdamW at matched learning rates and across learning rate sweeps. Measure convergence speed, final validation loss, gradient norm, update to weight ratio, failed runs, and sensitivity to hyperparameters.
Implementation
forgelm/optim.py
forgelm/position.py
Rotary Position Embeddings rotate query and key features according to token position. Unlike learned absolute embeddings, the positional relation is inserted directly into the attention dot product. This implementation supports position offsets, which are required when generating with a key value cache.
forgelm/norm.py
RMSNorm scales hidden states according to their root mean square magnitude without subtracting the mean. It uses fewer operations than LayerNorm and is common in modern decoder models. The implementation is written directly rather than calling a built in RMSNorm layer.
forgelm/activations.py
A normal SwiGLU feedforward network uses a SiLU activated gate. ForgeLM learns a scalar interpolation between SiLU and GELU:
activation = mix × SiLU(x) + (1 − mix) × GELU(x)
This tests whether the preferred gate shape changes during training.
forgelm/layers.py
A standard Transformer adds every attention and feedforward update at full strength. ForgeLM predicts a residual multiplier for each token. The gate starts near open so the initial network behaves similarly to a normal residual model.
The current gate changes update strength but does not skip computation. A future version could connect decisive gates to conditional execution.
forgelm/moe.py
The mixture of experts layer contains multiple feedforward experts. A router selects the best k experts for each token and combines their outputs. An auxiliary balancing loss discourages the router from sending nearly every token to the same expert.
Compared with a dense feedforward layer, MoE increases parameter capacity without activating every parameter for every token. The current implementation emphasizes clarity rather than optimized distributed execution.
forgelm/lora.py
Low Rank Adaptation freezes an existing linear layer and learns a small low rank update. This allows parameter efficient fine tuning. The repository includes a LoRALinear wrapper and utilities for injecting adapters into selected linear layers.
forgelm/quantization.py
The quantization utilities map floating point tensors to signed 8 bit integers with a symmetric scale and reconstruct them for error measurement. This is fake quantization for experimentation, not a custom high performance inference kernel.
forgelm/attention.py and forgelm/model.py
During autoregressive generation, earlier keys and values do not need to be recomputed for every new token. ForgeLM stores them per layer and appends only the new token state. Rotary position offsets ensure the new token receives the correct position.
forgelm/sampling.py
The sampler supports:
- Temperature scaling
- Greedy decoding when temperature is zero
- Top K filtering
- Nucleus or Top P filtering
- Min P filtering relative to the most likely token
- Repetition penalties
These controls are implemented separately so their effects can be tested rather than hidden inside a generation library.
forgelm/speculative.py
A small draft model proposes tokens, while a larger target model verifies them. The research question is whether several draft tokens can be accepted for each expensive target model pass. This implementation prioritizes understanding the algorithm and is not an optimized serving engine.
forgelm/interpretability.py
Attention entropy measures whether a head spreads probability over many keys or concentrates on a small number. It is useful for examining whether hybrid heads become specialized and whether attention collapses.
forgelm/evaluation.py
The repository includes perplexity and distinct N metrics. Perplexity measures predictive uncertainty on held out data. Distinct N estimates the diversity of generated token sequences. Neither metric alone measures output quality, so they should be combined with task specific evaluations.
forgelm/data.py and forgelm/trainer.py
The dataset creates next token prediction windows from a token sequence. The trainer handles optimization, gradient clipping, checkpoint creation, and device placement. It is deliberately small so the complete training path can be inspected.
ml-experiments/
│
├── forgelm/
│ ├── __init__.py Public package interface
│ ├── config.py Model and experiment configuration
│ ├── tokenizer.py Byte tokenizer and entropy aware BPE
│ ├── position.py Rotary position embeddings
│ ├── norm.py RMSNorm implementation
│ ├── activations.py Learned SiLU and GELU mixed SwiGLU
│ ├── attention.py Hybrid local and global causal attention
│ ├── layers.py Transformer block and residual gates
│ ├── moe.py Top K mixture of experts
│ ├── model.py Complete decoder only language model
│ ├── optim.py AdamFlux optimizer
│ ├── lora.py Low rank adaptation layers
│ ├── quantization.py Int8 quantization experiments
│ ├── sampling.py Generation filters and token sampling
│ ├── speculative.py Draft and target speculative decoding
│ ├── data.py Next token sequence dataset
│ ├── trainer.py Minimal training and checkpoint loop
│ ├── evaluation.py Perplexity and diversity metrics
│ └── interpretability.py Attention and parameter diagnostics
│
├── examples/
│ └── quickstart.py Small end to end model demonstration
│
├── experiments/
│ ├── optimizer_benchmark.py AdamW against AdamFlux
│ └── attention_probe.py Hybrid attention gate inspection
│
├── scripts/
│ ├── train_tiny.py Train a small model on a text file
│ └── generate.py Generate from a saved checkpoint
│
├── tests/
│ ├── test_tokenizer.py Tokenization round trip and vocabulary tests
│ ├── test_model.py Forward pass, loss, cache, and gradient tests
│ ├── test_optimizer.py AdamFlux optimization behavior
│ ├── test_moe.py Expert routing and balancing loss
│ └── test_extras.py LoRA, quantization, and sampling tests
│
├── docs/
│ └── experiments.md Controlled ablation plan
│
├── .github/workflows/
│ └── tests.yml Automated test workflow
│
├── pyproject.toml Package metadata and dependencies
├── CONTRIBUTING.md Contribution and experiment standards
└── LICENSE MIT License
For a batch of token identifiers, the model performs the following sequence:
- Convert token identifiers into learned embeddings.
- Pass hidden states through repeated Transformer blocks.
- Normalize the block input with RMSNorm.
- Compute rotary positioned queries, keys, and values.
- Form global and local causal attention distributions.
- Mix those distributions using learned per head gates.
- Add the token gated attention residual.
- Apply another RMSNorm.
- Run either Mixed SwiGLU or Top K MoE.
- Add the token gated feedforward residual.
- Apply the final RMSNorm.
- Project hidden states to vocabulary logits.
- Compute cross entropy loss for next token prediction when targets are supplied.
python3 -m venv .venv
source .venv/bin/activate
pip install -e '.[dev]'ForgeLM requires Python 3.10 or newer and PyTorch 2.2 or newer.
python examples/quickstart.pyThe quickstart creates a small model, runs a forward pass, reports its loss, and generates a short continuation.
Create a UTF 8 text file, then run:
python scripts/train_tiny.py \
--text data.txt \
--steps 500 \
--context 128 \
--batch-size 16 \
--device cpu \
--checkpoint checkpoints/final.ptFor an NVIDIA GPU, use --device cuda. For Apple Silicon with an appropriate PyTorch build, use --device mps.
python scripts/generate.py \
--checkpoint checkpoints/final.pt \
--prompt "The future of computing" \
--tokens 120 \
--device cpufrom forgelm import ForgeConfig, ForgeLM
config = ForgeConfig(
vocab_size=512,
context_length=128,
d_model=256,
n_heads=8,
n_layers=6,
d_ff=768,
local_window=32,
use_moe=False,
)
model = ForgeLM(config)Enable the mixture of experts path with:
config = ForgeConfig(
vocab_size=512,
use_moe=True,
n_experts=4,
experts_per_token=2,
)python experiments/optimizer_benchmark.py --steps 200 --lr 0.003 --seeds 5This is a small nonlinear regression sanity check. It is not evidence of language model superiority. Its purpose is to reveal obvious instability before expensive training.
python experiments/attention_probe.pyThis probe trains or evaluates the learned local and global gate values and reports whether heads begin to specialize.
pytestThe current suite checks:
- Unicode safe byte tokenization
- Entropy aware BPE vocabulary growth and decoding
- Model output dimensions and finite loss
- Cached generation output shape
- Gradient flow into attention gates
- AdamFlux progress on a convex objective
- MoE routing and balancing loss
- LoRA adapter behavior
- Quantization reconstruction error
- Sampling filter correctness
Every custom component should be treated as a hypothesis. A valid comparison should hold constant:
- Training corpus and token order
- Number of training tokens
- Model parameter count, or clearly report the difference
- Batch size and gradient accumulation
- Learning rate search budget
- Random seeds
- Validation split
- Hardware and precision
Report at least:
- Validation loss and perplexity
- Tokens processed per second
- Peak memory usage
- Gradient norm
- Update to parameter norm
- Mean and standard deviation across seeds
- Number of unstable or failed runs
Start with a small baseline decoder and alter one variable at a time.
| Experiment | Baseline | Variant | Main question |
|---|---|---|---|
| Tokenizer | Frequency BPE | Entropy aware BPE | Does contextual diversity improve compression or perplexity? |
| Attention | Full causal | Learned local and global mixture | Do heads consistently specialize? |
| Optimizer | AdamW | AdamFlux | Does surprise control improve stability without slowing convergence? |
| Activation | SwiGLU | Mixed SwiGLU | Does the learned activation mixture move away from its initialization? |
| Residual | Standard addition | Token adaptive gate | Do tokens learn meaningfully different update strengths? |
| Feedforward | Dense SwiGLU | Top K MoE | Does added expert capacity justify routing cost? |
- Does contextual entropy select merges that generalize across domains?
- Does hybrid attention learn a repeatable division of local and global functions?
- Can attention gates be converted into actual sparse computation after training?
- Does AdamFlux tolerate more aggressive learning rates than AdamW?
- Which layers benefit from parameter trust scaling, and which are slowed by it?
- Does Mixed SwiGLU converge toward SiLU, GELU, or a persistent mixture?
- Can residual gates predict tokens that require little additional computation?
- Do MoE experts specialize by syntax, topic, frequency, or position?
- How much quality is lost under symmetric int8 quantization?
- When does speculative decoding produce a real speedup after verification cost?
The repository is a functional experimental baseline with unit tests and small diagnostic scripts. It has not yet established that the custom mechanisms outperform standard methods on large language model benchmarks. Any performance claim should be supported by controlled, repeated experiments.
Contributions should include:
- A clearly stated hypothesis
- A standard baseline
- A minimal implementation
- Tests for correctness
- A reproducible experiment command
- Results including failures and limitations
See CONTRIBUTING.md for the complete process.
MIT