The Case Against Switching Optimizers Mid-Training: A Technical Deep Dive

Introduction

In the middle of a complex, multi-day machine learning training run, a critical question arises: should we switch optimizers? This moment—captured in message index 8165 of an opencode coding session—represents a fascinating intersection of technical judgment, engineering pragmatism, and the discipline required to resist the allure of novel techniques when the stakes are high. The message is the assistant's response to a user query about replacing the AdamW optimizer with the Muon optimizer for an ongoing DFlash speculative decoding drafter training pipeline. What unfolds is not merely a yes-or-no answer, but a masterclass in reasoning about optimization choices under real-world constraints.

The subject message (see [msg 8165]) contains both the assistant's internal reasoning and the final response to the user. The assistant recommends against switching, and the analysis is worth examining in depth because it reveals how an experienced practitioner evaluates trade-offs when a training run is already in progress, healthy, and producing results.

The Broader Context: What Is Being Trained?

To understand the significance of this optimizer question, we must first understand what is at stake. The project involves training a DFlash block-diffusion speculative decoding drafter for the Qwen3.6-27B language model. This is not a small undertaking. The drafter has 1.7 billion trainable parameters, and the training pipeline has been the subject of intense engineering effort across multiple prior segments of the conversation.

The training setup is running on a machine with 4× RTX PRO 6000 Blackwell GPUs (96 GB each), using a 3-1 GPU topology: three GPUs run the target model forward pass to generate hidden states, and one GPU trains the drafter. The pipeline has been optimized to run at approximately 16 Ktok/s with 100% GPU utilization on the target GPUs, reducing the estimated 6-epoch training time from an original 22.9 days to roughly 8 days. At the time of this message, training has reached approximately step 15,465, with the learning rate at only 40% of its peak value (2.43e-4 out of a 6e-4 peak). The loss has dropped from 12.5 to approximately 1.4, accuracy has risen from 0 to 0.17, and the estimated acceptance length of approximately 3.6 already exceeds the z-lab baseline of 3.1.

This is a training run that is demonstrably working. The user's question about Muon arises not from a crisis or failure, but from curiosity about whether a better optimizer could produce superior results.

What Is Muon, and Why Would Anyone Consider It?

Muon (short for "Momentum Orthogonalization Update") is a relatively recent optimizer that gained prominence in the NanoGPT speedrun community. Developed by Jordan Keller and collaborators, Muon applies Nesterov momentum followed by orthogonalization via Newton-Schulz iteration to the weight updates. The key insight is that for weight matrices, the optimal update direction is often orthogonal to the current weight matrix, and Muon approximates this by computing the nearest orthogonal matrix to the momentum-adjusted gradient.

Muon has shown impressive results in certain regimes, particularly for training smaller language models from scratch with limited data. In those settings, it can achieve comparable loss levels in significantly fewer steps than AdamW. This has led to considerable excitement in the ML community, with practitioners eager to apply it to their own training runs.

The user's question reflects this enthusiasm. If Muon could converge to the same accuracy in half the steps, that would translate to faster training, lower cost, and potentially better final model quality. It is a reasonable question to ask.

The Assistant's Reasoning Process

The assistant's response in [msg 8165] is structured as a two-part message: first, an internal "Agent Reasoning" section that reveals the deliberative process, and second, the actual response to the user. This dual structure is characteristic of the opencode assistant's approach and provides rare visibility into how the AI weighs competing factors.

The reasoning begins by acknowledging what Muon is and where it has shown success. The assistant notes that Muon "typically works best on weight matrices while other parameter types like embeddings, biases, and layer normalization still benefit from AdamW." This is an important nuance: Muon is not a drop-in replacement for AdamW on all parameters. In practice, implementations of Muon often use a hybrid approach where Muon is applied to weight matrices and AdamW is retained for everything else. This hybrid strategy introduces additional complexity and hyperparameter tuning.

The assistant then identifies the single most important constraint: the real bottleneck in our setup isn't optimizer compute time—it's the target model's forward pass generating hidden states for the drafter. This observation is crucial because it reframes the entire question. If the drafter training step is already faster than the target forward pass (the drafter GPU idles waiting for hidden states), then making the optimizer faster does nothing for wall-clock time. Even if Muon converged in half the steps, the pipeline would still be gated on the target model producing hidden states at the same rate.

This is a classic systems-thinking insight. The assistant recognizes that the optimization problem is not isolated to the drafter training loop; it is embedded in a larger pipeline with multiple interacting components. Optimizing a non-bottleneck component yields no system-level improvement.

The Four Pillars of the Argument Against Switching

The assistant's response to the user crystallizes into four main arguments, each addressing a different dimension of the decision.

1. The Bottleneck Argument

The first and most powerful argument is that the pipeline is bottlenecked on the target model's forward pass, not on optimizer convergence. The assistant states this clearly: "The drafter trains at ~0.33 batch/s but targets only produce ~0.24 batch/s. The drafter GPU idles waiting. Even if Muon converges in 2x fewer steps, wall-clock time barely changes — we're gated on target hidden state production."

This argument is grounded in the detailed throughput modeling that was done in earlier segments of the conversation. The pipeline's performance characteristics have been carefully measured and modeled. The target model's forward pass through its 64 layers (48 GDN layers and 16 attention layers) is computationally expensive, and the three target GPUs each produce batches at a rate that the single drafter GPU can consume without being the bottleneck. Switching to a faster optimizer would not change the rate at which hidden states arrive at the drafter.

2. The Mid-Run Switching Argument

The second argument addresses the practical difficulty of switching optimizers when a training run is already 15,000+ steps in. The assistant notes that "switching optimizer means losing Adam momentum/variance state and restarting warmup." This is not a trivial concern.

AdamW maintains two state variables per parameter: the first moment (momentum) and the second moment (variance) estimates. These accumulate information about the gradient history and are essential for AdamW's adaptive learning rate behavior. Switching to Muon mid-run means discarding this accumulated state. The learning rate schedule would also need to be reconsidered, as Muon has different optimal learning rate characteristics than AdamW.

Furthermore, the assistant notes that the learning rate is "only at 40% of peak — we haven't even hit the high-convergence phase yet." This is a critical observation. The training run is still in its warmup phase. The most significant accuracy improvements are expected to occur as the learning rate ramps to its peak value of 6e-4. Interrupting the run now, before the LR has even peaked, would be premature.

3. The Architecture Mismatch Argument

The third argument is perhaps the most technically subtle. The assistant points out that "Muon's Newton-Schulz orthogonalization is proven on standard transformer attention+FFN. Our drafter is 80% GDN layers (gated delta net from FLA) — unusual recurrent-style layers with different gradient dynamics. Muon is untested on this."

This is a genuinely novel concern. The DFlash drafter architecture, as implemented by the z-lab team and adapted for this project, uses GDN (Gated Delta Network) layers for the majority of its computation. GDN layers have a different structure than standard transformer layers—they involve recurrent-style computations with gating mechanisms that produce different gradient properties. Muon's orthogonalization step assumes that the optimal update direction for a weight matrix is approximately orthogonal to the current weights, an assumption that holds well for standard attention and feedforward layers but may not hold for GDN layers with their different dynamics.

The assistant's reasoning here demonstrates deep understanding of both the optimizer and the model architecture. It recognizes that the theoretical foundations of Muon may not transfer to non-standard architectures without empirical validation.

4. The Diminishing Returns Argument

The fourth argument addresses the training regime. Muon's biggest wins have been demonstrated in low-data, short-training regimes—the NanoGPT speedrun community focuses on training small models to convergence as quickly as possible. The assistant argues that "we're training 6 full epochs (~10B tokens for a 1.7B model) — AdamW converges fine given enough passes over the data."

This argument is about the relationship between optimizer choice and training duration. AdamW with a cosine learning rate schedule is a well-understood combination that reliably converges to good solutions given sufficient training steps. Muon's advantage is most pronounced when training steps are scarce. With 6 epochs of data and a well-tuned schedule, the marginal benefit of a "faster" optimizer diminishes.

Assumptions and Their Evaluation

The assistant's analysis rests on several key assumptions, and it is worth examining them critically.

Assumption 1: The target forward pass is the true bottleneck. This is well-supported by the pipeline measurements conducted in prior segments. The 3-1 topology was specifically chosen because modeling showed that three target GPUs produce batches at a rate that keeps the drafter busy without being overwhelmed. The assistant's confidence in this assumption is justified by empirical data.

Assumption 2: Switching optimizers mid-run is net negative. This is a strong claim, but it is well-reasoned. The loss of momentum/variance state, the need to restart warmup, and the disruption to a healthy training run all point in the same direction. However, one could imagine a scenario where Muon's different convergence properties could help escape a local minimum that AdamW is stuck in. The assistant implicitly assumes that the current convergence trajectory is good enough to continue, and the data supports this: accuracy is climbing, loss is decreasing, and the acceptance length estimate already exceeds the baseline.

Assumption 3: GDN layers are untested with Muon. This is factually correct—there is no published work on using Muon with GDN-based architectures. The assistant treats this as a risk factor, which is the prudent approach. An alternative perspective might argue that Muon's orthogonalization is architecture-agnostic and should work regardless of layer type, but this would be an untested hypothesis.

Assumption 4: 6 epochs of training reduces Muon's advantage. This is supported by the general observation that AdamW converges well given sufficient data passes. However, it is worth noting that Muon could potentially achieve better final accuracy even with 6 epochs, not just faster convergence to the same accuracy. The assistant implicitly assumes that AdamW and Muon would converge to similar final accuracy given enough steps, which is a reasonable but unverified assumption for this specific architecture and dataset.

What Knowledge Was Required to Understand This Message?

To fully appreciate the assistant's reasoning, a reader would need:

1. Understanding of AdamW and Muon optimizers: Knowledge of how momentum, variance estimation, and orthogonalization work, and what trade-offs exist between adaptive and non-adaptive optimization methods.
2. Knowledge of the DFlash architecture: Understanding that the drafter uses GDN layers (gated delta networks) which have different gradient dynamics than standard transformer layers, and that this matters for optimizer compatibility.
3. Pipeline architecture awareness: Understanding that the training pipeline is a multi-stage asynchronous system where the target model forward pass and drafter training are decoupled, and that bottlenecks can exist at any stage.
4. Training progress context: Knowing that the run is at step ~15,465, 17% through epoch 1, with LR at 40% of peak, and that accuracy is already at 0.17 with estimated acceptance length 3.6.
5. The z-lab baseline: Knowing that the z-lab drafter achieves acceptance length 3.1, and that our current training already exceeds this with only 17% of epoch 1 complete.
6. Understanding of training dynamics: Knowledge of warmup phases, cosine schedules, and how learning rate affects convergence speed and final quality.

What Knowledge Was Created by This Message?

The message creates several important pieces of output knowledge:

1. A decision record: The explicit recommendation against switching optimizers, with documented reasoning that can be referenced later.
2. A framework for evaluating optimizer switches: The four arguments (bottleneck, mid-run disruption, architecture mismatch, diminishing returns) provide a reusable template for similar decisions in the future.
3. A risk assessment of Muon for GDN architectures: The observation that Muon is untested on GDN layers is a novel contribution that may inform future work.
4. A prioritization of engineering effort: By ruling out the optimizer switch, the message implicitly directs attention to other levers for improvement: letting the current run finish, evaluating the trained drafter, and potentially exploring DDTree integration post-training.
5. A reinforcement of the "let it run" philosophy: The message reinforces the principle that a healthy training run should not be disrupted without compelling evidence of benefit.

The Decision-Making Philosophy

Beyond the specific technical arguments, the assistant's response reveals a broader decision-making philosophy that is worth examining.

First principle: Bottleneck-centric optimization. The assistant consistently identifies the actual bottleneck (target forward pass) before evaluating any proposed change. This prevents wasted effort on optimizing non-bottleneck components. It is a lesson that applies far beyond this specific context: always measure before optimizing.

Second principle: Risk-reward calibration. The assistant weighs the potential benefits of Muon (faster convergence, possibly better final accuracy) against the risks (disrupted training, lost state, untested on GDN layers, hyperparameter uncertainty). The conclusion is that the risk-reward ratio is unfavorable. This is a mature engineering judgment that prioritizes stability over speculative gains.

Third principle: Respect for momentum. The training run has momentum—both literally (Adam's momentum state) and figuratively (the trajectory of convergence, the engineering investment, the operational stability). Disrupting this momentum requires a compelling reason, and the assistant correctly judges that the case for Muon is not compelling enough.

Fourth principle: Defer experimentation to the right time. The assistant suggests that if we were starting fresh, an A/B test between AdamW and Muon would be worthwhile. This is not a rejection of experimentation; it is a recognition that mid-run is the wrong time for it. The message implicitly proposes a workflow: experiment offline, then apply the winning configuration to production runs.

Potential Weaknesses in the Analysis

No analysis is perfect, and it is worth considering potential blind spots in the assistant's reasoning.

Could Muon improve final accuracy beyond what AdamW can achieve? The assistant assumes that AdamW will converge to a good solution given 6 epochs, but this is not guaranteed. If the loss landscape has features that AdamW struggles with (e.g., sharp minima, saddle points), Muon's different update dynamics might navigate them more effectively. The assistant's argument about diminishing returns is plausible but not proven.

Is the bottleneck argument too static? The assistant assumes that the current bottleneck (target forward pass) will remain the bottleneck throughout training. But if the drafter's training dynamics change—for example, if accuracy improvements make the drafter's forward pass slower due to different computation patterns—the bottleneck could shift. This is a minor concern but worth noting.

Could the warmup restart actually be beneficial? The assistant treats the loss of Adam state as purely negative, but one could argue that a fresh start with Muon might escape any bad conditioning that has accumulated in Adam's state. This is speculative, but the assistant does not consider this possibility.

The GDN mismatch argument is untested. The assistant argues that Muon is untested on GDN layers, which is true, but it does not consider the possibility that Muon might work better on GDN layers than on standard layers. The uncertainty cuts both ways, and the assistant treats it as a risk factor rather than an opportunity.

These potential weaknesses do not invalidate the assistant's conclusion, but they highlight the inherent uncertainty in any such analysis. The assistant's recommendation is a judgment call, not a mathematical proof.

Conclusion

The subject message at [msg 8165] is a remarkable example of technical decision-making under uncertainty. The assistant synthesizes knowledge about optimizer theory, pipeline architecture, training dynamics, and model architecture to produce a nuanced recommendation that prioritizes stability over novelty.

The core insight is that the optimizer question cannot be evaluated in isolation. The answer depends on the pipeline's bottleneck structure, the training stage, the model architecture, and the risk tolerance of the project. By considering all of these factors together, the assistant arrives at a conclusion that is greater than the sum of its parts.

The message also serves as a case study in engineering discipline. It is tempting to chase the latest technique, especially when it has generated excitement in the community. But the assistant's response demonstrates the value of asking not "is this technique good?" but "is this technique good for our specific situation, right now?" The answer to the first question might be "yes, Muon is promising." The answer to the second question is "no, not mid-run, not for this architecture, not when the bottleneck is elsewhere."

This distinction—between absolute quality and contextual fit—is at the heart of effective engineering. The assistant's message is a reminder that the best decision is not always the one that uses the most advanced tool, but the one that produces the best outcome given the constraints.

The training continues with AdamW. The run is healthy. The LR is still ramping. Most of the learning is still ahead.