1. Executive Design Summary
W-1.1-125M is a modern, decoder-only autoregressive Transformer designed to be trained from scratch. It keeps the model small enough for an independent research project while using current architectural choices that make training and inference efficient: RoPE positional encoding, RMSNorm, SwiGLU feed-forward layers, causal self-attention, tied token embeddings, BF16 training, Flash-Attention compatibility, and KV caching.
Item | Frozen value |
|---|---|
Architecture | Decoder-only causal Transformer |
Layers | 14 |
Model width (d_model) | 768 |
Attention heads | 12 |
Head dimension | 64 |
SwiGLU hidden width | 2,048 |
Vocabulary | 32,000 |
Maximum context | 2,048 tokens for v1 |
Position method | RoPE |
Normalization | Pre-norm RMSNorm |
Parameters | 123.69M core trainable parameters |
Objective | Next-token prediction |
2. Exact Model Architecture
2.0 Why this model uses 14 layers
The 14-layer depth is a deliberate parameter-allocation choice. With a 32,000-token vocabulary, tied embeddings consume 24.58M parameters. Fourteen 768-wide blocks with a 2,048-wide SwiGLU MLP use the remaining budget efficiently and produce a 123.69M-parameter model without enlarging the vocabulary or introducing arbitrary dimensions. A 12-layer version with the same widths would be about 109.53M parameters, materially below the target class. The 14-layer choice must still be validated through training stability and quality measurements; it is a design rationale, not a claim of superiority over every 12-layer alternative.
2.1 Input and token embedding
- Input is a batch of integer token IDs with shape [batch, sequence_length].
- A learned token embedding matrix maps each token ID to a 768-dimensional vector.
- Embedding matrix shape: 32,000 × 768 = 24,576,000 parameters.
- No learned absolute positional embedding is used; position information is introduced inside attention through RoPE.
- Embedding and output LM-head weights are tied, so the same 32,000 × 768 matrix is reused for output logits.
2.2 Transformer block × 14
Every block uses pre-normalization. The block order is frozen as follows:
- RMSNorm on the residual stream
- Causal multi-head self-attention with RoPE
- Residual addition
- RMSNorm
- SwiGLU MLP
- Residual addition
This ordering improves training stability and keeps the residual path clean. The implementation should use residual additions in FP32 when practical, while the main matrix multiplications run in BF16 on supported GPUs.
2.3 Causal self-attention
Property | Value |
|---|---|
Input/output width | 768 |
Number of query heads | 12 |
Key/value heads | 12 (standard MHA) |
Head dimension | 64 |
Q projection | 768 → 768 |
K projection | 768 → 768 |
V projection | 768 → 768 |
Output projection | 768 → 768 |
Mask | Strict causal / no future-token attention |
Fast kernel | PyTorch SDPA / Flash Attention when available |
For each position, attention is computed from Q, K, and V. The scale factor is 1/√64 = 1/8 before the softmax. RoPE is applied to Q and K. Padding must be handled by the data pipeline or an explicit attention mask when batches contain padding.
2.4 Rotary positional encoding (RoPE)
- RoPE dimension: full 64 dimensions of each attention head.
- Base theta: 10,000 for the frozen v1 baseline.
- Maximum trained context: 2,048 tokens. Longer context is a future extension and should be evaluated rather than assumed to work.
- Cache RoPE sin/cos or equivalent frequency tensors to avoid recomputing them every forward pass.
2.5 RMSNorm
- RMSNorm before attention and before the MLP in every block, plus one final RMSNorm after block 14.
- Normalization width: 768. Use epsilon = 1e-5.
- RMSNorm has a learned scale vector and no learned bias.
2.6 SwiGLU MLP
Projection | Shape | Parameters/block |
|---|---|---|
Gate | 768 × 2,048 | 1,572,864 |
Up | 768 × 2,048 | 1,572,864 |
Down | 2,048 × 768 | 1,572,864 |
Total linear weights | 4,718,592 |
The MLP computes SiLU(gate(x)) ⊙ up(x), then projects the result back to 768 dimensions. Linear layers are bias-free in the frozen design.
2.7 Final norm and language-model head
- Apply final RMSNorm to the output of the last Transformer block.
- Project 768-dimensional states to 32,000 vocabulary logits using the transposed token-embedding matrix.
- Use cross-entropy loss against the next token. Shift inputs and targets by one position; ignore masked/padding positions if any exist.
3. Exact Parameter Budget
Component | Formula | Parameters |
|---|---|---|
Token embedding / tied LM head | 32,000 × 768 | 24,576,000 |
Attention weights per block | 4 × 768 × 768 | 2,359,296 |
SwiGLU weights per block | 3 × 768 × 2,048 | 4,718,592 |
Two RMSNorms per block | 2 × 768 | 1,536 |
One Transformer block | 7,079,424 | |
14 Transformer blocks | 14 × 7,079,424 | 99,111,936 |
Final RMSNorm | 768 | 768 |
TOTAL | 123,688,704 |
The project name “125M” is a model-class label. The frozen architecture contains approximately 123.69 million trainable parameters; this is preferable to distorting dimensions merely to hit an arbitrary exact count.
Transformer Architecture Flowchart
Conceptual encoder-decoder flow based on the Transformer design introduced by Vaswani et al. (2017). W-1.1-125M is decoder-only, so it keeps the causal self-attention/MLP pattern and does not use an encoder or cross-attention.
Figure 1. Original simplified Transformer flowchart, conceptually based on Attention Is All You Need; drawn for this specification rather than copied from the paper.
4. Tokenizer Specification
Item | Design |
|---|---|
Tokenizer family | Byte-level BPE |
Vocabulary size | 32,000 tokens |
Training source | Representative sample of the final pretraining mixture |
Languages | English + Urdu + Kashmiri + selected multilingual text |
Code | Include representative programming-language text |
Fallback | Byte-level coverage so unknown Unicode text remains representable |
Normalization | Minimal; preserve case, punctuation and meaningful Unicode distinctions |
Special-token IDs | PAD=0, BOS=1, EOS=2, DOC_SEP=3; remaining IDs assigned by tokenizer |
Training sample | At least 10–20B raw characters, stratified by intended mixture |
Quality checks | Round-trip Unicode tests; fertility by language; code, number and punctuation tests |
Versioning | Tokenizer vocabulary/merges are immutable after pretraining starts |
A tokenizer should be trained before the model. Its exact files, vocabulary order, special-token IDs, normalization behavior, and hash must be stored with every checkpoint so inference always uses the same token mapping.
5. Pretraining Dataset Target
The frozen compute target is 5 billion training tokens after tokenization. Dataset proportions are targets, not permission to ignore quality, licensing, or duplication. Every source should be filtered and tracked by provenance.
Data group | Target share | Approx. tokens | Purpose |
|---|---|---|---|
FineWeb-Edu / quality English | 65% | 3.25B | General knowledge and educational text |
Multilingual / FineWeb-2-derived | 20% | 1.00B | Urdu, Kashmiri where available, and multilingual robustness |
Licensed code corpus | 10% | 0.50B | Programming and code structure |
Curated STEM/math | 5% | 0.25B | Science and mathematical language |
- Total tokens: 5.0B, equal to about 40 training tokens per model parameter.
- Approximate clean UTF-8 text footprint: about 20–30 GB, strongly dependent on language mix and tokenizer efficiency.
- Packed token IDs occupy about 10 GB as uint16 (32K vocabulary fits in 16 bits), or about 20 GB as int32, before indexes/metadata.
- Keep train, validation, and test documents separated before packing to reduce leakage.
- The 5.0B figure means tokens actually consumed from the training split. Validation and test tokens are additional and never count toward the optimization budget.
- Create document-level held-out sets before packing: target at least 10M validation tokens and 10M test tokens, stratified by language and domain where volume permits.
5.1 Dataset governance and licensing
Requirement | Frozen rule |
|---|---|
Source approval | Every source must have a documented license or usage basis before ingestion. |
Provenance | Retain source name, snapshot/version, retrieval date, document ID and processing history. |
License filters | For code, allow only approved licenses; retain repository/file license metadata and attribution duties. |
Access terms | Respect applicable access restrictions and dataset terms; public availability is not automatically permission. |
Privacy | Filter credentials, private identifiers and sensitive personal data where feasible; document residual risk. |
Removal | Maintain document hashes and source IDs so records can be located and excluded from later dataset versions. |
Release | Publish a dataset card and processing statistics; do not redistribute source text unless permitted. |
6. Data Processing Pipeline
- Acquire only permitted datasets and record source/license metadata.
- Normalize encoding and remove broken/corrupt records while preserving meaningful Unicode.
- Apply language and quality filtering.
- Remove exact duplicates and perform near-duplicate/document-level deduplication.
- Remove obvious secrets, credential dumps, personally identifying data where feasible, spam, and low-quality boilerplate.
- Split documents into train/validation/test before sequence packing.
- Train/freeze the 32K byte-level BPE tokenizer on a representative sample.
- Tokenize documents and append EOS/document boundary tokens.
- Pack token streams into 2,048-token training sequences without silently mixing validation/test data.
- Save deterministic shards plus manifest, checksums, counts, and tokenizer version.
7. Training Configuration
Hyperparameter | Frozen starting value |
|---|---|
Objective | Autoregressive next-token cross entropy |
Precision | BF16 preferred; FP16 only with loss scaling if BF16 unavailable |
Optimizer | AdamW |
Adam betas | β1 = 0.9, β2 = 0.95 |
Epsilon | 1e-8 |
Peak learning rate | 6e-4 |
LR schedule | Linear warmup, then cosine decay |
Warmup | 96 optimizer updates |
Final LR | 6e-5 |
Weight decay | 0.1; exclude norm scales and embeddings |
Gradient clipping | Global norm 1.0 |
Dropout | 0.0 |
Sequence length | 2,048 |
Global batch | 524,288 non-padding tokens/update (256 full sequences) |
Optimizer steps | 9,537 updates to consume at least 5.0B tokens |
Initialization | N(0, 0.02); scale attention-output and MLP-down std by 1/√(2 × 14) |
Checkpoint cadence | Every 250 updates; permanent milestones every 1,000 updates |
Validation cadence | Every 100 updates on a fixed held-out token budget |
These values define the frozen full-run recipe. First run a separate 50–200-step pilot that does not count toward the 5B-token run. It must verify stable loss, finite gradients, checkpoint/resume equivalence, memory headroom and adequate input throughput. Gradient accumulation and data parallelism may change, but their product must preserve the 524,288-token global batch.
8. Training Runtime and Efficiency
- Use PyTorch DistributedDataParallel for multi-GPU training; use FSDP only if required by memory scale.
- Use scaled-dot-product attention / Flash-Attention kernels where supported.
- Enable gradient checkpointing only when memory pressure justifies the extra compute.
- Use fused AdamW and fused normalization/activation kernels when stable on the target stack.
- Pre-tokenize and memory-map/shard data so GPU training is not starved by text processing.
- Log tokens/second, model FLOP utilization where available, training/validation loss, LR, gradient norm, GPU memory, and checkpoint IDs.
9. Checkpoints and Reproducibility
- Checkpoint model weights, optimizer state, scheduler state, scaler state if used, global step, tokens seen, and RNG states.
- Store architecture config and tokenizer files beside checkpoints.
- Record dataset manifest/version, shard checksums, source proportions, git commit, package versions, and training command/config.
- Resume must reproduce the data position or use a documented resampling policy.
- Save weights in safetensors for distribution when possible; training-state checkpoints may use framework-native formats.
10. Validation and Evaluation
Stage | What to check |
|---|---|
Before full run | Forward/backward correctness, causal mask, RoPE, tied weights, checkpoint resume, deterministic tiny-run overfit |
During pretraining | Train/validation loss, perplexity, gradient norm, throughput, NaN/Inf detection |
Base-model evaluation | Held-out perplexity plus age-appropriate language, factual, reasoning, math and code benchmarks |
Multilingual evaluation | Separate English, Urdu and Kashmiri-capable evaluation where reliable test data exists |
Contamination checks | Search for overlap/near-duplicates between evaluation sets and training corpus |
Generation checks | Repetition, degeneration, EOS behavior, context handling and latency |
10.1 Release acceptance gates
Gate | Pass condition |
|---|---|
Correctness | All unit/integration tests pass, including causal masking, RoPE, tied weights, KV-cache equivalence and exact checkpoint resume. |
Numerical stability | No unexplained NaN/Inf event; validation loss is finite; gradient norms are monitored and recoverable. |
Learning | Final held-out loss beats documented untrained and pilot baselines, with no unexplained late-run regression. |
Data integrity | No train/validation/test document-ID overlap; near-duplicate leakage audit completed and recorded. |
Language coverage | Report tokenizer fertility and held-out loss separately for English, Urdu and any reliable Kashmiri test set. |
Benchmark reporting | Publish benchmark versions, prompts, decoding settings, sample counts and uncertainty where applicable. |
Serving | Reference and KV-cached logits match within chosen tolerance; context/output limits and malformed requests are tested. |
Documentation | Model card, data card, licenses, limitations, intended use, hashes and reproducibility configuration are present. |
Numeric quality thresholds for perplexity and downstream benchmarks must be set after the pilot establishes a credible baseline. This specification intentionally does not fabricate target scores or experimental results.
11. Post-Training: Turning the Base Model into a Chat Model
Pretraining produces a base next-token model, not a polished assistant. Chat behavior should be added as a separate, versioned stage so the original base model remains available.
- Supervised instruction tuning on a carefully licensed, high-quality conversation/instruction dataset.
- Define a stable chat template and reserve/version any required control tokens.
- Run safety and refusal-behavior evaluation appropriate to the deployment domain.
- Optionally run preference optimization after supervised tuning if a trustworthy preference dataset and evaluation suite are available.
- Re-evaluate general-language quality to detect catastrophic degradation from post-training.
12. Production Inference Architecture
Capability | Requirement |
|---|---|
Model loading | Config-driven loading; strict tokenizer/model version compatibility |
KV cache | Per-layer K/V caching for autoregressive decoding |
Batching | Continuous/dynamic batching if serving multiple users |
Precision | BF16/FP16 GPU serving; optional INT8/INT4 quantized build after quality validation |
Sampling | Temperature, top-p, top-k, repetition controls and max-new-token limit |
Streaming | Token streaming supported by server API |
Stop handling | EOS + application-defined stop sequences |
Observability | Request latency, tokens/s, queue time, errors and model version |
Limits | Context/token limits enforced server-side |
Packaging | Base model, tokenizer, generation config, model card and license/provenance notes |
13. Model Configuration Contract
The implementation should expose these values in a single version-controlled model configuration, rather than scattering constants through model.py:
Key | Value |
|---|---|
vocab_size | 32000 |
d_model | 768 |
n_layers | 14 |
n_heads | 12 |
head_dim | 64 |
ffn_hidden_size | 2048 |
max_seq_len | 2048 |
rope_theta | 10000.0 |
rms_norm_eps | 1e-5 |
attention_type | causal_mha |
activation | swiglu |
tie_word_embeddings | true |
linear_bias | false |
dropout | 0.0 |
14. Implementation Modules
Module | Responsibility |
|---|---|
config.py | Frozen architecture and training configuration |
tokenizer/ | Tokenizer training, serialization, encode/decode tests |
model.py | RMSNorm, RoPE, attention, SwiGLU, TransformerBlock, GPT model |
data.py | Shard reading, packing, distributed sampling |
train.py | Training loop, AMP/BF16, accumulation, distributed execution |
checkpoint.py | Save/resume and state validation |
evaluate.py | Loss/perplexity and benchmark adapters |
generate.py | KV-cached generation and sampling |
serve.py | Production inference API/streaming layer |
tests/ | Unit tests, shape tests, causal-mask tests, resume tests, tiny overfit test |
15. Definition of “Production Ready” for This Project
- Architecture is configuration-driven and parameter count is unit-tested.
- Tokenizer is frozen, versioned and packaged with every checkpoint.
- Dataset pipeline is reproducible, deduplicated, split safely, and source/licensing metadata is retained.
- Training can resume from checkpoint without corrupting optimizer or data state.
- Automated tests cover tensor shapes, masking, RoPE, weight tying, generation cache, serialization and a tiny overfit run.
- Evaluation is run on held-out data before and after post-training.
- Inference enforces context/output limits and exposes health/metrics endpoints.
- Quantized builds are released only after comparing quality against the reference precision model.
- A model card documents architecture, data categories, intended use, limitations, evaluation results and license information.
16. Final Frozen Blueprint
32K BPE tokens → 768-d embeddings → 14 × [RMSNorm → 12-head causal MHA + RoPE → residual → RMSNorm → 2,048-d SwiGLU → residual] → final RMSNorm → tied 32K LM head → next-token logits
Training target: 5 billion tokens at a maximum sequence length of 2,048. The result is the W-1.1-125M base model (~123.69M trainable parameters), after which instruction tuning and deployment are separate, versioned stages.
17. Architecture Freeze Checklist
- 14 layers / 768 width / 12 heads / 64 head dimension
- 2,048 SwiGLU hidden size
- 32,000 byte-level BPE vocabulary
- 2,048-token trained context
- RoPE θ = 10,000
- RMSNorm ε = 1e-5
- Bias-free linear layers
- Tied input/output embeddings
- 5B-token pretraining target
- BF16-first training
- AdamW + warmup + cosine schedule
- KV-cache inference
- Base checkpoint preserved before chat post-training
18. Research Basis and References
The architecture is informed by the following primary technical sources. They justify individual mechanisms; they are not experimental evidence that W-1.1-125M has achieved a particular quality level.
- Vaswani et al. (2017), “Attention Is All You Need.” arXiv:1706.03762.
- Radford et al. (2019), “Language Models are Unsupervised Multitask Learners.” OpenAI technical report.
- Su et al. (2021), “RoFormer: Enhanced Transformer with Rotary Position Embedding.” arXiv:2104.09864.
- Zhang and Sennrich (2019), “Root Mean Square Layer Normalization.” arXiv:1910.07467.
- Shazeer (2020), “GLU Variants Improve Transformer.” arXiv:2002.05202.
- Dao et al. (2022), “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” NeurIPS 2022.
- Hoffmann et al. (2022), “Training Compute-Optimal Large Language Models.” arXiv:2203.15556.
- Brown et al. (2020), “Language Models are Few-Shot Learners.” NeurIPS 2020.
End of specification. Any architecture, tokenizer, data-mixture, or optimizer change after the full run begins must create a new version rather than silently modifying W-1.1-125M.