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}
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.
The system integrates a JSON-driven preprocessing pipeline (BBPE
tokenization, multi-strategy packing, HDF5 and memory-mapped storage),
a 24-layer decoder-only architecture with Grouped Query Attention and
SwiGLU, and distributed training via DDP/FSDP with cosine scheduling.
A central focus is the numerical stability of BF16-precision training
in deep Transformers. Through variance propagation analysis, we show
that GPT-2 residual scaling ($\sigma_o = 0.02 / \sqrt{2L}$) on output
projections reduces per-block residual variance by a factor of $2L=48$,
containing post-24-layer variance at $1.34$ compared to $17.5$ without
scaling. Empirical evaluations over 15B training tokens demonstrate that
residual scaling consistently outperforms Kaiming initialization, with
the gap widening to $0.79$ in the mid-training regime before narrowing
to $0.38$ at convergence. These results establish residual scaling as a
practical necessity for BF16 Transformer training at scale.
\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.

\subsection{SFT Data Cleaning}

For supervised fine-tuning (SFT), raw data requires additional curation
beyond pretraining tokenization. {\sc Alembic}~\cite{alembic} is a companion
pipeline that handles SFT data generation, cleaning, and quality scoring:
three-generation strategies (topic-driven, seed-driven, self-instruct),
built-in cleaning (HTML/URL/markdown removal, char/word repetition filters),
and a MinHash-based near-duplicate detection system~\cite{broder1997syntactic}.
Given a set of $P$ hash functions ($P=128$) and a text $T$, the MinHash
pipeline proceeds as follows:

\begin{enumerate}[nosep,leftmargin=*]
  \item \textbf{Tokenization}: $T$ is split into character $n$-grams
    ($n=3$):
    \begin{equation}
      \Gamma(T) = \{\,c_i c_{i+1} c_{i+2} \mid i = 1,\dots,|T|-2 \,\}.
    \end{equation}
  \item \textbf{Signature}: For each hash function $h_k$, the minimum hash
    value over all $n$-grams forms the $k$-th element of the fingerprint:
    \begin{equation}
      s_k = \min_{t \in \Gamma(T)} h_k(t), \qquad
      h_k(t) = \operatorname{SHA256}(42 : k : t)_{[0:63]}.
    \end{equation}
    The full fingerprint is $\mathbf{s} = (s_1,\dots,s_P)$.
  \item \textbf{Similarity}: The Jaccard similarity between two sets is
    estimated by the fraction of agreeing fingerprint positions:
    \begin{equation}
      \widehat{J}(\mathbf{s}^{(a)},\mathbf{s}^{(b)}) =
      \frac{|\{\,k \mid s^{(a)}_k = s^{(b)}_k \,\}|}{P}.
    \end{equation}
  \item \textbf{Filtering}: Samples are processed sequentially; a sample
    is dropped if $\widehat{J}(\mathbf{s}, \mathbf{s}') \ge 0.7$ for any
    previously kept sample $\mathbf{s}'$.
\end{enumerate}

An optional LLM-as-Judge scoring module provides multi-dimensional
quality scores that can be used to filter low-quality samples.

\subsection{IFD-Based Instruction Difficulty Analysis}

Instruction Fulfillment Difficulty (IFD)~\cite{li2023ifd} quantifies
how challenging an instruction is for a model by comparing conditional
and unconditional losses:
\begin{equation}
\mathrm{IFD} = \frac{L_{\text{cond}}}{L_{\text{uncond}}}.
\end{equation}
An IFD $>1$ indicates the instruction increases the loss relative to
unconditional generation (the model struggles to follow it), while
IFD $<1$ means the instruction provides useful guidance.

We compute IFD for $N=3000$ SFT samples using both the pretrained
base model (after 15B tokens of pretraining) and a supervised
fine-tuned checkpoint (after 1K SFT steps).
Figure~\ref{fig:ifd} shows the distribution.

\begin{figure}[H]
\centering
\includegraphics[width=0.80\linewidth]{data/ifd_compare_clean.png}
\caption{IFD comparison: base model vs.\ trained checkpoint. The
diagonal line marks $\mathrm{IFD}_{\text{base}} = \mathrm{IFD}_{\text{ckpt}}$.}
\label{fig:ifd}
\end{figure}

The pretrained base model (15B tokens) has mean IFD $0.9625$;
$29.8\%$ of samples exceed $1.0$. After 1K SFT steps, mean IFD drops
to $0.7539$, with only $0.4\%$ of samples above $1.0$. The average
per-sample IFD reduction is $0.2086$. Conditional loss drops
$5.3\times$ more than unconditional loss, confirming that SFT teaches
instruction following rather than merely improving generic language
modeling. Detailed analysis is provided in Appendix~\ref{app:ifd}.

% ======================================================================
\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) and global L2 gradient clipping. The framework supports DDP and FSDP for multi-GPU distribution,
with gradient accumulation 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=1.5\times10^{-4}$ \\
Betas & $(0.9, 0.95)$, weight decay $0.1$ \\
Gradient clip & Global L2, max norm $1.0$ \\
Scheduler & Cosine, warmup ratio $0.02$ \\
Batch size & 4 per device $\times$ 4 GPUs $\times$ 32 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}.

% ======================================================================
\appendix
% ======================================================================

% ======================================================================
\section{IFD Data Examples}
\label{app:ifd}
% ======================================================================

Table~\ref{tab:ifd_examples} lists representative samples from the
IFD evaluation set, covering high, medium, and low IFD values
for the base model.

\begin{table}[H]
\centering
\caption{Representative IFD samples (base model sorted by descending IFD).}
\label{tab:ifd_examples}
\small
\begin{tabular}{@{}c c c c c c p{4.5cm}@{}}
\toprule
\textbf{Idx} &
  \textbf{Base IFD} &
  \textbf{Ckpt IFD} &
  \textbf{$L_{\text{cond}}^{\text{base}}$} &
  \textbf{$L_{\text{uncond}}^{\text{base}}$} &
  \textbf{$L_{\text{cond}}^{\text{ckpt}}$} &
  \textbf{Instruction (truncated)} \\
\midrule
0 & 4.605 & 1.525 & 12.38 & 2.69 & 3.77 & Complete the following analogy \dots \\
1 & 4.331 & 0.645 & 11.44 & 2.64 & 1.66 & Classify the following text \dots \\
2 & 3.741 & 0.702 & 11.75 & 3.14 & 2.17 & Label the following news article \dots \\
3 & 0.977 & 0.904 & 2.57 & 2.63 & 2.24 & Describe the role of a project manager \\
4 & 0.977 & 0.915 & 2.19 & 2.25 & 1.98 & Select a historical figure \dots \\
5 & 0.977 & 0.949 & 2.57 & 2.63 & 2.26 & Write the lyrics for an upbeat song \dots \\
6 & 0.977 & 0.925 & 2.94 & 3.00 & 2.62 & Explain how neural networks \dots \\
7 & 0.370 & 0.249 & 1.37 & 3.70 & 0.85 & Convert the given paragraph to a list \\
8 & 0.338 & 0.197 & 0.98 & 2.91 & 0.55 & Insert a suitable greeting \dots \\
9 & 0.307 & 0.062 & 0.70 & 2.29 & 0.15 & Remove third-person words \dots \\
\bottomrule
\end{tabular}
\end{table}

\subsection{Quantitative Summary}

Over $N=3000$ SFT samples:
\begin{itemize}[nosep]
  \item \textbf{Pretrained base model (15B tokens)}: mean IFD $= 0.9625$,
        median $= 0.9773$, std $= 0.1925$; $29.8\%$ of samples have
        IFD $> 1.0$.
  \item \textbf{SFT checkpoint (1K steps)}: mean IFD $= 0.7539$,
        median $= 0.8547$, std $= 0.2352$; only $0.4\%$ of samples
        exceed $1.0$.
  \item \textbf{Average IFD reduction}: $0.2086$ per sample.
  \item \textbf{Loss decomposition}: conditional loss drops by $0.9657$
        ($3.2424 \rightarrow 2.2767$), while unconditional loss drops by
        only $0.1838$ ($3.4142 \rightarrow 3.2303$). The $5.3\times$
        larger conditional reduction confirms the model primarily learns
        instruction following.
  \item \textbf{Correlation}: Pearson $r = 0.38$ between base and
        checkpoint IFD, indicating a moderate tendency for relatively
        hard instructions to remain relatively hard after training.
\end{itemize}

\subsection{Observed Patterns}

\paragraph{High-IFD samples (base IFD $> 3$, e.g.,~rows~0--2).}
These are tasks requiring task-intent comprehension: analogy completion,
text classification, article labeling. In the base model (15B pretraining), conditional
loss is extremely high ($L_{\text{cond}} \approx 12$), meaning the
instruction still acts as noise. After 1K SFT steps, IFD drops
sharply (e.g., $4.605 \rightarrow 1.525$), demonstrating
that SFT teaches the model to interpret and follow abstract task
descriptions.

\paragraph{Low-IFD samples (base IFD $< 0.4$, e.g.,~rows~7--9).}
These are formatting or extraction tasks: ``Convert paragraph to list,''
``Remove third-person words,'' ``Insert a greeting.'' Unconditional
loss is much higher than conditional loss even in the base model,
because the instruction naturally constrains the output space. The
pattern persists after SFT but with lower absolute values.

\paragraph{Mid-range samples (base IFD $\approx 0.98$, e.g.,~rows~3--6).}
These cover factual Q\&A and generation tasks: ``Describe the role of
a project manager,'' ``Write lyrics for a song,'' ``Explain how neural
networks work.'' In the base model IFD $\approx 1$ (instruction has
little effect); after SFT IFD drops to $\approx 0.9$, driven by
a clear reduction in conditional loss.

\paragraph{Cross-model correlation.}
The moderate Pearson correlation ($r = 0.38$) suggests that while
training reshapes the model's perception of instruction difficulty,
a residual signal persists: instructions that require complex reasoning
tend to remain non-trivially harder than simple rewrite or extraction
tasks even after SFT.

% ======================================================================
\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{alembic}
Alembic Contributors. \textit{Alembic: A lightweight LLM-driven SFT data
generation, cleaning, and scoring pipeline.}
\url{https://github.com/ViperEkura/Alembic}, 2026.

\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{broder1997syntactic}
A.~Z.~Broder. On the resemblance and containment of documents.
\textit{SEQUENCES '97}, 1997.

\bibitem{li2023ifd}
M.~Li, Y.~Zhang, S.~Li, Z.~Li, Z.~Li, L.~Zhu.
From quantity to quality: Boosting LLM performance with self-guided data
selection for instruction tuning.
\textit{NeurIPS}, 2024.

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

\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}