Show HN: Mixlab, an ML arch lab in Go. JSON config, Metal and CUDA, 1.6s builds

Show HN: Mixlab, an ML arch lab in Go. JSON config, Metal and CUDA, 1.6s builds(github.com)

2 points by mrothroc 25 days ago | 1 comment

I built a tool for quickly testing different ML architectures. Define a model in JSON, train on your Mac (Metal) or ship the same config to a cloud GPU (CUDA). No code changes between platforms.

Why: I wanted to compare attention vs Mamba vs GQA at different parameter budgets without writing PyTorch for each experiment. Edit a JSON config, hit enter, see loss numbers. It will race different configs for you. The number one goal is iteration speed.

JSON config lets you chain together common ML blocks (attention, GQA, mamba, RetNet, and several more) and optimizers (muon, adamw) and compiles them to MLX IR, which can either run on Metal or CUDA backends.

Why Go: 1.6s builds, built-in profiling (mixlab -cpuprofile gives you a flame graph), import-based extensibility for custom blocks. No C++ extensions, no custom build systems. And personally I prefer strongly-typed, compiled languages.

On a Shakespeare benchmark matching nanoGPT (6L, 6H, d=384, 10.8M params): val loss 1.5527 on M1 Max, 1.5588 on A40. PyTorch numerical parity confirmed to 8 decimal places.

brew install mrothroc/tap/mixlab

https://github.com/mrothroc/mixlab

mrothroc 25 days ago |

Simple example to show how configs are defined:

{ "name": "plain_3L",

  // Minimal causal transformer baseline: 3 attention layers plus 3 SwiGLU layers.
  "model_dim": 128,
  "vocab_size": 1024,
  "seq_len": 128,

  // Blocks execute sequentially, alternating token mixing and feed-forward mixing.
  "blocks": [
    {"type": "plain", "heads": 4},
    {"type": "swiglu"},
    {"type": "plain", "heads": 4},
    {"type": "swiglu"},
    {"type": "plain", "heads": 4},
    {"type": "swiglu"}
  ],

  // Slightly longer than smoke-test configs so the baseline loss moves visibly.
  "training": {
    "steps": 200,
    "lr": 3e-4,
    "grad_clip": 1.0,
    "weight_decay": 0.01,
    "seed": 42,
    "batch_tokens": 1024
  }
}