Settings

Theme

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

github.com

2 points by mrothroc 10 days ago · 1 comment · 1 min read

Reader

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

mrothrocOP 10 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
  }
}

Keyboard Shortcuts

j
Next item
k
Previous item
o / Enter
Open selected item
?
Show this help
Esc
Close modal / clear selection