\documentclass[11pt,a4paper]{article} % ===== Packages ===== \usepackage[utf8]{inputenc} \usepackage[T1]{fontenc} \usepackage{newtxtext,newtxmath} \usepackage[margin=1in]{geometry} \usepackage{amsmath} \usepackage{booktabs} \usepackage{graphicx} \usepackage{hyperref} \usepackage{float} \usepackage{caption} \usepackage{enumitem} \usepackage{url} \usepackage{microtype} \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{} \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 on output projections reduces per-block residual variance by a factor of 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\_\allowbreak{}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_{ 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,~1).} These are tasks requiring task-intent comprehension: analogy completion and 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~5,~6).} These are formatting or extraction tasks: ``Convert paragraph to list,'' ``Remove third-person words.'' 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 with large drop (e.g.,~rows~2,~3).} These are factual QA or grammar correction tasks. Base IFD is $\approx 1.05$ (instruction has little effect), but after SFT IFD drops to $\approx 0.1$ as the model learns the precise answer (e.g., ``Madrid'' for ``capital of Spain''), making conditional loss near-zero while unconditional loss remains high. \paragraph{Mid-range with small drop (e.g.,~row~4).} These are open-ended generation tasks (``Describe the role of a project manager''). Base IFD $\approx 0.98$; after SFT it drops only modestly to $\approx 0.9$, since both conditional and unconditional losses decrease proportionally without a memorized target. \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}