Tighten abstract, add Muon optimizer mention, update ckpt_comparison figure
This commit is contained in:
parent
5726bf6c84
commit
4a143b056d
Binary file not shown.
|
Before Width: | Height: | Size: 243 KiB After Width: | Height: | Size: 273 KiB |
28
main.tex
28
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}
|
||||
|
||||
% ======================================================================
|
||||
|
|
|
|||
Loading…
Reference in New Issue