🔬Modal Research Exchange

Since I keep getting asked about my Dimuon experiments: I benchmarked Polar Switch (https://arxiv.org/html/2610.26742v1) as a replacement for Polar Express in our DiMuon candidate, then pushed it further. The new candidate is TriMuon, with the refined polar approximation Polar Express Junction. The recovered momentum and preconditioning rules are unchanged.

The refined candidate is both faster and more accurate than our compiled PE baseline and upstream PS on the original saved Qwen3-8B rank-32 LoRA workload. For 288 factors, final-update error against the ideal FP64 SVD-based polar reference is: PE 14.41% mean / 41.02% p99 / 47.88% worst. PS3 11.18% / 35.25% / 41.26%. TriMuon (FP16 polar core + four refinements) 2.22% mean / 8.22% p99 / 9.90% worst, mean cosine 0.999619.
Matched optimizer GPU time, B200 / H100: PE 7.121 / 7.228 ms. PS 6.487 / 6.239 ms. TriMuon 5.986 / 6.028 ms (15.9% / 16.6% less time than PE).
Why both faster and more accurate? Two outer Polar Switch passes instead of five PE iterations, then four cubic Newton-Schulz refinements folded into the small rank x rank multiplier before its final rectangular apply.
Still to establish: better downstream models, convergence and end-to-end training cost. These are CUDA-graph optimizer-only measurements. We need paired Spindle LoRA training/hparam sweeps against PE, PS and AdamW before claiming customer benefit or releasing broadly.

[peyton] what experiments or justification would we need to actually release this?
[aaron] for normal pretraining I would do an hparam suite autoresearch. For this I guess we would try to replicate and beat Lora-MUON's benchmarks: https://arxiv.org/abs/2606.12921
[aaron] I also disabled TF32 for safety but I should see if that matters in practice
[aaron] I guess we should also test some big boy MoEs as I already did Qwen level widths
[peyton] in general, i'm wary about releasing absent having a significant body of work showing effectiveness
[aaron] We need an autoresearch harness for Spindle. Also that could be deployed to customers for LR sweeps etc
[aaron] Like DiMuon this has no beta2 no weight decay so that's nice
[kailash] https://x.com/anpaure/status/2115520552970649700 Re the dimuon reconstruction

{
  "adjacent_experiment": "",
  "reproducibility_reason": "Needs a Spindle fork: add the TriMuon optimizer (Polar Switch passes plus Newton-Schulz refinements) next to dimuon-lite. A short tinker SDK LoRA SFT script on a served Qwen model then tests the claim.",
  "reproducible_on_spindle": true,
  "sources": [
    {
      "brief": null,
      "error": "ReadTimeout: The read operation timed out",
      "full_text": false,
      "kind": "link",
      "title": "https://arxiv.org/html/2610.26742v1",
      "url": "https://arxiv.org/html/2610.26742v1"
    },
    {
      "brief": null,
      "error": "ReadTimeout: The read operation timed out",
      "full_text": false,
      "kind": "link",
      "title": "https://arxiv.org/abs/2606.12921",
      "url": "https://arxiv.org/abs/2606.12921"
    },
    {
      "brief": null,
      "error": "digest failed: ask_json failed after 2 attempts: no JSON object in response: '{\"title\":\"Tweet by @anpaure (thread opener, no formal title): \\\\\"Tinker\\'s Dimuon is Muon with delayed diagonal conditioning, Polar Express and a rank-dependent update norm. I recovered its update rule '",
      "full_text": true,
      "kind": "tweet",
      "title": "Post by @anpaure",
      "url": "https://x.com/anpaure/status/2115520552970649700"
    }
  ],
  "spec": {
    "base_model": "Qwen/Qwen3.5-4B",
    "baseline": "dimuon-lite (DiMuon with Polar Express), beta1=0.9, grad_clip_norm=0.0. Same model, rank, data order, seed, renderer, batch size, steps and LR grid. The best-LR run of each arm is compared.",
    "batch_size": 64,
    "budget": {
      "gpu_hours": 1.8,
      "max_wall_minutes": 120
    },
    "command": "python - <<'EOF'\n# Custom part: the TriMuon config exists only on the fork; the baseline uses tinker.types.DimuonOptimizerConfig() / DimuonParams.\nimport os, tinker\nfrom tinker_cookbook.recipes.chat_sl import chat_datasets  # renderer + no_robots builder\nsc = tinker.ServiceClient(base_url=os.environ[\"TINKER_BASE_URL\"])  # fork URL for the TriMuon arm\nARM = os.environ.get(\"ARM\", \"pe\"); LR = float(os.environ.get(\"LR\", \"1e-4\"))\nopt_cfg = tinker.types.DimuonOptimizerConfig() if ARM == \"pe\" else tinker.types.TrimuonOptimizerConfig()  # Trimuon* added by the fork\ntc = sc.create_lora_training_client(base_model=\"Qwen/Qwen3.5-4B\", rank=32, optimizer=opt_cfg)\ntrain, test = chat_datasets.no_robots_batches(tc.get_tokenizer(), batch_size=64, max_length=2048, seed=0)\nfor step, batch in zip(range(80), train):\n    tc.forward_backward(batch, loss_fn=\"cross_entropy\").result()\n    P = tinker.types.DimuonParams if ARM == \"pe\" else tinker.types.TrimuonParams\n    tc.optim_step(P(learning_rate=LR, beta1=0.9, grad_clip_norm=0.0)).result()\nprint(\"eval_nll\", sum(tc.forward(b, loss_fn=\"cross_entropy\").result().metrics[\"loss:mean\"] for b in test) / len(test))\ntc.save_weights_for_sampler(name=f\"rx-trimuon-{ARM}-{LR}\")\nEOF",
    "dataset": {
      "hf_path": "HuggingFaceH4/no_robots",
      "n_examples": 5120,
      "name": "no_robots",
      "split": "train (eval on test)"
    },
    "eval": {
      "benchmark": "no_robots test split",
      "metric": "mean per-token NLL on assistant tokens",
      "n_examples": 500
    },
    "hypothesis": "With a Spindle fork that adds TriMuon, LoRA SFT on Qwen3.5-4B reaches held-out no_robots loss at least 0.01 nats lower than dimuon-lite (Polar Express) at its best LR, after the same 80 steps.",
    "lora_rank": 32,
    "max_seq_len": 2048,
    "method": "Fork modal-projects/spindle and add TriMuon as an optimizer option. It keeps DiMuon's recovered momentum and preconditioning rules. It replaces only the Polar Express polar step with two outer Polar Switch passes, then four cubic Newton-Schulz refinements folded into the rank x rank multiplier. Use the FP16 polar core and TF32 disabled, as in aaron's benchmark. Deploy the fork as a throwaway copy. Train rank-32 LoRA on Qwen/Qwen3.5-4B with a short tinker SDK SFT loop on no_robots. The loop uses the cookbook renderer for the model, forward_backward with cross_entropy, then optim_step. The treatment arm runs on the fork with TriMuon. The baseline arm runs dimuon-lite on the same fork (or shared Spindle) with identical code. Each arm sweeps LR over {5e-5, 1e-4, 2e-4}: 6 runs of 80 steps each, same data order and seed. Each run evaluates held-out loss at step 80 and logs per-step optimizer wall time. Qwen3.5-4B is a scaled-down stand-in for the source's Qwen3-8B rank-32 workload. Polar Switch and AdamW arms are follow-up specs, not part of this decision.",
    "optimizer": {
      "lr": 0.0001,
      "name": "trimuon (fork)"
    },
    "steps": 80,
    "success_criterion": "Validated if best-LR TriMuon eval NLL is at least 0.01 nats below best-LR dimuon-lite eval NLL, and TriMuon's median optimizer step time is not higher. Failed if the gap is under 0.01 nats or reversed. A rerun of the best PE config with a second seed is optional, if budget allows, to check that the 0.01 margin exceeds seed noise.",
    "title": "TriMuon vs dimuon-lite (Polar Express) for LoRA SFT on no_robots",
    "training_type": "sft"
  },
  "summary": "TriMuon swaps DiMuon's Polar Express step for a faster, more accurate polar approximation; we compare held-out SFT loss against shared dimuon-lite at matched steps.",
  "tags": [
    "optimizer",
    "sft",
    "lora",
    "spindle-fork"
  ],
  "title": "Does TriMuon beat dimuon-lite's Polar Express on LoRA SFT eval loss?",
  "value_estimate": 4,
  "value_rationale": "A validated TriMuon win would decide whether Spindle's default DiMuon gets a faster, more accurate polar step that customers inherit, but one small SFT run is only a first gate before broader sweeps. Triage only saw the links for the Polar Switch paper, the LoRA-Muon paper and the @anpaure tweet; all three fetches failed."
}