From 6f9278b35b17c2ede049b6ee25af2a31ae38e797 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Thu, 25 Jun 2026 15:51:05 +0800
Subject: [PATCH] Initial commit: BF16 weight lock-in paper for 1.2B
 Transformer training with AstrAI

---
 .gitignore |   6 +
 main.tex   | 323 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 329 insertions(+)
 create mode 100644 .gitignore
 create mode 100644 main.tex

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..0afa7a9
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,6 @@
+# Whitelist: ignore everything
+*
+
+# Only track .tex files and .gitignore itself
+!.gitignore
+!*.tex
diff --git a/main.tex b/main.tex
new file mode 100644
index 0000000..9bffb37
--- /dev/null
+++ b/main.tex
@@ -0,0 +1,323 @@
+\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}