diff --git a/data/ckpt_comparison.png b/data/ckpt_comparison.png index 3e2214a..6d138cc 100644 Binary files a/data/ckpt_comparison.png and b/data/ckpt_comparison.png differ diff --git a/main.tex b/main.tex index 4c7d91e..fa42dfc 100644 --- a/main.tex +++ b/main.tex @@ -28,22 +28,18 @@ \maketitle \begin{abstract} -Training billion-parameter language models requires careful co-design of -data infrastructure, distributed execution, and numerical precision -management. This paper presents {\sc AstrAI}, an open-source framework -for end-to-end training of a 1.2B-parameter autoregressive Transformer -over 15B tokens. We describe the full pipeline: JSON-driven preprocessing -with BBPE tokenization and multi-strategy packing, HDF5 and memory-mapped -storage backends, and a companion SFT pipeline ({\sc Alembic}) with -MinHash-based near-duplicate detection and LLM-as-Judge scoring. The -model is a 24-layer decoder-only Transformer with Grouped Query Attention, -SwiGLU, RoPE, and RMSNorm, trained with AdamW and cosine scheduling via -DDP/FSDP. A central focus is BF16 numerical stability: variance propagation -analysis shows that GPT-2 residual scaling reduces per-block residual -variance by a factor of 48, containing post-24-layer variance at 1.34 -compared to 17.5 without scaling. Empirical evaluations confirm that -residual scaling consistently outperforms Kaiming initialization, with the -loss gap peaking at 0.79 in the mid-training regime. +We present {\sc AstrAI}, an open-source framework for end-to-end training +of a 1.2B-parameter Transformer on 15B tokens. The pipeline covers +JSON-driven BBPE preprocessing with multi-strategy packing, HDF5/mmap +storage backends, and a companion SFT pipeline ({\sc Alembic}) with MinHash +deduplication and LLM-as-Judge scoring. The 24-layer decoder uses GQA, SwiGLU, +RoPE, and RMSNorm, trained with a hybrid Muon/AdamW optimizer and cosine scheduling under DDP/FSDP. +A focused BF16 stability analysis shows that GPT-2 residual scaling +($\sigma = 0.02/\sqrt{2L}$) reduces per-block residual variance by a factor +of 48, containing post-24-layer variance at 1.34 versus 17.5 under standard +initialization. Empirically, this scaling yields a sustained loss advantage +over Kaiming initialization, with the gap peaking at $\Delta = 0.79$ in the +mid-training regime. \end{abstract} % ======================================================================