diff --git a/data/dpo_metric.png b/data/dpo_metric.png index 4fd4b33..4c2e7c2 100644 Binary files a/data/dpo_metric.png and b/data/dpo_metric.png differ diff --git a/data/sft_metric.png b/data/sft_metric.png index 2e4d4c0..d05ee12 100644 Binary files a/data/sft_metric.png and b/data/sft_metric.png differ diff --git a/main.tex b/main.tex index 38cf0dc..9ee1f98 100644 --- a/main.tex +++ b/main.tex @@ -35,9 +35,9 @@ 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 +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 -converges stably without over-optimisation. A BF16 stability analysis +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 @@ -285,7 +285,7 @@ Table~\ref{tab:train_params} lists the key hyperparameters. \textbf{Hyperparameter} & \textbf{Value} \\ \midrule Precision & BF16 (weights + optimizer states) \\ - Optimizer & Hybrid Muon/AdamW$^a$, $\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 & WSD (warmup 2\%, stable, decay) \\ @@ -338,21 +338,21 @@ 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 1{,}000 fine-tuning steps on a +\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 loss drops rapidly in the first 200 -steps and then enters a slower decay phase. The learning rate -follows a WSD schedule (2\% warmup, stable, decay). Gradient norms stabilize -after approximately 300 steps, indicating that the fine-tuning -process has reached a stable optimization regime.} +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 1K-step SFT checkpoint used in the IFD analysis +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.5 to -$\sim$1.6 over 1{,}000 steps, with the gradient norm converging to a +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} @@ -373,10 +373,9 @@ preferred response $y_w$ and dispreferred response $y_l$, the loss is: \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. Unlike the pretraining and SFT -phases, which both employ WSD (warmup--stable--decay) scheduling to -maintain a long high-learning-rate plateau, DPO alignment uses a -cosine schedule with short linear warmup. The shorter ~3{,}000-step +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 @@ -392,17 +391,14 @@ schedule, and gradient norm.} \end{figure} Figure~\ref{fig:dpo_metric} summarises the DPO training dynamics. -The raw preference loss (left panel) starts near $0.69$ and is visibly -noisy, a hallmark of pairwise preference sampling. The 100-step moving -average reveals a steady downward trend that reaches a minimum of -$\sim$0.51 near step 1{,}200, after which it gently rebounds to -$\sim$0.52 and oscillates, suggesting the policy has converged to a -stable preference boundary rather than over-optimising away from the -reference distribution. The learning-rate schedule (centre panel) peaks -at $5\times10^{-6}$ after a short linear warmup and then follows cosine +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 50 with occasional spikes above 55, then gradually decline and -stabilise in the 35--40 range after step 400, indicating consistent +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. % ====================================================================== @@ -554,7 +550,12 @@ spread.} \begin{figure}[H] \centering \includegraphics[width=0.95\linewidth]{data/ckpt_weight_density_per_run.png} -\caption{Per-checkpoint weight density breakdowns.} +\caption{Per-checkpoint weight density breakdowns. The three checkpoints +are \texttt{norm-15bt} (AdamW, Normal init, 15B tokens), \texttt{kami-15bt} +(AdamW, GPT-2 residual scaling, 15B tokens), and \texttt{muon-25bt} +(Muon, 25B tokens). Note: the iteration labels inside the figure reflect +legacy script metadata and should be ignored in favour of the token +budgets stated here.} \label{fig:ckpt_weight_density_per_run} \end{figure} @@ -577,10 +578,9 @@ non-scaled weights across all checkpoints 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 +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) converges stably without over-optimisation, reaching a minimum -preference loss near step~1{,}200. +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