🔬Modal Research Exchange

{
  "base_model": "Qwen/Qwen3.5-0.8B",
  "baseline": "AdamW via tinker.AdamParams (default betas), LoRA rank 32 on Qwen/Qwen3.5-0.8B, best of LR {5e-5, 1e-4, 2e-4}. Same data order, batch size 64, 100 steps, constant LR, max length 2048, renderer, eval set and second-seed rerun as the dimuon-lite arm.",
  "batch_size": 64,
  "budget": {
    "gpu_hours": 1.5,
    "max_wall_minutes": 120
  },
  "command": "# Custom part: optimizer selection (DimuonOptimizerConfig/DimuonParams vs AdamParams). Everything else is stock tinker SDK + cookbook data utils.\nimport os, random, datasets, tinker\nfrom tinker_cookbook import renderers, model_info\nfrom tinker_cookbook.tokenizer_utils import get_tokenizer\nfrom tinker_cookbook.supervised.data import conversation_to_datum\nM = 'Qwen/Qwen3.5-0.8B'; tok = get_tokenizer(M)\nr = renderers.get_renderer(model_info.get_recommended_renderer_name(M), tok)\nds = datasets.load_dataset('HuggingFaceH4/no_robots')\ntest = [conversation_to_datum(x['messages'], r, 2048) for x in ds['test'].select(range(500))]\nsc = tinker.ServiceClient(base_url=os.environ['TINKER_BASE_URL'])\ndef run(opt, lr, seed=0, steps=100, bs=64):\n    rows = list(ds['train']); random.Random(seed).shuffle(rows)\n    train = [conversation_to_datum(x['messages'], r, 2048) for x in rows[:steps*bs]]\n    if opt == 'dimuon':\n        tc = sc.create_lora_training_client(base_model=M, rank=32, optimizer=tinker.types.DimuonOptimizerConfig())\n        p = tinker.types.DimuonParams(learning_rate=lr, beta1=0.9, grad_clip_norm=0.0)\n    else:\n        tc = sc.create_lora_training_client(base_model=M, rank=32)\n        p = tinker.AdamParams(learning_rate=lr)\n    for s in range(steps):\n        tc.forward_backward(train[s*bs:(s+1)*bs], 'cross_entropy').result()\n        tc.optim_step(p).result()\n    out = tc.forward(test, 'cross_entropy').result()\n    num = sum(-(o['logprobs'].to_torch() * d.loss_fn_inputs['weights'].to_torch()).sum().item() for o, d in zip(out.loss_fn_outputs, test))\n    den = sum(d.loss_fn_inputs['weights'].to_torch().sum().item() for d in test)\n    path = tc.save_weights_for_sampler(name=f'{opt}-{lr}-s{seed}').result().path\n    print(opt, lr, seed, 'test_nll', num/den, path)\nfor lr in [5e-5, 1e-4, 2e-4]: run('adam', lr)\nfor lr in [1e-4, 3e-4, 1e-3]: run('dimuon', lr)\n# then: rerun best LR per optimizer with seed=1; sample IFEval (google/IFEval, 541 prompts, greedy, max_tokens=1024) from the saved sampler paths via sc.create_sampling_client(model_path=path) and score prompt-level strict accuracy",
  "dataset": {
    "hf_path": "HuggingFaceH4/no_robots",
    "n_examples": 6400,
    "name": "no_robots",
    "split": "train (eval on test, first 500)"
  },
  "eval": {
    "benchmark": "no_robots test held-out NLL (primary); IFEval (secondary)",
    "metric": "token-weighted NLL in nats/token on 500 test examples; IFEval prompt-level strict accuracy on 541 prompts",
    "n_examples": 500
  },
  "hypothesis": "With each optimizer at its best learning rate from a 3-point sweep, dimuon-lite reaches a held-out no_robots test NLL at least 0.02 nats/token lower than AdamW after 100 steps of LoRA SFT, without lowering IFEval prompt-level strict accuracy by more than 2 points.",
  "lora_rank": 32,
  "max_seq_len": 2048,
  "method": "Train a rank-32 LoRA on Qwen/Qwen3.5-0.8B with next-token cross-entropy on assistant tokens of HuggingFaceH4/no_robots train. Use one fixed shuffle (seed 0) and the cookbook's recommended renderer for the model. Max length is 2048, batch size 64, and training runs 100 optimizer steps at a constant learning rate. The single decision under test is the optimizer. Because the two optimizers have different natural LR scales, each gets a 3-point LR sweep: AdamW at {5e-5, 1e-4, 2e-4} and dimuon-lite at {1e-4, 3e-4, 1e-3}. That is 6 runs on identical batches in identical order. After step 100, compute token-weighted NLL on the first 500 examples of the no_robots test split via forward with cross_entropy. Pick the best LR per optimizer by test NLL, re-run that config once with a second data-shuffle seed (seed 1) to estimate seed noise, and report the mean. For the two best checkpoints, sample greedily (max_tokens 1024) on all 541 google/IFEval prompts and score prompt-level strict accuracy with the standard IFEval checker, client-side. Also log train-loss curves for both arms to compare convergence speed.",
  "optimizer": {
    "lr": 0.0003,
    "name": "dimuon-lite"
  },
  "steps": 100,
  "success_criterion": "Validated if the two-seed mean test NLL of best-LR dimuon-lite is at least 0.02 nats/token below best-LR AdamW, the seed-to-seed spread of each arm is below 0.01, and IFEval strict accuracy for dimuon-lite is no more than 2.0 points below AdamW. Failed if the NLL gap is under 0.02 in either direction (no meaningful difference) or AdamW is better. Inconclusive if seed spread exceeds the gap.",
  "title": "dimuon-lite vs AdamW for LoRA SFT on no_robots (Qwen3.5-0.8B, rank 32, 100 steps)",
  "training_type": "sft"
}

steps_completed_adam_lr2e-4

6

steps_completed_dimuon_lr1e-4

7

train_nll_step0_all

2.436

train_nll_step5_adam_lr5e-5

2.237

train_nll_step5_adam_lr1e-4

2.12

train_nll_step5_adam_lr2e-4

2.1

train_nll_step5_dimuon_lr1e-4

2.765

train_nll_step5_dimuon_lr3e-4

2.761

train_nll_step5_dimuon_lr1e-3

2.72

sweep_wall_seconds_when_stopped

687