PaperlayerView exact revision on arXiv

Source-linked AI summary

The Ephemeral Cost of Switching Optimizers Across LLM Training Stages

Xinyu Zhao, Davis Wertheimer, Ashutosh Sharma, Rameswar Panda, Paul Pu Liang

arXiv:2610.22650v1cs.LG

TL;DR

Prior work suggested optimizer switching might incur a cost, but it remained unclear whether that cost extends beyond forgetting to final performance after task-general SFT. The paper varies AdamW and Muon updates across training stages and evaluates open-weight models. It finds little evidence of a meaningful mismatch penalty, with Muon remaining competitive for SFT.

  • Problem

    It remained unclear whether optimizer switching causes a final-performance cost beyond the forgetting cost suggested for task-specific fine-tuning.

    Source: §1 · p. 1

  • Method

    The authors vary AdamW and Muon update proportions across pretraining, continued pretraining, and SFT, and evaluate task-general SFT on open-weight models.

    Source: §1 · p. 1, §3 · p. 3

  • Results

    There is little evidence of a meaningful mismatch penalty; Muon SFT generally improves downstream performance regardless of pretraining optimizer, while optimizer consistency reduces forgetting without benchmark gains.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Post-training optimizer choice can be made largely independently of the pretraining optimizer, and Muon is viable for AdamW-pretrained models.

    Source: §5 · p. 6

Abstract

from arXiv · show

Recent work suggests that switching optimizers between pretraining and subsequent training stages incurs a performance cost, but systematic investigation of this phenomenon is lacking. We study the mismatch between AdamW and Muon optimizers by pretraining 1.2B Llama-3 models from scratch while varying the proportion of updates taken with each optimizer, and similarly interpolating between the two optimizers during continued pretraining and supervised fine-tuning (SFT). Surprisingly, we find little evidence of a meaningful optimizer mismatch penalty: final performance is largely explained by the independent effects of the pretraining and post-training optimizers, with downstream performance generally improving as the proportion of Muon updates during SFT increases. Although matching the pretraining optimizer can reduce forgetting during SFT, this does not translate into improved final benchmark performance. We further evaluate AdamW and Muon for task-general SFT of open-weight models, across the Llama and Qwen families up to 14B parameters, where Muon remains competitive with AdamW, despite these models having been pretrained with AdamW. Together, our results suggest that the choice of optimizer for post-training can be made largely independently of the optimizer used during pretraining.

1 Introduction

The introduction asks whether switching AdamW and Muon between training stages imposes a final-performance cost beyond forgetting. It previews controlled optimizer-mixture experiments and reports little evidence of a meaningful mismatch penalty.

1 Introduction

  • The study tests whether optimizer switching reduces final performance, beyond the previously suggested forgetting cost during task-specific fine-tuning.The introduction notes that the cost for larger-scale, multi-task SFT remained unclear.Source: §1 · p. 1
  • The authors vary AdamW and Muon update proportions across pretraining, continued pretraining, and task-general SFT to separate mismatch from optimizer and data-shift effects.They also evaluate task-general SFT on released open-weight models up to 14B parameters.Source: §1 · p. 1
  • Across six downstream benchmarks, SFT optimizer interactions are weak and inconsistent, while forgetting residuals show a diagonal matched-optimizer preference.The figure uses best-learning-rate settings for benchmark averages and a fixed learning rate for forgetting residuals.Source: Figure 1 · p. 2
  • Optimizer consistency reduces forgetting during SFT, but does not improve final benchmark performance; increasing Muon updates during SFT generally improves performance regardless of pretraining optimizer.Across open-weight Llama and Qwen models up to 14B parameters, Muon remains competitive with AdamW for task-general SFT.Source: §1 · p. 2

2 Experimental Setup

The experimental setup compares AdamW, Muon, and mixtures of their updates across pretraining and post-training stages. It evaluates 1.2B Llama-3 models using specified model, data, and SFT configurations.

2 Experimental Setup

  • The 1.2B Llama-3 models use 32 layers, 24 full multi-head-attention heads, and hidden dimension 1536.AdamW uses decoupled weight decay and gradient clipping; Muon uses 0.2 RMS update scaling and the same base pretraining learning rate.Source: §2 · p. 2
  • Mixed optimizers select a Muon update independently at each step with probability p ∈ [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0], otherwise using AdamW.The expected AdamW update RMS is used to unify update scales and permit a common base learning rate across mixture settings.Source: §2 · p. 2
  • CPT residuals compare shared 5 × 10^-5 learning-rate settings with per-setting best learning rates; matched-optimizer preference vanishes after tuning.Rows vary pretraining Muon proportion, and columns vary CPT Muon proportion.Source: Figure 2 · p. 3
  • Pretraining uses the same data mixture at 1× and 8× Chinchilla scales, with approximately 24B and 193B tokens, respectively.The models are pretrained from scratch.Source: §2 · p. 3
  • SFT uses the Tulu-3 mixture for two epochs, selects learning rates on held-out validation loss, and evaluates six benchmarks.Sequences are packed to 4096 tokens for 1.2B and Llama-2-7B models, and 8192 tokens for Qwen models.Source: §2 · p. 3

3 Cost of Optimizer Mismatch

Across the optimizer mixtures, final benchmark performance showed little evidence of a meaningful mismatch penalty and was largely explained by each stage’s optimizer effects. Matching optimizers reduced forgetting, while learning-rate tuning removed apparent switching costs during continued pretraining.

  • The study interpolated between AdamW and Muon updates and decomposed each score-grid cell into a global mean, optimizer marginal effects, and a residual interaction.This separates mismatch effects from each optimizer’s independent contribution.Source: §3 · p. 3
  • Optimizer interaction effects were mild and lacked a consistent pattern, while final performance was dominated by the independent effects of the pretraining and SFT optimizers.The matched AdamW/AdamW residual was −0.2 and Muon/Muon was +0.2; diagonal interactions ranged from −0.6 to +0.3.Source: §3 · p. 4
  • More Muon updates during SFT generally improved downstream performance regardless of pretraining optimizer; models with some Muon pretraining also performed better on average, largely as a pretraining artifact.The latter models already achieved lower pretraining validation loss than AdamW-pretrained models.Source: §3 · p. 4
  • During continued pretraining, matched optimizers appeared favorable at a shared 5e-5 learning rate, but the interaction effects disappeared after learning rates were tuned for each setting.This suggests apparent switching costs may stem from hyperparameter mismatch rather than the optimizers themselves.Source: §3 · p. 4
  • During SFT, using more Muon updates generally increased forgetting of the pretraining distribution, while Muon pretraining and optimizer matching reduced it.The authors suggest switching costs may manifest as forgetting rather than reduced learning capability.Source: §3 · p. 4

4 Does model scale change the preferred SFT optimizer?

The paper tested whether Muon’s task-general SFT performance held at greater data scale and across AdamW-pretrained open-weight models. Muon had a negligible gap versus AdamW in one 8× setting but a higher benchmark average in four of five open-weight settings.

  • At 8× Chinchilla, the SFT-optimizer performance gap was negligible for AdamW-pretrained models, while Muon remained better for Muon-pretrained models at both 1× and 8×.Source: §4 · p. 5
  • Muon-SFT produced a higher mean benchmark average in 4 of 5 settings across AdamW-pretrained Llama and Qwen models up to 14B parameters.All gaps except perhaps Llama2-7B were small relative to run-to-run variation.Source: §4 · p. 5
  • The open-weight comparison covered AdamW-pretrained Llama2-7B, Qwen2.5-7B, and Qwen3 models at 4B, 8B, and 14B parameters.Each setting’s reported average was measured over three seeds.Source: §4 · p. 5

5 Conclusion

The conclusion finds little evidence that switching between AdamW and Muon adds a meaningful performance penalty beyond the optimizers’ independent effects. Muon-heavy SFT generally improves downstream performance, while optimizer consistency can reduce forgetting without improving final task-general SFT performance.

  • Final performance is largely determined by the independent effects of the pretraining and post-training optimizers, with little evidence of a meaningful mismatch penalty.Source: §5 · p. 5
  • Optimizer consistency can reduce forgetting of the pretraining distribution during SFT, but does not improve final performance under task-general SFT.Source: §5 · p. 5
  • SFT shows no significant interaction effects at either a fixed learning rate or the best learning rate.Source: §5 · p. 5
  • Muon remains competitive with AdamW for SFT of AdamW-pretrained open-weight Llama and Qwen models up to 14B parameters.Source: §5 · p. 6

A.1 Model Architecture

The appendix specifies the Llama-3 model architecture, optimizer algorithms, and pretraining setup used in the experiments.

A.1 Model Architecture

  • The 1.2B Llama-3 model has 32 layers, hidden dimension 1536, and 24 attention heads; its approximately 1.60B total parameters include approximately 1.21B non-embedding parameters.Source: §A.1 · p. 7

A.2 Optimizers

  • The study compares decoupled-weight-decay AdamW with Muon, which uses RMS update scaling and orthogonalization for matrix-valued updates.Source: §A.2 · pp. 7–8

A.3.1 Pretraining

  • Models are pretrained on the Wen et al. [2026] data mixture at 1× and 8× Chinchilla scales, with mixed runs sampling Muon updates at probability p and AdamW otherwise.Source: §A.3.1 · p. 8

A.3.2 SFT

The appendix describes SFT procedures for the 1.2B and open-weight models, including optimizer-specific learning-rate selection. Preferred post-training learning rates vary with both the pretraining optimizer and the post-training optimizer mixture.

A.3.2 SFT

  • For 1.2B models, SFT uses two epochs of the Tulu-3 mixture, and the learning rate is selected separately for each optimizer setting using held-out validation loss.Source: §A.3.2 · pp. 8–9
  • Across SFT and CPT, the optimal post-training learning rate depends on both the pretraining optimizer and the post-training optimizer mixture.Source: §A.3.2 · p. 10, Table 4 · p. 10, Table 5 · p. 10
  • Open-weight Llama and Qwen models use packed Tulu-3 SFT data, with sequence lengths of 4096 tokens for Llama-2 and 8192 for Qwen models.Source: §A.3.2 · pp. 8, 10

B.1 SFT benchmark performance for 1.2B models

The benchmark results report SFT performance at the best learning rate for each setting, with averages and standard deviations over three random seeds.

  • Table 7 results are used to calculate row-wise averages, column-wise averages, and residuals in Figure 1 (left).Source: §B.1 · p. 10, Table 7 · p. 11
  • Table 7 reports SFT benchmark performance at the best learning rate for each setting.Source: Table 7 · p. 11
  • Benchmark averages and standard deviations are reported across three random SFT seeds.Source: §B.1 · p. 10, Table 7 · p. 11

B.2 SFT benchmark performance at uniform LR for 1.2B models

At the majority-optimal uniform learning rate of 4e-4, increasing Muon updates during SFT generally improves final benchmark performance, with only mild optimizer-consistency interactions.

B.2 SFT benchmark performance at uniform LR for 1.2B models

  • Increasing the proportion of Muon updates during SFT generally increases final benchmark performance.The heatmap uses the same learning rate, 4e-4, for all settings.Source: §B.2 · p. 11
  • Residual benchmark cells show only mild interaction effects from optimizer consistency across training stages.Source: §B.2 · p. 11
  • The uniform learning rate of 4e-4 is the majority-optimal setting used for this companion heatmap.Source: §B.2 · p. 11

B.3 SFT benchmark performance by task type for 1.2B models

Across downstream task categories, SFT performance follows the aggregate pattern: more Muon updates during SFT generally improve performance, while optimizer-choice interactions are relatively small. The accompanying tables report forgetting and continued-pretraining loss summaries used in the paper’s analyses.

B.3 SFT benchmark performance by task type for 1.2B models

  • Across task categories, increasing Muon updates during SFT generally improves performance, with relatively small interaction effects from matching optimizers across stages.Figure 5 breaks down the Figure 1 results by task type; its rows index pretraining Muon proportions and columns index SFT Muon proportions.Source: §B.3 · p. 11, Figure 5 · p. 12

B.4 Post-SFT forgetting for 1.2B models

  • Table 8 reports change in pretraining validation loss after SFT, where smaller values indicate less forgetting and more retained pretraining knowledge.The reported values average three random seeds and are used to calculate Figure 1’s right-side marginal and residual effects.Source: Table 8 · p. 12, §B.4 · p. 12

B.5 CPT performance for 1.2B models

  • Tables 9 and 10 report three-seed average and standard deviation for continued-pretraining final loss, used to calculate Figure 2 residuals.Table 9 uses a uniform learning rate of 5e-5; Table 10 uses the best learning rate for each setting.Source: §B.5 · p. 12, Table 9 · p. 13, Table 10 · p. 13

Related papers

Loading 2610.22650v1…