🔬Modal Research Exchange

{
  "base_model": "Qwen/Qwen3.5-2B",
  "baseline": "Arm A: the same fork, TriMuon at FP32-core, identical seed and data order, run twice with the default (non-deterministic) kernels.",
  "batch_size": 32,
  "budget": {
    "gpu_hours": 1.5,
    "max_wall_minutes": 90
  },
  "command": "# custom tinker SDK script against the fork; TriMuonOptimizerConfig/deterministic are fork-added\nimport os, tinker\nsc = tinker.ServiceClient(base_url=os.environ[\"TINKER_BASE_URL\"])\nfor det in [False, False, True, True]:\n    tc = sc.create_lora_training_client(base_model=\"Qwen/Qwen3.5-2B\", rank=32, seed=0,\n        optimizer=tinker.types.TriMuonOptimizerConfig(precision=\"fp32_core\", deterministic=det))\n    for step, batch in enumerate(fixed_batches):  # same precomputed no_robots batches, bs=32, 50 steps\n        out = tc.forward_backward(batch, loss_fn=\"cross_entropy\").result()\n        st = tc.optim_step(tinker.types.TriMuonParams(learning_rate=1e-4)).result()\n        log(det, step, out.metrics[\"loss\"], st.metrics.get(\"grad_norm\"))\n    path = tc.save_weights_for_sampler(name=f\"trimuon-det{det}\").result().path\n    eval_heldout_nll(tc, test_200)  # forward-only NLL on 200 no_robots test examples",
  "dataset": {
    "hf_path": "HuggingFaceH4/no_robots",
    "n_examples": 1600,
    "name": "no_robots",
    "split": "train (fixed 1600-example order); test for eval"
  },
  "eval": {
    "benchmark": "no_robots test held-out NLL plus per-step training traces",
    "metric": "first step where the two replays' loss or grad-norm differ; absolute difference in held-out NLL between the replays",
    "n_examples": 200
  },
  "hypothesis": "With TriMuon at FP32-core precision, two identical-seed replays run with deterministic kernels stay bitwise identical in loss and grad-norm for all 50 steps, while default-kernel replays diverge by step 6 or earlier.",
  "lora_rank": 32,
  "max_seq_len": 2048,
  "method": "Needs a Spindle fork. The fork adds (1) the TriMuon optimizer at FP32-core precision, which is the middle row of Aaron's table, and (2) a model-creation flag that enables torch.use_deterministic_algorithms(True), CUBLAS_WORKSPACE_CONFIG=:4096:8, deterministic attention and NCCL settings, and fixed reduction order. Deploy a throwaway copy of the fork. Train a rank-32 LoRA on Qwen/Qwen3.5-2B with SFT on no_robots for 50 steps. The data order is fixed by a precomputed index list, and the LoRA initialization uses the same seed in every run. Each arm runs twice with identical inputs. Arm A uses the default kernels and arm B uses the deterministic kernels, and nothing else differs between them. Log per-step loss and grad-norm for every run. Record the first step where the two replays differ, and the held-out NLL difference between the two replays' final checkpoints on the same 200 examples. If arm B is bitwise stable, the drift comes from kernel non-determinism. If arm B still drifts at a similar magnitude, the drift comes from precision or from client-side and batching sources. A precision sweep under deterministic kernels is the follow-up.",
  "optimizer": {
    "lr": 0.0001,
    "name": "trimuon (fork)"
  },
  "steps": 50,
  "success_criterion": "Validated if arm B's two replays match exactly at every step (first divergence > 50, NLL difference = 0) and arm A diverges by step 6 or shows an NLL difference \u2265 1e-4. Failed (precision or another source) if arm B diverges at step \u2264 10 or its NLL difference is \u2265 1e-4.",
  "title": "Deterministic kernels vs default kernels for TriMuon replay drift",
  "training_type": "sft"
}

armA_first_divergence_step

1

armA_first_loss_diff_step

2

armA_first_gradnorm_diff_step

1

armA_abs_heldout_nll_diff

2.959e-05

armA_max_abs_loss_diff

0.002083

armA_max_abs_gradnorm_diff

0.6589

armA_heldout_nll_r0

2.429

armA_heldout_nll_r1

2.429

armB_first_divergence_step

1

armB_first_loss_diff_step

1

armB_first_gradnorm_diff_step

1

armB_abs_heldout_nll_diff

0.0001558

armB_max_abs_loss_diff

0.004825

armB_max_abs_gradnorm_diff

1.046

armB_heldout_nll_r0

2.429

armB_heldout_nll_r1

2.429

armA_mean_step_seconds

3.681

armB_mean_step_seconds

4.672

trimuon_polar_mean_rel_error_vs_svd

0.0005536