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{enumitem}
\usepackage{url}

\DeclareMathOperator{\Var}{Var}

\title{End-to-End Training of a 1.2B Transformer with AstrAI \\
       \large Data Pipeline, Distributed Training, and a BF16 Numerical Case Study}

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

\begin{document}

\maketitle

\begin{abstract}
We document the end-to-end process of training a 1.2B-parameter autoregressive
language model from scratch using the {\sc AstrAI} open-source framework. The
pipeline covers data preprocessing (JSONL $\rightarrow$ BBPE tokenization
$\rightarrow$ HDF5/mmap storage), a 24-layer GQA-SwiGLU decoder-only
architecture, and distributed training with DDP/FSDP and cosine
scheduling. During training, we encountered
a BF16 numerical pathology: approximately 73,500 weights irreversibly locked at
$1.0$ within 100k steps. We analyze this as a three-stage cascade---variance
accumulation in unscaled residuals, gradient saturation at deep layers, and
BF16 precision loss blocking recovery---and show that residual scaling
($\sigma_o = 0.02 / \sqrt{2L}$) on output projections reduces per-block
residual variance by a factor of 48, preventing the lock-in.
\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.}
\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:bf16}).

% ======================================================================
\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{Case Study: BF16 Weight Lock-in}
\label{sec:bf16}
% ======================================================================

During the above training run, we encountered a numerical failure: at the first
checkpoint (iteration 100k), approximately 73,500 weights had locked at
exactly $1.0$ and 5,928 embedding dimensions had overflowed to
$\sim2\times10^{37}$. None recovered by iteration 950k.

\subsection{Observation}

Table~\ref{tab:locked} shows the affected parameters. The pattern is
structured: the same rows of $\mathbf{W}_k$ (rows 6--7) lock across layers
1--23; the same row of $\mathbf{W}_{\text{down}}$ (row~1) locks in all 24
layers; the overflow is confined to tokens 0--7 (BOS/EOS/PAD tokens).

\begin{table}[H]
\centering
\caption{Parameters locked at $w = 1.0$ or overflowed.}
\label{tab:locked}
\begin{tabular}{@{}lrr@{}}
\toprule
\textbf{Parameter} & \textbf{Count} & \textbf{Location} \\
\midrule
$\mathbf{W}_k$ (layers 1--23) & 35,328 & Rows 6--7 \\
$\mathbf{W}_{\text{down}}$ (layers 0--23) & 36,864 & Row 1, cols 3164--4699 \\
LM head & 1,536 & Rows 6--7 \\
Embedding overflow & 5,928 & Tokens 0--7 \\
\midrule
Total affected & 73,536 & 0.006\% of parameters \\
\bottomrule
\end{tabular}
\end{table}

\subsection{Root Cause: Three-Stage Cascade}

\textbf{Stage 1: Variance accumulation.} 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}
\text{Without scaling: } \Var(\mathbf{x}_{24}) \approx 1 + 24 \times 0.689 = 17.5,
\qquad
\text{With scaling: } \Var(\mathbf{x}_{24}) \approx 1 + 24 \times 0.014 = 1.34.
\end{equation}

\textbf{Stage 2: Saturation.} The 17.5$\times$ inflated variance at deep
layers saturates softmax (one-hot, gradient $\to 0$) and SiLU (derivative
$\to 0$ or 1). This creates a noisy gradient landscape where isolated
weights receive disproportionately large updates---a single batch can push a
weight from $\sim 0.06$ to $1.0$.

\textbf{Stage 3: BF16 dead zone.} BF16 has a 7-bit mantissa~\cite{ieee754}. The ULP at
$w = 1.0$ is $2^{-7} = 0.0078125$. The per-step AdamW update for a
locked weight is $1.2\times10^{-6}$ to $3.5\times10^{-8}$, which is at
least $1000\times$ smaller than the ULP. The weight cannot change.
The BF16 momentum buffers lose precision at large magnitudes, making
recovery impossible even if a corrective gradient arrives.

Global L2 clipping ($\ell_2 \leq 1.0$) fails here: a single-element
gradient spike of magnitude $10^4$ across $1.2\times10^9$ parameters
produces $\|\nabla\|_2 \approx \sqrt{10^8 + 1.2\times10^9} \approx 36056$.
After clipping ($\times 1/36056$), the spike remains at $0.28$---still
enough to jump a weight from $0.06$ to $1.0$---while all other gradients
are scaled to near-zero.

\subsection{Solution}

Residual scaling of output projections---setting
$\sigma_o = \sigma_{\text{down}} = 0.02 / \sqrt{2L}$
on $\mathbf{W}_o$ and $\mathbf{W}_{\text{down}}$---reduces the per-block
residual contribution by $1/(2L) = 1/48$, bringing post-24-block variance
from $17.5$ down to $1.34$ and preventing the saturation--explosion--lock-in
cascade. Additionally, we recommend storing AdamW moments in FP32 rather
than BF16, and logging per-layer gradient histograms during early training.

% ======================================================================
\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. A BF16 numerical failure encountered during
training was traced to variance accumulation in unscaled residuals, gradient
saturation, and BF16 precision loss---and resolved via the residual scaling
pattern already present in the codebase. 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, M.~Ahmed, Y.~Lu, S.~Pan, W.~Bo, Y.~Liu.
Roformer: Enhanced transformer with rotary position embedding.
\textit{Neurocomputing}, 568:127063, 2024.

\end{thebibliography}

\end{document}