---
library_name: transformers
license: apache-2.0
datasets:
- HuggingFaceFW/fineweb-edu
language:
- en
pipeline_tag: text-generation
tags:
- causal-lm
- base-model
- custom-architecture
- pytorch
- tiktoken
---

<p align="center">
  <img src="./axiom_logo.png" width="220">
</p>

# Axiom-Dense-380M-Base

Axiom-Dense-380M-Base is a decoder-only causal language model trained from scratch for general-purpose next-token prediction on English web text. This is a base pretrained model, not an instruction-tuned chat model.

## Model Summary

- Model type: decoder-only Transformer (causal LM)
- Parameter count: 385,849,344
- Context length: 1,024 tokens
- Vocabulary: 100,277 (`tiktoken` `cl100k_base`)
- Training objective: autoregressive next-token prediction
- Special handling: tied input/output embeddings (`embed.weight` tied to `lm_head.weight`)

## Architecture

This model follows a dense Transformer stack with grouped-query attention and rotary positional embeddings.

- Hidden size: 1024
- Layers: 24
- Attention heads: 16
- KV heads: 8 (GQA)
- FFN multiplier: 2.6667 (rounded to hardware-friendly multiple)
- Normalization: RMSNorm
- Positional encoding: RoPE (`theta=10000`)
- Activation: SwiGLU
- Dropout: 0.0

Implementation details are defined in:
- `model.py` (core architecture and generation)
- `config.py` (`ModelConfig`, `TrainConfig`)

## Training Data

- Source dataset: `HuggingFaceFW/fineweb-edu`, `sample-10BT` split
- Local dataset path during training: `data/fineweb-edu-10BT`
- Text field: `text`
- Validation split strategy: deterministic hash split with `val_fraction=0.001` and `split_seed=1337`
- Document boundary treatment: EOS token appended after each document

## Training Setup

- Target tokens: 8,000,000,000
- Effective tokens per optimizer step: 327,680 (`batch_size=1`, `seq_len=1024`, `grad_accum=320`)
- Computed optimizer steps: 24,414
- Planned tokens represented by training schedule: 7,999,979,520
- Optimizer: AdamW8bit (fallback to AdamW if unavailable)
- LR schedule: warmup, constant phase, cosine decay
- Warmup steps: 2,000
- LR max/min: 3e-4 / 1e-5
- Weight decay: 0.1
- Betas: (0.9, 0.95)
- Gradient clipping: 1.0
- Precision: bfloat16
- Gradient checkpointing: enabled
- Compile: disabled in provided config

## Evaluation Snapshot

Validation metrics in this repo are tracked in `eval.csv` at interval checkpoints.

- Best observed eval loss: 2.7394 at step 15,000
- Best observed eval perplexity: 15.4780 at step 15,000
- Final logged eval loss: 2.8972 at step 24,000
- Final logged eval perplexity: 18.1233 at step 24,000

These are internal development metrics on the project validation split, not a broad benchmark suite.

## Intended Use

- Continued pretraining
- Supervised finetuning or instruction tuning
- Research and experimentation on medium-scale dense LMs
- Educational use for studying custom Transformer implementations

## Out-of-Scope / Not Recommended

- Safety-critical or high-stakes decisions (medical, legal, financial)
- Direct deployment as a reliable assistant without task-specific alignment and evaluation
- Use cases requiring guaranteed factual accuracy

## Limitations

- Base model behavior: may produce repetitive, off-topic, or hallucinatory outputs
- No instruction tuning by default
- English-centric training distribution
- Context window limited to 1,024 tokens
- Bias/toxicity risks inherited from web-scale text data

## Safety and Risk Notes

Potential harms include generation of incorrect, biased, or unsafe text. Downstream users should add:

- Domain-specific evaluation
- Prompt and output safety filtering
- Human oversight for sensitive workflows
- Red-teaming before production release

## Tokenization

- Tokenizer backend: `tiktoken`
- Encoding: `cl100k_base`
- Vocab size: 100,277
- EOS token: tokenizer `eot_token`

## Reproducibility

Core files relevant to reproducibility:

- `train.py` (training loop, checkpointing, metrics)
- `data.py` (dataset packing/streaming and deterministic split logic)
- `model.py` (architecture)
- `config.py` (model/training hyperparameters)

Seed configuration:

- Python / NumPy / PyTorch seed: 1337

## Usage

This repository contains custom model/tokenizer code paths. Load with the project code or with Hugging Face `transformers` remote code support if published with matching `auto_map` files.