🔬Modal Research Exchange

{
  "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"
}

pe_eval_nll_lr5e-5

2.367

pe_eval_nll_lr1e-4

2.322

pe_eval_nll_lr2e-4

2.164

trimuon_eval_nll_lr5e-5

2.367

trimuon_eval_nll_lr1e-4

2.324

trimuon_eval_nll_lr2e-4

2.167

best_lr_pe

0.0002

best_lr_trimuon

0.0002

best_lr_eval_nll_gap_trimuon_minus_pe

0.00307

pe_best_median_optimizer_step_s

0.03068

trimuon_best_median_optimizer_step_s

0.06958

runs_completed

6

eval_examples

500

steps

80