Thinking Machines' Dimuon Optimizer Beats AdamW on LoRA Fine-Tuning Steps

Thinking Machines Lab has added Dimuon to Tinker, its managed LoRA fine-tuning API. In the company’s benchmarks, Dimuon reached lower loss in fewer supervised fine-tuning steps than AdamW while taking roughly the same wall-clock time per step. The optimizer is available now and can be selected when creating a training client.

Dimuon gives developers a second optimizer for low-rank adaptation, or LoRA, which freezes the base model and trains smaller adapter matrices. Thinking Machines reports the clearest gains in supervised fine-tuning (SFT). Its reinforcement learning (RL) tests showed performance similar to AdamW. These results come from internal benchmarks, so workloads should be tested against a tuned AdamW baseline.

Dimuon joins AdamW

Tinker now supports AdamW, which remains the default, and Dimuon. The optimizer is fixed for the lifetime of a training client, so new jobs can enable Dimuon with one configuration change:

training_client = service_client.create_lora_training_client(
    base_model="Qwen/Qwen3-8B",
    rank=64,
    optimizer=types.DimuonOptimizerConfig(),
)

training_client.forward_backward(data, "cross_entropy")
training_client.optim_step(
    types.DimuonParams(learning_rate=1e-3)
)

The name references Muon, Keller Jordan’s matrix-aware optimizer. Thinking Machines developed Dimuon for its LoRA training system.

Dimuon exposes a small set of hyperparameters:

Retune before comparing

Dimuon generally needs a learning rate 10 to 100 times larger than AdamW. Thinking Machines found that roughly 15 times the tuned AdamW rate worked best in its SFT experiments, while its RL tests favored a ratio closer to 60 times. A practical initial comparison is 1e-3 for Dimuon against 1e-4 for AdamW, followed by a sweep around those values.

Thinking Machines recommends scaling the SFT learning rate with the square root of batch size. A rate tuned at a small LoRA rank can also serve as the starting point for higher ranks. In the company’s experiments, Dimuon’s preferred learning rate changed slightly less across ranks than AdamW’s.

Results favor supervised tuning

Across the published experiments, Dimuon’s advantage depended on the training regime:

LoRA rank controls adapter capacity and trainable parameter count. Increasing it can improve results, but it also raises memory and compute requirements. Rank comparisons should therefore record training cost alongside loss and task quality.

Thinking Machines also reports that linear learning-rate decay to zero improved SFT step efficiency by 1.3 times over a constant rate in its tests. Tinker lets callers implement the schedule by changing the learning_rate passed to each optim_step.

Client boundaries shape migration

Changing the optimizer during a run requires saving a checkpoint and creating a new training client because each client keeps its original optimizer configuration:

sft_checkpoint = training_client.save_state(
    "sft-complete"
).result()

rl_client = service_client.create_training_client_from_state(
    sft_checkpoint.path,
    optimizer=types.AdamOptimizerConfig(),
)

Optimizer state, including momentum and running averages, resets by default when the new client is created. Use create_training_client_from_state_with_optimizer when the next stage should retain that state.

Sampler checkpoints created with save_weights_for_sampler contain weights for inference and do not carry optimizer state. Training workflows that may resume or change stages should retain state checkpoints separately.

Fewer steps can lower SFT cost

LoRA reduces fine-tuning memory and storage by training compact adapters, but the number of optimization steps still affects runtime and expense. If Dimuon reaches a target validation score in fewer steps without increasing step time, an SFT job can finish sooner at lower cost. The available evidence supports that claim for some internal SFT workloads, with no comparable gain reported for RL-only training.

Thinking Machines has previously published broader LoRA guidance arguing that well-configured adapters can match full fine-tuning across many SFT and RL tasks. Dimuon targets the optimization efficiency of those adapter runs, especially when larger supervised datasets or higher ranks require more training.

Run a controlled A/B test

A controlled comparison should tune both optimizers and measure validation quality, stability, total steps, and wall-clock time. A useful test plan is:

1. Keep the base model, dataset, data order, batch size, LoRA rank, and evaluation procedure fixed.
2. Establish a tuned AdamW baseline.
3. Sweep Dimuon learning rates from roughly 10 to 100 times the AdamW rate, including a point near 15 times for SFT.
4. Compare constant and linear-decay schedules separately.
5. Repeat promising SFT runs at higher ranks if memory and compute allow.
6. Select the configuration using validation or task metrics, total runtime, and failure rate, not training loss alone.

Thinking Machines says Dimuon remains under active development and is seeking user feedback. Developers can inspect the Tinker repository and optimizer documentation before adding it to production training workflows.