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
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user