A small code language model trained from scratch.
Pretraining data: CodeSearchNet (Python subset)
Tokenizer: Byte-level BPE trained with the Hugging Face tokenizers library (hf_bpe_tokenizer.py), using a code-aware pretokenization regex (splits on word/number/punctuation/whitespace runs before byte-level encoding). Byte alphabet occupies ids 0-255 directly; <|endoftext|> is appended as the last id after training, matching GPT-2's layout. A from-scratch reference implementation lives in bpe_tokenizer.py for learning purposes but isn't used for production training.
Vocab size was chosen via a compression-rate sweep (bytes-per-token on held-out test_set, tokenizer trained on train_set) — see 01-tokenizer.ipynb for the full table and reasoning. Current pick: vocab_size = 2,000.
Dataset: dataset.py holds CustomDataset, which slices a flat token-id stream into fixed-size, non-overlapping (context, target) windows via a single vectorized reshape. 02-dataset.ipynb encodes each split with the production tokenizer, joins documents with <|endoftext|>, and caches the resulting id lists as data/*.pt. Both the tokenizer and cached indices are hosted on a private Hugging Face Hub dataset repo so the same notebook runs unmodified locally or on Colab.
Training budget: full train_set encoded with the production tokenizer = 163,272,100 tokens/epoch. Planned 4 epochs ≈ 653M training tokens.
Model size: targeting 30-40M params (incl. embeddings), matching the Chinchilla-optimal point for a 1x L4 GPU / 1 hour run at 30% MFU (~33M params, ~660M tokens, 20 tokens/param — see 03-model.ipynb).
Baseline: first working end-to-end run, 2 epochs on an L4 GPU (plain SGD, no momentum) — train loss 7.78 → 6.16, val loss 6.46 → 6.13. Confirms the full pipeline (tokenizer → dataset → model → optimizer) trains correctly. Model and training code have since moved to model.py/train.py, with F.cross_entropy for the loss, GPT-style weight initialization (scaled residual output projections), and AdamW as the optimizer.
Local iteration run: for fast iteration on a MacBook (mps backend, no full-size GPU), 04-train.ipynb slices the cached training indices down to a fixed NUMBER_OF_TRAIN_TOKENS (10M tokens) instead of the full ~163M-token train set, and uses a much smaller model (1,628,672 trainable params: embedding_dim=128, proj_dim=256, n_heads=4, n_layers=4, context_size=512) — a smoke-test setup, not the production-scale target from 03-model.ipynb. 4 epochs with AdamW: train loss 4.21 → 2.51, val loss 3.67 → 2.73. Checkpoints save model_state_dict, optimizer_state_dict, config, and epoch together to checkpoints/*.pt so a model can be reconstructed and reloaded without re-deriving hyperparameters.
Next steps: write a proper predict() function with top-k sampling to concretely evaluate how the small local models are performing → build a Colab-only notebook for production-scale runs (needs GPU beyond the local MacBook) → run it for a first production-close model → then split-test RoPE, an alternative activation function in place of ReLU, grouped query attention, and weight tying.