AstrAI-paper/main.tex

\documentclass[11pt,a4paper]{article}

% ===== Packages =====
\usepackage[utf8]{inputenc}
\usepackage[T1]{fontenc}
\usepackage[margin=1in]{geometry}
\usepackage{amsmath,amssymb}
\usepackage{booktabs}
\usepackage{graphicx}
\usepackage{hyperref}
\usepackage{float}
\usepackage{caption}
\usepackage{enumitem}
\usepackage{url}

\DeclareMathOperator{\Var}{Var}

\title{End-to-End Training of a 1.2B Transformer with AstrAI \\
       \large Data Pipeline, Distributed Training, and BF16 Numerical Stability via Residual Scaling}

\author{AstrAI Contributors}
\date{June 2026}

\begin{document}

\maketitle

\begin{abstract}
We present an end-to-end framework for training a 1.2B-parameter autoregressive
language model using the {\sc AstrAI} open-source toolkit. The pipeline
encompasses data preprocessing (JSONL tokenization to BBPE, HDF5 and
memory-mapped storage), a 24-layer decoder-only architecture with Grouped
Query Attention and SwiGLU feed-forward blocks, and distributed training
via DDP/FSDP with cosine scheduling. We further examine BF16 numerical
stability in deep Transformers, demonstrating that GPT-2 residual scaling
($\sigma_o = 0.02 / \sqrt{2L}$) on output projections reduces per-block
residual variance by a factor of 48, yielding a post-24-layer variance of
$1.34$ versus $17.5$ without scaling. Empirical results across 15B training
tokens confirm that residual scaling maintains superior loss reduction over
Kaiming initialization, with a widening gap in the mid-training regime.
\end{abstract}

% ======================================================================
\section{Introduction}
% ======================================================================

Training a billion-parameter language model end-to-end involves far more than
model architecture. Data must be preprocessed and stored efficiently, the
training loop must handle distributed parallelism, gradient accumulation,
checkpointing, and logging---and numerical pitfalls must be diagnosed and
fixed. This paper describes the complete workflow using {\sc AstrAI}~\cite{astrai}, an
open-source framework for Transformer training and inference, and highlights a
BF16 precision issue encountered along the way.

% ======================================================================
\section{Data Pipeline}
% ======================================================================

\subsection{Preprocessing}

Raw data arrives as JSONL files. The preprocessing pipeline is configured
via a JSON specification that defines:

\begin{itemize}[nosep]
  \item \textbf{Tokenization}: BBPE tokenizer (100K vocabulary) with standard
        special tokens.
  \item \textbf{Masking}: Declarative loss mask assignment per section
        (e.g.,~mask user input, compute loss on assistant response).
  \item \textbf{Packing}: Documents concatenated via \texttt{simple}
        (sequential), \texttt{bfd} (best-fit decreasing), or
        \texttt{bfd\_split} strategies.
  \item \textbf{Position IDs}: \texttt{none}, \texttt{doc\_reset} (per-document
        boundary), or \texttt{continuous}.
  \item \textbf{Output}: Tokenized sequences written to \texttt{.h5} or
        \texttt{.bin} shards, auto-split at 100M tokens per shard.
\end{itemize}

Samples shorter than 50~chars or longer than 2M~chars are filtered out.

\subsection{Storage Backends}

Two storage backends serve the DataLoader:

\begin{itemize}[nosep]
  \item \textbf{H5Store}: HDF5-based, memory-loaded with shared-memory support
        for multi-worker access.
  \item \textbf{MmapStore}: Zero-copy memory-mapped \texttt{.bin} files shared
        via OS page cache.
\end{itemize}

A resumable distributed sampler provides seed-based shuffle with
epoch/iteration resume.

% ======================================================================
\section{Model Architecture}
% ======================================================================

The model is a 24-layer decoder-only Transformer with Grouped Query Attention
(GQA)~\cite{ainslie2023gqa} and SwiGLU feed-forward blocks~\cite{shazeer2020glu},
with Rotary Position Embedding (RoPE)~\cite{su2024roformer}.

\begin{table}[H]
\centering
\caption{Model configuration. Total: $\sim$1.2B parameters.}
\label{tab:model_config}
\begin{tabular}{@{}lrlr@{}}
\toprule
\textbf{Parameter} & \textbf{Value} & \textbf{Parameter} & \textbf{Value} \\
\midrule
Vocabulary ($V$) & 100,000 & Hidden dim ($d$) & 1,536 \\
Layers ($L$) & 24 & FFN dim ($d_{\textit{ffn}}$) & 6,912 \\
Query heads & 24 & KV heads & 4 \\
Head dim & 64 & Max length & 2,048 \\
Norm & RMSNorm ($\epsilon=10^{-5}$) & RoPE $\theta$ & 10,000 \\
\bottomrule
\end{tabular}
\end{table}

Each decoder block $\ell$ computes:
\begin{equation}
\mathbf{h}_\ell = \mathbf{x}_\ell + \operatorname{GQA}\bigl(\operatorname{RMSNorm}(\mathbf{x}_\ell)\bigr), \qquad
\mathbf{x}_{\ell+1} = \mathbf{h}_\ell + \operatorname{MLP}\bigl(\operatorname{RMSNorm}(\mathbf{h}_\ell)\bigr),
\end{equation}
where $\operatorname{MLP}(\mathbf{x}) = \mathbf{W}_{\text{down}}(\mathbf{W}_{\text{up}}\mathbf{x} \odot \operatorname{SiLU}(\mathbf{W}_{\text{gate}}\mathbf{x}))$.

\subsection{Initialization}

Linear weights follow $\mathcal{N}(0, 0.02)$; embeddings follow
$\mathcal{N}(0, 0.02)$. The output projection $\mathbf{W}_o$ and FFN
down-projection $\mathbf{W}_{\text{down}}$ use residual-scaled
initialization~\cite{radford2019gpt2}:
\begin{equation}
\sigma_o = \sigma_{\text{down}} = 0.02 / \sqrt{2L}.
\end{equation}
This scaling is critical for BF16 stability (Section~\ref{sec:num-stability}).

% ======================================================================
\section{Training Configuration}
% ======================================================================

The model is trained on next-token cross-entropy loss:
\begin{equation}
\mathcal{L} = -\sum_{t=1}^{T} \log P(x_t \mid x_{<t}; \theta).
\end{equation}

Training uses AdamW~\cite{loshchilov2019adamw} with cosine learning rate
scheduling (5\% warmup), global L2 gradient clipping, and periodic
validation. The framework supports DDP and FSDP for multi-GPU distribution,
with gradient accumulation and activation checkpointing to manage memory.
Table~\ref{tab:train_params} lists the key hyperparameters.

\begin{table}[H]
\centering
\caption{Training hyperparameters for the 1.2B run.}
\label{tab:train_params}
\begin{tabular}{@{}lr@{}}
\toprule
\textbf{Hyperparameter} & \textbf{Value} \\
\midrule
Precision & BF16 (weights + AdamW states) \\
Optimizer & AdamW, $\eta=7.5\times10^{-6}$ \\
Betas & $(0.9, 0.95)$, weight decay $0.01$ \\
Gradient clip & Global L2, max norm $1.0$ \\
Scheduler & Cosine, warmup ratio $0.05$ \\
Batch size & 4 per device $\times$ 4 GPUs $\times$ 8 accumulation \\
Sequence length & 2,048 tokens \\
Total steps & 950,000 \\
\bottomrule
\end{tabular}
\end{table}

% ======================================================================
\section{Numerical Stability via Residual Scaling}
\label{sec:num-stability}
% ======================================================================

Deep Transformers trained in BF16 face numerical stability challenges from
residual variance accumulation across layers. We evaluate the GPT-2
residual-scaling initialization~\cite{radford2019gpt2} as a mitigation
strategy.

\subsection{Variance Analysis}

At initialization with $\mathcal{N}(0, 0.02)$, a linear projection output has:
\begin{equation}
\Var(\mathbf{W}\mathbf{x}) = d_{\text{in}} \cdot (0.02)^2 \cdot \Var(\mathbf{x})
= 0.6144 \cdot \Var(\mathbf{x}) \quad (\text{for } d=1536).
\end{equation}

Within one block, attention and FFN each add a residual term. The variances
at each sub-stage are:

\begin{center}
\begin{tabular}{@{}lcc@{}}
\toprule
\textbf{Component} & \textbf{Operation} & $\Var$ (scaled by $\Var(\mathbf{x})$) \\
\midrule
Q/K/V proj & Linear(1536, $n_{\text{heads}}\cdot64$) & 0.6144 \\
Attention out & SDPA + $\mathbf{W}_o$ (scaled) & $0.378 / L$ \\
Gate/Up proj & Linear(1536, 6912) & 0.6144 \\
SiLU gate & $\operatorname{SiLU}(z) \approx 0.5z$ & $0.6144 \times 0.298 = 0.1831$ \\
Gated product & element-wise $\odot$ & $\approx 0.6144 \times 0.1831 = 0.1125$ \\
Down proj & Linear(6912, 1536) (scaled) & $0.311 / L$ \\
\midrule
Per-block residual & $\mathbf{R}_\ell = R_{\text{attn}} + R_{\text{ffn}}$ & $0.689 / L$ (scaled) \\
\bottomrule
\end{tabular}
\end{center}

Without the $1/\sqrt{2L}$ factor on $\mathbf{W}_o$ and
$\mathbf{W}_{\text{down}}$, the per-block residual variance becomes $0.689$
instead of $0.689/L \approx 0.014$. After $L=24$ blocks:
\begin{equation}
\begin{aligned}
\text{Without scaling: } \Var(\mathbf{x}_{24}) &\approx 1 + 24 \times 0.689 = 17.5,\\
\text{With scaling: } \Var(\mathbf{x}_{24}) &\approx 1 + 24 \times 0.014 = 1.34.
\end{aligned}
\end{equation}

\subsection{GPT-2 Residual Scaling}

The GPT-2 initialization~\cite{radford2019gpt2} scales output projections by
$1/\sqrt{2L}$:
\begin{equation}
\sigma_o = \sigma_{\text{down}} = 0.02 / \sqrt{2L}.
\end{equation}

This reduces per-block residual variance contribution from $0.689$ to
$0.689/L \approx 0.014$, a factor of $2L = 48$. The post-24-block variance
drops from $17.5$ to $1.34$, a $13.1\times$ improvement. In BF16
($7$-bit mantissa, ULP $= 0.0078$ at $w = 1.0$)~\cite{ieee754},
this keeps weight magnitudes within stable precision bounds. We further
recommend storing AdamW moments in FP32 and logging per-layer gradient
histograms during early training.

\subsection{Empirical Training Results}

\begin{figure}[H]
\centering
\includegraphics[width=0.50\linewidth]{data/loss_compare.png}
\caption{Training loss curves: GPT-2 residual scaling vs.~Kaiming
initialization over 15B tokens.}
\label{fig:loss}
\end{figure}

Figure~\ref{fig:loss} shows both loss curves; GPT-2 residual scaling (lower
curve) maintains a clear advantage, particularly in the 0.3--0.8B token region.

\begin{table}[H]
\centering
\caption{Loss at 0.125B-interval milestones, 0--1B tokens.}
\label{tab:loss_milestones}
\begin{tabular}{@{}lccc@{}}
\toprule
\textbf{Tokens (B)} &
  \textbf{GPT-2 scaling} &
  \textbf{Kaiming init} &
  \textbf{$\Delta$} \\
\midrule
0.125 & 7.37 & 7.66 & 0.29 \\
0.250 & 5.80 & 6.14 & 0.34 \\
0.375 & 4.82 & 5.38 & 0.56 \\
0.500 & 4.06 & 4.80 & 0.74 \\
0.625 & 3.50 & 4.29 & 0.79 \\
0.750 & 3.24 & 3.80 & 0.56 \\
0.875 & 3.21 & 3.43 & 0.22 \\
1.000 & 2.80 & 3.18 & 0.38 \\
\bottomrule
\end{tabular}
\end{table}

Table~\ref{tab:loss_milestones} quantifies the per-milestone gap. GPT-2
residual scaling leads at every interval, with $\Delta$ growing from 0.29
at 0.125B to a peak of 0.79 at 0.625B, then narrowing to 0.38 at 1B.
The widening mid-range gap aligns with the variance accumulation region
identified in the theoretical analysis (Section~\ref{sec:num-stability}).

% ======================================================================
\section{Conclusion}
% ======================================================================

We have described the end-to-end pipeline for training a 1.2B Transformer with
{\sc AstrAI}: data preprocessing with JSON-driven tokenization and packing,
a 24-layer GQA-SwiGLU architecture, callback-based training with DDP/FSDP
executors, and cosine scheduling. We further analyzed numerical stability
under BF16, showing that GPT-2 residual scaling ($\sigma_o = 0.02/\sqrt{2L}$)
reduces per-block residual variance by a factor of 48, keeping post-24-layer
variance at $1.34$ versus $17.5$ without scaling. The complete framework and model
weights are available at \url{https://github.com/ViperEkura/AstrAI}.

% ======================================================================
\begin{thebibliography}{99}

\bibitem{ainslie2023gqa}
J.~Ainslie, J.~Lee-Thorp, M.~de Jong, Y.~Zemlyanskiy, F.~Lebr\'on, S.~Sanghai.
GQA: Training generalized multi-query transformer models from multi-head
checkpoints. \textit{EMNLP}, 2023.

\bibitem{astrai}
AstrAI Contributors. \textit{AstrAI: An open-source training and inference
framework for Transformer language models.}
\url{https://github.com/ViperEkura/AstrAI}, 2026.

\bibitem{ieee754}
IEEE Computer Society. \textit{IEEE Standard for Floating-Point Arithmetic},
IEEE Std 754-2008, 2008.

\bibitem{loshchilov2019adamw}
I.~Loshchilov, F.~Hutter. Decoupled weight decay regularization.
\textit{ICLR}, 2019.

\bibitem{radford2019gpt2}
A.~Radford, J.~Wu, R.~Child, D.~Luan, D.~Amodei, I.~Sutskever.
Language models are unsupervised multitask learners.
\textit{OpenAI Blog}, 2019.

\bibitem{shazeer2020glu}
N.~Shazeer. GLU variants improve Transformer.
\textit{arXiv:2002.05202}, 2020.

\bibitem{su2024roformer}
J.~Su, A.~Murtadha, Y.~Lu, S.~Pan, B.~Wen, Y.~Liu.
Roformer: Enhanced transformer with rotary position embedding.
\textit{Neurocomputing}, 568:127063, 2024.

\end{thebibliography}

\end{document}