Tighten abstract, add Muon optimizer mention, update ckpt_comparison figure

This commit is contained in:
ViperEkura 2026-07-08 22:27:33 +08:00
parent 5726bf6c84
commit 4a143b056d
2 changed files with 12 additions and 16 deletions

Binary file not shown.

Before

Width:  |  Height:  |  Size: 243 KiB

After

Width:  |  Height:  |  Size: 273 KiB

View File

@ -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}
% ======================================================================