diff --git a/data/weight_dist_by_component.png b/data/weight_dist_by_component.png deleted file mode 100644 index fc9c117..0000000 Binary files a/data/weight_dist_by_component.png and /dev/null differ diff --git a/main.tex b/main.tex index 64c3b27..b66a419 100644 --- a/main.tex +++ b/main.tex @@ -18,7 +18,7 @@ \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} + \large Data Pipeline, Distributed Training, and Ablations on Optimizer, Initialization, and BF16 Numerical Stability} \author{AstrAI Contributors} \date{} @@ -34,21 +34,24 @@ JSON-driven BBPE preprocessing with multi-strategy packing, HDF5/mmap storage backends, and a companion SFT pipeline ({\sc Alembic}) with MinHash deduplication and LLM-as-Judge scoring. The 24-layer decoder uses GQA, SwiGLU, RoPE, and RMSNorm, trained with a hybrid Muon/AdamW optimizer and -cosine scheduling under DDP/FSDP. A BF16 stability analysis shows that +WSD (Warmup--Stable--Decay) scheduling under DDP/FSDP. A BF16 stability analysis shows that GPT-2 residual scaling substantially reduces per-block residual variance accumulation, keeping post-training variance well below the overflow threshold of standard initialization; empirically this yields a sustained loss advantage over Kaiming initialization throughout training. -Post-training weight distribution analysis across three -checkpoints---varying optimizer, initialization, and training -budget---confirms that residual-scaled projections maintain narrow -distributions throughout training, preserving the numerical stability -established at initialization. An SVD-based effective rank analysis -further reveals that the model operates near its representational -capacity, with attention Q/O projections consistently showing lower -utilization than K/V projections, a pattern stable across all -configurations and consistent with the low-rank structure induced by -grouped query attention. +Systematic ablations show that the hybrid Muon/AdamW optimizer +outperforms pure AdamW on 2D weight matrices, and that GPT-2 residual +scaling yields a sustained loss advantage over both Kaiming and Normal +initialization throughout training. Post-training weight distribution +analysis across three checkpoints---varying optimizer, initialization, +and training budget---confirms that residual-scaled projections +maintain narrow distributions throughout training, preserving the +numerical stability established at initialization. An SVD-based +effective rank analysis further reveals that the model operates near +its representational capacity, with attention Q/O projections +consistently showing lower utilization than K/V projections, a pattern +stable across all configurations and consistent with the low-rank +structure induced by grouped query attention. \end{abstract} % ====================================================================== @@ -60,8 +63,10 @@ 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. +open-source framework for Transformer training and inference. Beyond the +pipeline, we conduct systematic ablations on optimizers (hybrid Muon/AdamW +versus pure AdamW) and initializations (GPT-2 residual scaling, Kaiming, +and Normal), and analyse a BF16 precision issue encountered along the way. % ====================================================================== \section{Data Pipeline} @@ -140,10 +145,33 @@ pipeline proceeds as follows: An optional LLM-as-Judge scoring module provides multi-dimensional quality scores that can be used to filter low-quality samples. -An IFD (Instruction Fulfillment Difficulty) analysis is provided in -Appendix~\ref{app:ifd}. +\subsection{DPO Data Generation} -% ====================================================================== +To construct pairwise preference data for Direct Preference +Optimization, we start from a bilingual (Chinese--English) instruction +set curated by the same MinHash pipeline described above. For each +prompt $x$ in the instruction set, we generate two responses: + +\begin{itemize}[nosep] + \item \textbf{Chosen} $y_w$: generated by the reference model + $\pi_{\text{ref}}$, i.e.~the SFT checkpoint after supervised + fine-tuning on the curated instruction data. + \item \textbf{Rejected} $y_l$: generated by the base model + $\pi_{\text{base}}$, i.e.~the model at the end of pretraining + before any instruction tuning. +\end{itemize} + +Both generations use greedy decoding to eliminate sampling variance and +ensure that the preference signal reflects model capability rather than +decoding randomness. The resulting preference pairs +$(x, y_w, y_l)$ are stored in the same JSONL format used for SFT and +fed directly into the DPO training loop +(Section~\ref{sec:dpo}). Because the base model and the reference +model share the same architecture but differ in instruction-following +ability, the contrast between $y_w$ and $y_l$ is sharp and consistent, +which stabilises the DPO gradient updates. All instructions are +balanced across Chinese and English domains to preserve bilingual +alignment capability during preference optimisation. \section{Model Architecture} % ====================================================================== @@ -231,8 +259,8 @@ The model is trained on next-token cross-entropy loss: \mathcal{L} = -\sum_{t=1}^{T} \log P(x_t \mid x_{