From a83555f326a12e3e5948b1a065c1ffb5ff72ad98 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Mon, 20 Jul 2026 11:51:48 +0800 Subject: [PATCH] =?UTF-8?q?=EF=BB=BFfix:=20align=20conclusion,=20abstract,?= =?UTF-8?q?=20introduction=20with=20body=20data=20and=20full=20paper=20sco?= =?UTF-8?q?pe?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Fix sigma=0.003 (init value) to consistently narrower than non-scaled (body 4.4, conclusion, abstract) - Add SFT/DPO to conclusion, abstract, and introduction - Fix checkpoint label: kami-15bt (Muon) -> (AdamW, GPT-2 residual scaling) in effective rank appendix - Fix table weight_std caption: Muon -> residual scaling constrains drift - Fix 4.2 math: 0.689/L -> 0.689/(2L) - Fix length filter: per-sample -> per-text-field skip --- main.tex | 78 +++++++++++++++++++++++++++++++------------------------- 1 file changed, 43 insertions(+), 35 deletions(-) diff --git a/main.tex b/main.tex index cc0cc9c..7d64a4c 100644 --- a/main.tex +++ b/main.tex @@ -29,29 +29,26 @@ \begin{abstract} We present {\sc AstrAI}, an open-source framework for end-to-end training -of a 1.2B-parameter Transformer on $\sim$20B tokens. The pipeline covers +of a 1.2B-parameter Transformer on $\sim$25B tokens. The pipeline covers JSON-driven BBPE preprocessing with multi-strategy packing, tiered -storage backends (in-memory HDF5, memory-mapped binary, and lazy JSONL), 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 -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. -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. +storage backends, and a companion SFT pipeline ({\sc Alembic}) with +MinHash deduplication. The 24-layer GQA-SwiGLU decoder is trained with a +hybrid Muon/AdamW optimizer and WSD scheduling under DDP/FSDP. +Supervised fine-tuning on deduplicated bilingual instructions reduces loss +from $\sim$2.5 to $\sim$1.6 over 1{,}000~steps; DPO alignment +($\beta=0.1$, cosine schedule) on model-generated preference pairs +converges stably without over-optimisation. A BF16 stability analysis +shows that GPT-2 residual scaling ($\sigma_0 = 0.02/\sqrt{2L}$) reduces +per-block activation variance by a factor of 48, and post-training weight +analysis across three checkpoints confirms that residual-scaled +projections remain consistently narrower than non-scaled weights across +optimizers, initializations, and training budgets. Optimizer ablations +demonstrate that the hybrid Muon/AdamW outperforms pure AdamW, with +2D weight matrices benefiting from Muon's orthogonalisation and 1D +parameters from AdamW's second-moment adaptation. An SVD effective-rank +analysis reveals near-capacity weight utilization, with Q/O projections +showing consistently lower effective rank than K/V projections under +grouped query attention. \end{abstract} % ====================================================================== @@ -63,8 +60,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. Beyond the -pipeline, we conduct systematic ablations on optimizers (hybrid Muon/AdamW +open-source framework for Transformer training and inference, from JSONL +ingestion through pretraining, supervised fine-tuning (SFT) on deduplicated +bilingual instructions, and direct preference optimization (DPO) alignment. +We also 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. @@ -463,7 +462,7 @@ $1/\sqrt{2L}$: \end{equation} 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 +$0.689/(2L) \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 @@ -538,10 +537,11 @@ are visible in the per-component weight std constraints. \end{itemize} -Critically, the residual-scaled projections remain bounded at -$\sigma \approx 0.003$ across all checkpoints regardless of training -duration, confirming that the $1/\sqrt{2L}$ scaling continues to -enforce its design constraint throughout training. +Critically, the residual-scaled projections maintain consistently +lower standard deviations than their non-scaled counterparts across +all three checkpoints (Table~\ref{tab:weight_std}), confirming that +the $1/\sqrt{2L}$ scaling continues to enforce its design constraint +throughout training. \begin{figure}[H] \centering @@ -574,8 +574,16 @@ 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. Post-training weight distribution analysis confirms that this scaling constraint persists throughout training, with residual-scaled -projections maintaining narrow distributions ($\sigma \approx 0.003$) -regardless of optimizer or training duration. +projections maintaining consistently lower standard deviations than +non-scaled weights across all checkpoints +(Table~\ref{tab:weight_std}). + +Supervised fine-tuning on deduplicated bilingual instructions (processed +by the companion {\sc Alembic} pipeline with MinHash deduplication) +reduces training loss from $\sim$2.5 to $\sim$1.6 over 1{,}000~WSD-scheduled +steps. Subsequent DPO alignment on preference pairs ($\beta=0.1$, cosine +schedule) converges stably without over-optimisation, reaching a minimum +preference loss near step~1{,}200. Optimizer ablations (Figure~\ref{fig:ckpt_comparison}) demonstrate that the hybrid Muon/AdamW configuration consistently outperforms pure AdamW @@ -773,7 +781,7 @@ selection signal without re-evaluating after fine-tuning. To assess how well the trained parameters utilize their allocated capacity, we perform an SVD-based effective rank analysis on three -checkpoints: \texttt{kami-15bt} (Muon, 15B tokens), \texttt{norm-15bt} +checkpoints: \texttt{kami-15bt} (AdamW, GPT-2 residual scaling, 15B tokens), \texttt{norm-15bt} (Normal init, 15B tokens), and \texttt{muon-25bt} (Muon, 25B tokens). For each 2D weight matrix $\mathbf{W} \in \mathbb{R}^{m\times n}$ with SVD $\mathbf{W} = \mathbf{U}\boldsymbol{\Sigma}\mathbf{V}^{\mkern-1mu\mathsf{T}}$, @@ -839,9 +847,9 @@ distribution analysis in Section~\ref{sec:num-stability}. \begin{table}[H] \centering -\caption{Weight std by component across checkpoints. Non-scaled -weights broaden with training duration; Muon constrains drift at -equal token count (15B) but is overtaken by longer training (25B). +\caption{Weight std by component across checkpoints. Non-scaled weights broaden with training duration; residual scaling +constrains drift at equal token count (15B) but is overtaken by +longer training (25B). Residual-scaled projections ($\mathbf{W}_o$, $\mathbf{W}_{\text{down}}$) remain bounded.} \label{tab:weight_std}