Compare commits

...
15 Commits
Author SHA1 Message Date
ViperEkura e6be33aa53 doc: add SFT length-filter rationale — high PPL + high variance from A.3 IFD figure
- Add 15-token length floor to SFT samples, with explicit reference
  to Appendix A.3 / Figure 5 (ifd_length_grid)
- Short replies (<10 tokens) show both high per-token perplexity
  (L_uncond ~6-8, PPL ~400-3000 vs long replies ~2-3, PPL ~7-20)
  and wide variance in L_cond (span 0-17.5), which distorts
  downstream IFD-based difficulty estimates.
2026-07-24 06:37:35 +08:00
ViperEkura 0a1d0573ae doc: add SFT length filter (15-token floor) with IFD bias rationale
Per-field length filtering in SFT drops instruction--response pairs
with responses shorter than 15 tokens. Short replies exhibit high
per-token variance in both conditional and unconditional loss
(Appendix A.3 / Figure 5), which would distort downstream IFD-based
difficulty estimates.
2026-07-24 06:27:10 +08:00
ViperEkura 0c2bc916f2 fix: loss_compare caption token count (20B -> 5B)
The loss_compare.png x-axis only spans 0-5B tokens. The caption
incorrectly claimed ~20B tokens.
2026-07-24 06:21:52 +08:00
ViperEkura 4e70e827ff remove: drop ckpt_weight_density_per_run.png figure
The per-run weight density figure contained misleading legacy
iteration labels (500k/1M iter) that contradicted the paper's
stated token budgets (15B). Removing it avoids confusion; the
per-category density plot (ckpt_weight_density.png) and Table 7
remain as the primary evidence.
2026-07-24 06:20:08 +08:00
ViperEkura a7bbc7b29f fix: align SFT/DPO figures and text with actual training data
- SFT: 1,000 steps/WSD → ~3,800 steps/cosine; loss ~2.5→1.6 → ~2.1→1.5
- DPO: fix preference-loss narrative to match training-loss curve;
  correct initial grad-norm (~50 → ~200)
- Table 2: peak LR 1.5e-4 → 2.0e-4 (matches pt_metric.png)
- Clarify scheduling: pretraining=WSD, SFT+DPO=cosine
- Add disclaimer to ckpt_weight_density_per_run caption for legacy iter labels
2026-07-24 06:15:42 +08:00
ViperEkura 4b37f289c0 fix: attribute muon-25bt larger std to Muon characteristic, not training duration alone
The muon-25bt (Muon, 25B) vs kami/norm (AdamW, 15B) comparison confounds
optimizer choice with training duration. Reframe as "Muon produces larger
post-convergence weight variance" rather than "training duration drives
variance growth."
2026-07-20 11:59:09 +08:00
ViperEkura a83555f326 fix: align conclusion, abstract, introduction with body data and full paper scope
- 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
2026-07-20 11:51:48 +08:00
ViperEkura d93ff48320 docs: fix storage backends (2->3) and per-field length filter
- Storage Backends: add JsonlStore (lazy on-the-fly tokenization)
- Describe memory-efficiency tradeoff across three backends
- Fix length filter: per-text-field skip, not per-sample global
- Abstract: HDF5/mmap -> tiered storage backends
- Rename metric figures (singular) and update references
2026-07-20 11:27:39 +08:00
ViperEkura da0f536526 Update DPO iteration count: 1,500 -> 3,000 2026-07-20 03:59:55 +08:00
ViperEkura c775a2b3e0 Update DPO metrics figure 2026-07-20 03:59:04 +08:00
ViperEkura 3f0ff911a8 Fix lingering cosine->WSD in Conclusion; update AGENTS.md with checkpoint labels and scheduler rules 2026-07-19 15:30:49 +08:00
ViperEkura e149997200 Add DPO data generation & training; fix checkpoint labels; switch to WSD scheduler; update optimizer ablations in abstract/conclusion 2026-07-19 15:25:42 +08:00
ViperEkura 6c2e04a86f Add DPO training subsection with dpo_metrics figure and Rafailov et al. citation 2026-07-19 15:01:23 +08:00
ViperEkura 411354eeb1 Update sft_metrics.png 2026-07-18 00:52:52 +08:00
ViperEkura a0c39601a0 Add effective rank analysis, weight distribution evolution, Muon/SFT training figures
- Add SVD effective rank (ER@99%) and condition number analysis across
  three checkpoints (kami-15bt, norm-15bt, muon-25bt) in appendix
- Add post-training weight distribution analysis as §5.4, verifying
  residual scaling persists throughout training
- Add per-component weight std table in appendix showing Muon constrains
  early-stage drift vs Normal init
- Add Muon optimizer training dynamics figure (muon_pt.png)
- Add SFT training metrics figure (sft_metrics.png) on mixed CN-EN data
- Restructure: move effective rank to appendix, weight distribution to
  §5.4, delete standalone §6, tighten abstract to prose-only
- Update conclusion to reference new analyses
2026-07-17 23:05:21 +08:00
7 changed files with 354 additions and 60 deletions
Binary file not shown.

After

Width:  |  Height:  |  Size: 940 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 315 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 86 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 68 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 84 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 100 KiB

+354 -60
View File
@@ -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{}
@@ -29,17 +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
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 focused BF16 stability analysis shows that GPT-2 residual scaling
($\sigma = 0.02/\sqrt{2L}$) reduces per-block residual variance by a factor
of 48, containing post-24-layer variance at 1.34 versus 17.5 under standard
initialization. Empirically, this scaling yields a sustained loss advantage
over Kaiming initialization, with the gap peaking at $\Delta = 0.79$ in the
mid-training regime.
of a 1.2B-parameter Transformer on $\sim$25B tokens. The pipeline covers
JSON-driven BBPE preprocessing with multi-strategy packing, tiered
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.1 to $\sim$1.5 over $\sim$3{,}800~steps; DPO alignment
($\beta=0.1$, cosine schedule) on model-generated preference pairs
shows stable training-loss convergence 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}
% ======================================================================
@@ -51,8 +60,12 @@ 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, 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.
% ======================================================================
\section{Data Pipeline}
@@ -77,17 +90,30 @@ via a JSON specification that defines:
\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.
In text-mode sections, individual fields shorter than 50~chars or
longer than 2M~chars are skipped during tokenization.
\subsection{Storage Backends}
Two storage backends serve the DataLoader:
Three storage backends serve the DataLoader, trading off memory
footprint against access speed for datasets ranging from fine-tuning
scale to TB-level pretraining:
\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.
\item \textbf{H5Store}: HDF5-based, fully loaded into RAM at init
with \texttt{share\_memory\_()} for cross-worker sharing.
Fastest random access; requires the dataset to fit in memory.
\item \textbf{MmapStore}: Zero-copy memory-mapped \texttt{.bin} files
via \texttt{np.memmap}. Data stays on disk, managed by the OS
page cache; multiple workers share physical pages without
duplication. Suited for large pretraining corpora that exceed
RAM.
\item \textbf{JsonlStore}: Reads raw \texttt{.jsonl} directly. Lazy
mode (\texttt{processor=fn}) keeps only raw text records in
memory and defers per-sample tokenization to
\texttt{fetch\_record()}, avoiding a pre-tokenized copy
entirely---used by DPO/GRPO for on-the-fly training from
source files.
\end{itemize}
A resumable distributed sampler provides seed-based shuffle with
@@ -128,13 +154,44 @@ pipeline proceeds as follows:
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.
A length filter is applied to SFT samples based on the IFD
length-bias analysis in Appendix~\ref{sec:ifd_bias}
(Figure~\ref{fig:length_bias}): instruction--response pairs
whose response contains fewer than 15 tokens are discarded,
because short replies exhibit both high per-token perplexity
($L_{\text{uncond}} \approx 6\text{--}8$, PPL~$\approx 400\text{--}3000$)
and wide variance in both $L_{\text{cond}}$ and $L_{\text{uncond}}$,
which would distort downstream IFD-based difficulty estimates.
The threshold is applied per-field, analogous to the pretraining
filter described above.
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}
% ======================================================================
@@ -222,8 +279,8 @@ The model is trained on next-token cross-entropy loss:
\mathcal{L} = -\sum_{t=1}^{T} \log P(x_t \mid x_{<t}; \theta).
\end{equation}
Training uses a hybrid optimizer: Muon for 2D weight matrices and AdamW~\cite{loshchilov2019adamw} for 1D parameters (embeddings, biases, LayerNorm), with cosine learning rate
scheduling (2\% warmup) and global L2 gradient clipping. The framework supports DDP and FSDP for multi-GPU distribution,
Training uses a hybrid optimizer: Muon for 2D weight matrices and AdamW~\cite{loshchilov2019adamw} for 1D parameters (embeddings, biases, LayerNorm), with WSD (Warmup--Stable--Decay)
learning rate scheduling (2\% warmup) and global L2 gradient clipping. The framework supports DDP and FSDP for multi-GPU distribution,
with gradient accumulation to manage memory.
Table~\ref{tab:train_params} lists the key hyperparameters.
@@ -236,14 +293,16 @@ Table~\ref{tab:train_params} lists the key hyperparameters.
\textbf{Hyperparameter} & \textbf{Value} \\
\midrule
Precision & BF16 (weights + optimizer states) \\
Optimizer & AdamW, $\eta=1.5\times10^{-4}$ \\
Optimizer & Hybrid Muon/AdamW$^a$, $\eta=2.0\times10^{-4}$ \\
Betas & $(0.9, 0.95)$, weight decay $0.1$ \\
Gradient clip & Global L2, max norm $1.0$ \\
Scheduler & Cosine, warmup ratio $0.02$ \\
Scheduler & WSD (warmup 2\%, stable, decay) \\
Batch size & 4 per device $\times$ 4 GPUs $\times$ 32 accumulation \\
Sequence length & 2,048 tokens \\
Total steps & 19,000 \\
Total tokens & $\sim$25B ($\approx$23k steps) \\
\bottomrule
\multicolumn{2}{@{}l@{}}{\footnotesize $^a$Muon applied to 2D weight matrices (attention projections and FFN layers);}\\
\multicolumn{2}{@{}l@{}}{\footnotesize \phantom{$^a$}AdamW applied to 1D parameters (embeddings, biases, LayerNorm scales).}\\
\end{tabular}
\end{table}
@@ -251,7 +310,7 @@ Total steps & 19,000 \\
\centering
\includegraphics[width=0.50\linewidth]{data/loss_compare.png}
\caption{Training loss curves: GPT-2 residual scaling vs.~Kaiming
initialization over $\sim$20B tokens.}
initialization over $\sim$5B tokens.}
\label{fig:loss}
\end{figure}
@@ -264,6 +323,92 @@ initialization over $\sim$20B tokens.}
Figure~\ref{fig:ckpt_comparison} compares four configurations. The left panel shows training loss for Muon (Embedding Adam + 1D Adam), Muon (Embedding Muon + 1D Adam), Kaiming init, and Normal init; the center panel zooms in on the two current Muon variants; and the right panel shows gradient norms over optimizer steps. The older Kaiming and Normal initializations converge more slowly and plateau at higher loss. Between the current variants, using Adam for the embedding layer yields lower loss and more stable gradients than using Muon embeddings.
\begin{figure}[H]
\centering
\includegraphics[width=0.85\linewidth]{data/pt_metric.png}
\caption{Muon optimizer training dynamics: training loss and learning
rate schedule for the hybrid Muon/AdamW configuration across $\sim$25B
tokens. The Muon optimizer is applied to all 2D weight matrices
(attention projections and FFN layers), while AdamW handles 1D
parameters (embeddings, biases, and normalization scales). The WSD
schedule (warmup, stable phase, and final decay) is visible in the lower panel.}
\label{fig:pt_metric}
\end{figure}
Figure~\ref{fig:pt_metric} shows the extended training dynamics of the
Muon optimizer configuration over $\sim$25B tokens. The loss curve
exhibits the expected power-law decay in the early phase (0--5B tokens),
followed by a gradual plateau as the model approaches convergence on
the pretraining distribution. The WSD schedule holds the learning rate
constant during the long stable phase and then decays at the end of
training, ensuring stable weight updates in the final phase.
\begin{figure}[H]
\centering
\includegraphics[width=0.85\linewidth]{data/sft_metric.png}
\caption{SFT training metrics over $\sim$3{,}800 fine-tuning steps on a
mixed Chinese--English instruction dataset: training loss, learning
rate, and gradient norm. The smoothed loss decreases from $\sim$2.1 to
$\sim$1.5, dropping rapidly in the first 500 steps and then entering a
slower decay phase. The learning rate follows a cosine schedule with short linear warmup. Gradient norms stabilise after approximately 500 steps,
indicating that the fine-tuning process has reached a stable optimization
regime.}
\label{fig:sft_metric}
\end{figure}
Figure~\ref{fig:sft_metric} shows the supervised fine-tuning metrics
for the SFT checkpoint used in the IFD analysis
(Appendix~\ref{app:ifd}). The training loss on the mixed
Chinese--English instruction dataset decreases from $\sim$2.1 to
$\sim$1.5 over $\sim$3{,}800 steps, with the gradient norm converging to a
stable range after the warmup phase.
\subsection{Direct Preference Optimization}
\label{sec:dpo}
Following supervised fine-tuning, we align the model with pairwise
human preferences via Direct Preference Optimization
(DPO)~\cite{rafailov2023dpo}. DPO avoids explicit reward-model training
by optimizing the policy $\pi$ directly against a frozen reference
policy $\pi_{\text{ref}}$ (the SFT checkpoint). For a prompt $x$ with
preferred response $y_w$ and dispreferred response $y_l$, the loss is:
\begin{equation}
\mathcal{L}_{\text{DPO}} = -\log \sigma\!\left(
\beta \Bigl[
\log\tfrac{\pi(y_w\mid x)}{\pi_{\text{ref}}(y_w\mid x)}
-\log\tfrac{\pi(y_l\mid x)}{\pi_{\text{ref}}(y_l\mid x)}
\Bigr]\right),
\end{equation}
where $\beta$ controls the KL-divergence penalty against the reference.
We use $\beta=0.1$, a batch of 64 preference pairs per step, and the
same hybrid Muon/AdamW optimizer. The pretraining phase employs WSD (warmup--stable--decay) scheduling
to maintain a long high-learning-rate plateau. Both SFT and DPO
alignment use cosine schedules with short linear warmup. The shorter ~3{,}000-step
alignment run does not benefit from an extended stable phase; instead,
cosine decay lowers the learning rate steadily, which discourages
over-optimisation away from the reference distribution and matches the
observed convergence pattern in Figure~\ref{fig:dpo_metric}.
\begin{figure}[H]
\centering
\includegraphics[width=0.95\linewidth]{data/dpo_metric.png}
\caption{DPO training metrics over $\sim$3{,}000 alignment steps:
preference loss (raw and 100-step moving average), learning rate
schedule, and gradient norm.}
\label{fig:dpo_metric}
\end{figure}
Figure~\ref{fig:dpo_metric} summarises the DPO training dynamics.
The raw training loss (left panel) starts near $0.7$ and is visibly
noisy; the smoothed curve reveals a steady downward trend that reaches
$\sim$0.1--0.15 by step 3{,}000 without rebound, indicating stable
convergence. The learning-rate schedule (centre panel) peaks at
$5\times10^{-6}$ after a short linear warmup and then follows cosine
decay to a floor of $\sim$0.5\,$\times\,$10$^{-6}$. Gradient norms (right panel) start
near 200 with occasional spikes above 250, then gradually decline and
stabilise in the 35--50 range after step 1{,}000, indicating consistent
gradient magnitudes throughout alignment.
% ======================================================================
\section{Numerical Stability via Residual Scaling}
\label{sec:num-stability}
@@ -321,7 +466,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
@@ -362,6 +507,55 @@ 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}).
\subsection{Post-Training Weight Distribution}
To verify that the residual-scaling constraint persists throughout
training---not just at initialization---we compare weight value
distributions across three checkpoints:
\texttt{kami-15bt} (AdamW, GPT-2 residual scaling, 15B tokens), \texttt{norm-15bt} (AdamW, Normal
init, 15B tokens), and \texttt{muon-25bt} (Muon, 25B tokens).
Figure~\ref{fig:ckpt_weight_density} shows the kernel density
estimates. Training preserves the bimodal structure: residual-scaled
projections ($\mathbf{W}_o$, $\mathbf{W}_{\text{down}}$) remain
narrowly concentrated near zero, while non-scaled weights broaden
from their initial $\mathcal{N}(0,0.02)$ distribution. Two trends
are visible in the per-component weight std
(Table~\ref{tab:weight_std}, Appendix~\ref{app:weight_std}):
\begin{itemize}[nosep]
\item \textbf{Muon produces larger post-convergence weight
variance}: the \texttt{muon-25bt} checkpoint (Muon, 25B tokens)
exhibits the largest weight std ($\sim$0.024), exceeding both
15B AdamW checkpoints ($\sim$0.015--0.021), consistent with
Muon allowing wider parameter distributions after convergence.
\item \textbf{Residual scaling constrains early-stage drift}: at
equal training budget (15B tokens) and with the same AdamW
optimizer, the GPT-2-scaled \texttt{kami-15bt} shows
\emph{smaller} std ($\sim$0.015) than the Normal-init
\texttt{norm-15bt} ($\sim$0.021), indicating that residual
scaling itself limits weight drift beyond its initialization
effect.
\end{itemize}
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
\includegraphics[width=0.95\linewidth]{data/ckpt_weight_density.png}
\caption{Weight value density estimates across three checkpoints,
grouped by component. Residual-scaled projections
($\mathbf{W}_o$, $\mathbf{W}_{\text{down}}$) maintain narrow
distributions near zero, while non-scaled weights show broader
spread.}
\label{fig:ckpt_weight_density}
\end{figure}
% ======================================================================
\section{Conclusion}
% ======================================================================
@@ -369,11 +563,33 @@ identified in the theoretical analysis (Section~\ref{sec:num-stability}).
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 a hybrid
Muon/AdamW optimizer under DDP/FSDP 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}.
Muon/AdamW optimizer under DDP/FSDP executors, and WSD 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. Post-training weight distribution analysis confirms that
this scaling constraint persists throughout training, with residual-scaled
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.1 to $\sim$1.5 over $\sim$3{,}800~cosine-scheduled
steps. Subsequent DPO alignment on preference pairs ($\beta=0.1$, cosine
schedule) shows stable training-loss convergence without over-optimisation.
Optimizer ablations (Figure~\ref{fig:ckpt_comparison}) demonstrate that
the hybrid Muon/AdamW configuration consistently outperforms pure AdamW
with both Kaiming and Normal initializations, achieving lower training loss
and faster convergence. Within the Muon family, applying AdamW to the
embedding layer rather than Muon yields a further loss reduction and more
stable gradient norms, confirming that 1D parameters benefit from
second-moment adaptation while 2D weight matrices gain from Muon's
orthogonalisation step. An SVD-based effective rank analysis
(Appendix~\ref{app:eff_rank}) further reveals that the model
operates near its representational capacity. The complete framework and
model weights are available at \url{https://github.com/ViperEkura/AstrAI}.
% ======================================================================
\appendix
@@ -553,32 +769,105 @@ This stability justifies using the base-model IFD as a data
selection signal without re-evaluating after fine-tuning.
% ======================================================================
\section{Weight Distribution by Component}
\label{app:weight_dist}
\section{Effective Rank Analysis}
\label{app:eff_rank}
% ======================================================================
Figure~\ref{fig:weight_dist} shows the distribution of weight
magnitudes at initialization, grouped by component type. Embeddings
and non-residual-scaled projections (QKV, attention output, FFN
gate/up) follow $\mathcal{N}(0, 0.02)$, producing near-identical
bell curves centered at zero. The residual-scaled projections
(output projection $\mathbf{W}_o$ and FFN down-projection
$\mathbf{W}_{\text{down}}$) use $\sigma = 0.02 / \sqrt{2L} \approx 0.0029$,
visible as the narrow, sharply peaked distribution concentrated
near zero. This factor-48 variance reduction is the mechanism by
which GPT-2 residual scaling prevents BF16 underflow in deep
Transformers (Section~\ref{sec:num-stability}).
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} (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}}$,
we compute the effective rank at 99\% energy:
\begin{equation}
\text{ER@99\%} = \frac{1}{\min(m,n)} \min_k \left\{ k \;\middle|\; \frac{\sum_{i=1}^{k} \sigma_i^2}{\sum_{i=1}^{\min(m,n)} \sigma_i^2} \ge 0.99 \right\}.
\end{equation}
\begin{figure}[H]
Table~\ref{tab:eff_rank} summarizes the results. All three checkpoints
exhibit a high overall ER@99\% ($\sim$90\%), indicating that the 1.2B
model operates close to its representational capacity. Key findings:
\begin{itemize}[nosep]
\item Q/O projections show lower ER@99\% ($\sim$0.73--0.77) and
high condition numbers ($\kappa > 10^4$), consistent with the
low-rank structure of GQA (24 query heads sharing 4 KV heads).
\item K/V projections, FFN layers, and embeddings maintain high
ER@99\% ($\sim$0.96--0.98) and low condition numbers
($\kappa < 10$).
\item The overall ER@99\% varies by less than 0.01 across
checkpoints, indicating these properties are determined
primarily by architecture rather than optimizer or training
duration.
\end{itemize}
\begin{table}[H]
\centering
\includegraphics[width=0.85\linewidth]{data/weight_dist_by_component.png}
\caption{Weight distribution by component at initialization.
Each panel shows the histogram of weight values for a specific
module group (embedding, attention projections, FFN projections,
output projections). The narrow peaks correspond to the
residual-scaled $\mathbf{W}_o$ and $\mathbf{W}_{\text{down}}$
projections.}
\label{fig:weight_dist}
\end{figure}
\caption{SVD effective rank (ER@99\%) and mean condition number
($\kappa$) by component across three checkpoints.}
\label{tab:eff_rank}
\small
\begin{tabular}{@{}lcccccc@{}}
\toprule
& \multicolumn{3}{c}{\textbf{ER@99\%}}
& \multicolumn{3}{c}{\textbf{Cond.\ Number $\kappa$}} \\
\cmidrule(lr){2-4} \cmidrule(lr){5-7}
\textbf{Component} & \textbf{kami} & \textbf{norm} & \textbf{muon}
& \textbf{kami} & \textbf{norm} & \textbf{muon} \\
\midrule
attn.k\_proj & 0.967 & 0.971 & 0.960 & 8.4 & 5.4 & 5.7 \\
attn.o\_proj & 0.766 & 0.708 & 0.730 & 31{,}589 & 57{,}284 & 22{,}838 \\
attn.q\_proj & 0.754 & 0.764 & 0.756 & 46{,}644 & 32{,}072 & 14{,}472 \\
attn.v\_proj & 0.976 & 0.976 & 0.971 & 2.5 & 2.4 & 2.9 \\
embed\_tokens & 0.985 & 0.987 & 0.984 & 4.9 & 1.9 & 3.5 \\
lm\_head & 0.969 & 0.981 & 0.980 & 21.0 & 13.5 & 16.1 \\
mlp.down & 0.961 & 0.961 & 0.963 & 6.4 & 7.2 & 7.2 \\
mlp.gate & 0.966 & 0.967 & 0.965 & 5.7 & 5.3 & 5.7 \\
mlp.up & 0.967 & 0.968 & 0.965 & 4.9 & 5.9 & 4.9 \\
\midrule
\textbf{Overall ER@99\%} & \textbf{0.909} & \textbf{0.903} & \textbf{0.903} & & & \\
\bottomrule
\end{tabular}
\end{table}
% ======================================================================
\section{Per-Component Weight Statistics}
\label{app:weight_std}
% ======================================================================
Table~\ref{tab:weight_std} reports the weight standard deviation by
component for each checkpoint, supplementing the post-training weight
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; residual scaling constrains drift at equal
token count (15B). Muon produces larger post-convergence weight
variance than AdamW.
Residual-scaled projections ($\mathbf{W}_o$, $\mathbf{W}_{\text{down}}$)
remain bounded.}
\label{tab:weight_std}
\small
\begin{tabular}{@{}lccc@{}}
\toprule
\textbf{Component} & \textbf{kami-15bt} & \textbf{norm-15bt} & \textbf{muon-25bt} \\
& (AdamW, 15B) & (AdamW, 15B) & (Muon, 25B) \\
\midrule
attn.q\_proj & 0.0154 & 0.0207 & 0.0230 \\
attn.k\_proj & 0.0153 & 0.0206 & 0.0238 \\
attn.v\_proj & 0.0146 & 0.0202 & 0.0244 \\
attn.o\_proj$^*$ & 0.0148 & 0.0084 & 0.0177 \\
mlp.up & 0.0153 & 0.0204 & 0.0237 \\
mlp.gate & 0.0155 & 0.0204 & 0.0235 \\
mlp.down$^*$ & 0.0100 & 0.0089 & 0.0180 \\
embed\_tokens & 0.0205 & 0.0205 & 0.0239 \\
lm\_head & 0.0224 & 0.0257 & 0.0298 \\
\bottomrule
\end{tabular}
\\[2pt]
\footnotesize $^*$Residual-scaled projection ($\sigma_0 = 0.02/\sqrt{2L}$).
\end{table}
% ======================================================================
\begin{thebibliography}{99}
@@ -627,6 +916,11 @@ A.~Radford, J.~Wu, R.~Child, D.~Luan, D.~Amodei, I.~Sutskever.
Language models are unsupervised multitask learners.
\textit{OpenAI Blog}, 2019.
\bibitem{rafailov2023dpo}
R.~Rafailov, A.~Sharma, E.~Mitchell, C.~D.~Manning, S.~Ermon, C.~Finn.
Direct Preference Optimization: Your language model is secretly a reward model.
\textit{NeurIPS}, 2023.
\bibitem{shazeer2020glu}
N.~Shazeer. GLU variants improve Transformer.
\textit{arXiv:2002.05202}, 2020.