diff --git a/.gitignore b/.gitignore index 0afa7a9..f49d932 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,8 @@ # Whitelist: ignore everything * +!*/ # Only track .tex files and .gitignore itself !.gitignore !*.tex +!data/** diff --git a/data/loss_compare.png b/data/loss_compare.png new file mode 100644 index 0000000..086c776 Binary files /dev/null and b/data/loss_compare.png differ diff --git a/main.tex b/main.tex index 5408e5a..dce904c 100644 --- a/main.tex +++ b/main.tex @@ -16,7 +16,7 @@ \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} + \large Data Pipeline, Distributed Training, and BF16 Numerical Stability via Residual Scaling} \author{AstrAI Contributors} \date{June 2026} @@ -26,18 +26,18 @@ \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 +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, preventing the lock-in. +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} % ====================================================================== @@ -132,7 +132,7 @@ 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}). +This scaling is critical for BF16 stability (Section~\ref{sec:num-stability}). % ====================================================================== \section{Training Configuration} @@ -170,44 +170,18 @@ Total steps & 950,000 \\ \end{table} % ====================================================================== -\section{Case Study: BF16 Weight Lock-in} -\label{sec:bf16} +\section{Numerical Stability via Residual Scaling} +\label{sec:num-stability} % ====================================================================== -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. +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{Observation} +\subsection{Variance Analysis} -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: +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). @@ -243,35 +217,63 @@ instead of $0.689/L \approx 0.014$. After $L=24$ blocks: \end{aligned} \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$. +\subsection{GPT-2 Residual Scaling} -\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. +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} -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. +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{Solution} +\subsection{Empirical Training Results} -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. +\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} @@ -280,10 +282,10 @@ than BF16, and logging per-layer gradient histograms during early training. 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 +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}. % ======================================================================