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

"""LoRA SFT on a Spindle server through the stock Tinker SDK.

dimuon-lite vs AdamW on HuggingFaceH4/no_robots.

Reads TINKER_BASE_URL, TINKER_API_KEY, MODAL_PROXY_TOKEN_ID and
MODAL_PROXY_TOKEN_SECRET. Trains two LoRA arms (dimuon primary, AdamW baseline)
in lockstep on identical batches in identical order. Measures token-weighted
held-out NLL on the first 500 no_robots test examples before and after
training, and prints one final line `RESULT {...}` with the metrics.
"""

import json
import math
import os
import time

import tinker
from datasets import load_dataset
from tinker import types

# Deviations from the spec, forced by this server and the script budget:
# - Qwen/Qwen3.5-0.8B is not served here, so this uses Qwen/Qwen3.5-2B, the documented fallback.
# - CONFIG["steps"]=100 is a hard cap on total optimizer steps. Both arms run inside it,
#   so each arm gets steps // 2 = 50 steps (n_examples = 50 * 64 = 3200 per arm, the same
#   batches for both arms).
# - No 3-point LR sweep and no second-seed rerun fit in the budget. Each optimizer runs once,
#   with the spec's dimuon LR (3e-4) and the midpoint of the spec's AdamW sweep (1e-4).
# - IFEval is skipped. The standard checker needs packages outside the allowed set.
# - The data split is train with seed-0 shuffle. Eval uses the test split, first 500 rows.
#   Prompts use the tokenizer chat template in place of the cookbook renderer. Loss covers
#   the first assistant turn, with the system prompt included when present.
# CONFIG-START
CONFIG = {
    "base_model": "Qwen/Qwen3.5-2B",
    "lora_rank": 32,
    "dataset": "HuggingFaceH4/no_robots",
    "split": "train",
    "eval_split": "test",
    "shuffle_seed": 0,
    "n_examples": 3200,
    "n_eval": 500,
    "steps": 100,
    "batch_size": 64,
    "max_seq_len": 2048,
    "lr": 0.0003,
    "optimizer": "dimuon",
    "baseline_optimizer": "adam",
    "baseline_lr": 0.0001,
    "eval_chunk": 64,
    "max_wall_seconds": 2700
}
# CONFIG-END

STARTED = time.monotonic()
# Leave part of the wall budget for the two final held-out evals.
TRAIN_WALL_FRACTION = 0.85


def elapsed():
    return time.monotonic() - STARTED


def row_to_pair(row):
    turns = row.get("messages") or []
    system = next((t["content"] for t in turns if t["role"] == "system"), None)
    user = next((t["content"] for t in turns if t["role"] == "user"), None)
    assistant = next((t["content"] for t in turns if t["role"] == "assistant"), None)
    if user and assistant:
        return (system, user, assistant)
    return None


def load_examples():
    """Return (train, eval) lists of (system, prompt, response) triples."""
    train_rows = load_dataset(CONFIG["dataset"], split=CONFIG["split"]).shuffle(seed=CONFIG["shuffle_seed"])
    train = []
    for row in train_rows:
        pair = row_to_pair(row)
        if pair:
            train.append(pair)
        if len(train) == CONFIG["n_examples"]:
            break
    test_rows = load_dataset(CONFIG["dataset"], split=CONFIG["eval_split"])
    test_rows = test_rows.select(range(min(CONFIG["n_eval"], len(test_rows))))
    evals = [p for p in (row_to_pair(r) for r in test_rows) if p]
    return train, evals


def make_datum(tokenizer, system, prompt, response):
    """Tokenize one chat turn. Loss weights cover only the response tokens."""
    msgs = []
    if system:
        msgs.append({"role": "system", "content": system})
    msgs.append({"role": "user", "content": prompt})
    prompt_ids = tokenizer.apply_chat_template(msgs, tokenize=True, add_generation_prompt=True)
    if hasattr(prompt_ids, "keys"):
        prompt_ids = prompt_ids["input_ids"]
    prompt_ids = list(prompt_ids)
    response_ids = tokenizer.encode(response, add_special_tokens=False) + [tokenizer.eos_token_id]
    tokens = (prompt_ids + response_ids)[: CONFIG["max_seq_len"] + 1]
    weights = ([0.0] * len(prompt_ids) + [1.0] * len(response_ids))[: len(tokens)]
    datum = types.Datum(
        model_input=types.ModelInput.from_ints(tokens[:-1]),
        loss_fn_inputs={"target_tokens": tokens[1:], "weights": weights[1:]},
    )
    return datum, sum(weights[1:])


def mean_loss(output, batch):
    """Per-token loss over a batch of (datum, n_loss_tokens) pairs."""
    loss_tokens = sum(n for _, n in batch)
    return output.metrics["loss:sum"] / max(loss_tokens, 1.0)


def eval_loss(training, data):
    """Token-weighted NLL (nats/token) over data, sent in chunks."""
    chunk = CONFIG["eval_chunk"]
    pieces = [data[i : i + chunk] for i in range(0, len(data), chunk)]
    futures = [training.forward([d for d, _ in p], "cross_entropy") for p in pieces]
    total = sum(f.result().metrics["loss:sum"] for f in futures)
    tokens = sum(n for _, n in data)
    return total / max(tokens, 1.0)


def optim_params(name, lr):
    if name == "dimuon":
        return types.DimuonParams(learning_rate=lr)
    return types.AdamParams(learning_rate=lr)


def main():
    service = tinker.ServiceClient(
        base_url=os.environ["TINKER_BASE_URL"],
        api_key=os.environ["TINKER_API_KEY"],
        default_headers_overrides={
            "Modal-Key": os.environ["MODAL_PROXY_TOKEN_ID"],
            "Modal-Secret": os.environ["MODAL_PROXY_TOKEN_SECRET"],
        },
    )
    arms = {
        "primary": {
            "optimizer": CONFIG["optimizer"],
            "lr": CONFIG["lr"],
            "client": service.create_lora_training_client(
                base_model=CONFIG["base_model"], rank=CONFIG["lora_rank"]
            ),
            "losses": [],
        },
        "baseline": {
            "optimizer": CONFIG["baseline_optimizer"],
            "lr": CONFIG["baseline_lr"],
            "client": service.create_lora_training_client(
                base_model=CONFIG["base_model"], rank=CONFIG["lora_rank"]
            ),
            "losses": [],
        },
    }
    tokenizer = arms["primary"]["client"].get_tokenizer()
    print(f"training clients ready for {CONFIG['base_model']} after {elapsed():.0f}s", flush=True)

    train_ex, eval_ex = load_examples()
    train_data = [make_datum(tokenizer, s, p, r) for s, p, r in train_ex]
    eval_data = [make_datum(tokenizer, s, p, r) for s, p, r in eval_ex]
    print(f"loaded {len(train_data)} train / {len(eval_data)} eval examples", flush=True)

    # LoRA starts as a zero delta, so the base-model loss is the same for both arms.
    eval_before = eval_loss(arms["primary"]["client"], eval_data) if eval_data else None
    print(f"eval loss before training: {eval_before}", flush=True)

    per_arm_steps = CONFIG["steps"] // len(arms)
    bs = CONFIG["batch_size"]
    train_budget = CONFIG["max_wall_seconds"] * TRAIN_WALL_FRACTION
    for step in range(per_arm_steps):
        if elapsed() > train_budget:
            print(f"stopping at step {step}: wall-time budget reached", flush=True)
            break
        start = step * bs % len(train_data)
        batch = (train_data + train_data)[start : start + bs]
        datums = [d for d, _ in batch]
        step_started = time.monotonic()
        fb = {k: a["client"].forward_backward(datums, "cross_entropy") for k, a in arms.items()}
        outs = {k: f.result() for k, f in fb.items()}
        opt = {
            k: a["client"].optim_step(optim_params(a["optimizer"], a["lr"])) for k, a in arms.items()
        }
        for f in opt.values():
            f.result()
        msg = []
        for k, a in arms.items():
            loss = mean_loss(outs[k], batch)
            if not math.isfinite(loss):
                raise RuntimeError(f"step {step} ({k}): loss is {loss}")
            a["losses"].append(loss)
            msg.append(f"{a['optimizer']} {loss:.4f}")
        print(
            f"step {step}: train loss {' | '.join(msg)} ({time.monotonic() - step_started:.1f}s)",
            flush=True,
        )

    evals_after = {}
    for k, a in arms.items():
        evals_after[k] = eval_loss(a["client"], eval_data) if eval_data else None
        print(f"eval loss after training ({a['optimizer']}): {evals_after[k]}", flush=True)

    p, b = arms["primary"], arms["baseline"]
    losses, blosses = p["losses"], b["losses"]
    gap = (
        evals_after["baseline"] - evals_after["primary"]
        if evals_after["baseline"] is not None and evals_after["primary"] is not None
        else None
    )
    result = {
        "train_loss_first": losses[0] if losses else None,
        "train_loss_last": losses[-1] if losses else None,
        "train_loss_min": min(losses) if losses else None,
        "eval_loss_before": eval_before,
        "eval_loss_after": evals_after["primary"],
        "steps": len(losses),
        "n_examples": len(train_data),
        "n_eval": len(eval_data),
        "base_model": CONFIG["base_model"],
        "lora_rank": CONFIG["lora_rank"],
        "lr": CONFIG["lr"],
        "optimizer": CONFIG["optimizer"],
        "dataset": CONFIG["dataset"],
        "wall_seconds": round(elapsed(), 1),
        "train_losses": [round(x, 5) for x in losses],
        "total_optimizer_steps": len(losses) + len(blosses),
        "baseline": {
            "optimizer": b["optimizer"],
            "lr": b["lr"],
            "train_loss_first": blosses[0] if blosses else None,
            "train_loss_last": blosses[-1] if blosses else None,
            "train_loss_min": min(blosses) if blosses else None,
            "eval_loss_before": eval_before,
            "eval_loss_after": evals_after["baseline"],
            "steps": len(blosses),
            "train_losses": [round(x, 5) for x in blosses],
        },
        "nll_gap_baseline_minus_primary": gap,
    }
    print("RESULT " + json.dumps(result), flush=True)


if __name__ == "__main__":
    main()

command

python experiment_run_0ea02f3aa6.py

config

{'base_model': 'Qwen/Qwen3.5-0.8B', 'lora_rank': 32, 'dataset': 'HuggingFaceH4/no_robots', 'split': 'train (eval on test, first 500)', 'n_examples': 4096, 'n_eval': 256, 'steps': 100, 'batch_size': 64, 'max_seq_len': 2048, 'lr': 0.0003, 'optimizer': 'dimuon', 'max_wall_seconds': 2700}

modal_call_id

fc-01M6FGCPTKMERT6BJDPK09BMTB

script_source

Claude (claude-opus-5-5) adapted the reference script