GitHub - fumishiki/rosetta-moe

Same MoE Transformer. 4 languages. 22 benchmark scenarios. One hardware.

Rust, Go, Python, Julia — each implements the same MoE Transformer from scratch (forward, backward, optimizer, inference). All matmul hits the same Apple Accelerate BLAS on the same M1. No ML frameworks — every gradient is hand-derived.

Demo

Results

Rust leads T4 inference throughput and kernel softmax. Julia leads single-thread forward and training. Both hit ~600 GFLOPS on 64x64 matmul in this run. At hidden=256, forward spread narrows to 1.46x.

Full results with 22 scenarios across 5 axes: docs/bench-results.md

Quick Start

make test          # test all 4 languages
make bench         # benchmark + summary table
make convergence   # verify loss convergence
make convergence-plots  # run each language + render convergence demo GIF

One-tap verification (macOS only):

make test-rust / make test-go / make test-python / make test-julia
make bench-rust / make bench-julia

# First-time setup
cd python && pip install -e ".[dev]" && cd ..
cd julia && julia --project=. -e 'using Pkg; Pkg.instantiate()' && cd ..

Why This Exists

What Makes This Different

5 Research Axes

Architecture

All 4 implementations share the same model contract:

Input token_ids [batch, seq]
  → Embedding lookup                    [batch, seq, hidden]
  → N × TransformerBlock:
      RMSNorm → MQ Attention (RoPE)   → Residual add
      RMSNorm → MoE (Router + SwiGLU) → Residual add
  → Final RMSNorm
  → Linear (LM head)                   [batch, seq, vocab]

Model Specification (benchmark config)

Equation-to-Code Traceability

Every math formula in the model is mapped to concrete source code:

All source files are annotated with inline comments mapping math notation to code. See each language's README for the full table.

Technology Stack

All matrix multiplications route through Apple Accelerate's cblas_sgemm, which uses the AMX coprocessor on Apple Silicon. Non-matmul operations (softmax, RMSNorm, attention masking, MoE routing, AdamW optimizer) are hand-written loops — these are where language differences show up most.

Key Findings

1. BLAS levels the playing field. All 4 languages hit the same AMX — the language only matters for non-BLAS code.
2. GC was effectively invisible in this run. Train-step gc_throughput is ~1.000 for all 4 languages.
3. Julia keeps the fastest train step. 0.99ms vs Rust 1.29ms in h=64 mem-train scenario.
4. Rust keeps the strongest inference throughput. 4,579 inf/s at T4.
5. Language gaps shrink at larger hidden sizes. Forward spread drops to 1.46x at h=256.
6. RSS still diverges strongly. Rust ~20MB vs Julia ~536MB in warm forward scenario.

Why Julia Beats Rust at Training

Rust leads forward (0.56ms vs 0.59ms) — but Julia leads training (0.98ms vs 1.44ms). The reversal comes from backward pass implementation.

Before: Julia 7.36ms (Rust was already ~1.4ms — Rust winning by 5x)
After:  Julia 0.98ms (Rust 1.44ms — Julia winning by 1.5x)

The fix was one pattern change — @simd → @. broadcast:

# Before: @simd — 9K-11K heap allocations per loop (Julia 1.12 SIMD lane metadata)
@inbounds @simd for i in eachindex(grad)
    grad[i] = grad[i] * mask[i]
end

# After: @. broadcast — compiler fuses into single loop, zero intermediate allocation
@. grad = grad * mask

Why Rust can't match this. Rust's backward is zero-alloc (1.4 MB/step) — every gradient buffer is pre-allocated. But without broadcast fusion, each element-wise operation runs as a separate loop. Julia's compiler fuses grad = grad * mask * scale + bias into a single pass over memory. Rust would need a hand-written fused kernel for each combination — possible, but trades developer velocity for performance.

The insight: Zero allocation is not the same as zero overhead. Julia's broadcast fusion eliminates not just allocations but also loop overhead, cache misses from multiple passes, and function call boundaries. This is why a GC language beats a zero-GC language on training — the compiler does work that the programmer can't practically do by hand for every operation combination.

Rust still leads inference (no backward pass), all non-BLAS kernels (softmax 1.83us vs 4.83us), and memory footprint (20MB vs 523MB).

When to Choose What

Scaling Behavior

As hidden size grows, BLAS (AMX) share of total compute increases and language overhead shrinks. At h=256, forward spread already halves from 4.2x to 1.9x.

How Much GPU Time Does Python Waste?

Our benchmark isolates the overhead gap outside BLAS compute — the host-language tax on every training step. All 4 languages call the same cblas_sgemm, so the difference is pure language overhead: interpreter dispatch, GC pauses, FFI bridge cost, and memory management.

h=64, train step, Apple M1 (AMX) measured. BLAS = step time − overhead. GPU utilization = BLAS / step time.

Note: Julia's 98% utilization reflects broadcast fusion (@.) on M1 AMX — but this is not merely a CPU trick. Julia's advantage is language-level fusion compilation: the compiler fuses multiple element-wise operations into a single pass, eliminating intermediate allocations. On GPU, this same capability manifests as Reactant.jl (XLA backend) — whole-graph compilation and automatic kernel fusion without manual CUDA code. Rust has no equivalent automatic fusion; GPU acceleration requires hand-written CUDA kernels or frameworks like Burn. Rust's 32% overhead here is dominated by per-parameter gradient update allocations (a CPU memory access pattern), and its zero-copy abstraction (&mut [f32], no GC, no runtime) would close the kernel launch overhead gap on GPU — but not the fusion gap. These figures are M1 CPU results; treat the absolute numbers as reference, not GPU predictions.

Python spends 90% of every training step waiting for the CPython interpreter. The GPU sits idle. At h=64 this means 10.4x slower training than Julia for the same matmul work.

Switching from Python — what you save at 13T training scale:

Based on measured h=64 M1 CPU overhead ratios — reference estimates, not production GPU predictions. On production GPUs, Rust's kernel launch overhead would match Julia's (both zero-GC, direct memory control), but Julia retains an architectural advantage: Reactant.jl compiles the full compute graph into fused XLA kernels automatically, whereas Rust requires manual kernel fusion or framework support. The Python vs {Rust, Julia, Go} gap is the robust signal; the Rust vs Julia gap is M1-specific and would narrow (but not necessarily close) on GPU.

In concrete terms: a GPT-4-class run (13T tokens, $100M) wastes **$90M on the CPython interpreter** — enough to fund the entire training run again in Julia or Rust.

"But torch.compile fixes this"

Yes — and that's exactly the point. The Python ML ecosystem has spent enormous engineering effort building layers to work around the language:

• CUDA Graphs — bypass Python by replaying GPU command buffers directly
• torch.compile / XLA — trace the compute graph to eliminate per-op Python dispatch
• C++ DataLoaders — hide the GIL behind multiprocessing
• NCCL — GPU-direct all-reduce so Python never touches gradient data

These are not features of Python. They are escape hatches from Python. Every one of these optimizations exists because CPython is too slow to be in the critical path. The $90M waste is real — it's just that PyTorch has learned to avoid paying most of it by doing the real work in C++/CUDA.

The question this benchmark answers: what if the host language were fast enough that you didn't need escape hatches?

Loss Convergence

All 4 implementations converge strongly (from ~7.2-7.9 to <=0.025, Rust to ~1e-4), confirming gradient flow and optimizer behavior. Run make convergence to verify. Detailed engineering revision log: docs/convergence-revision-history.md

Generate convergence artifacts and the animated demo (runs each language sequentially):

Verification Environment

Platform requirement: This benchmark/verification flow is designed for macOS (Apple Silicon) with Apple Accelerate. Cross-platform runs are possible for some tasks, but published benchmark numbers and make verify expectations are macOS-based.

Known Limitations

• Small model size: hidden=64 amplifies per-call overhead. At hidden=256, BLAS share already grows significantly (see Scaling Behavior above).
• Single hardware: Apple M1 only. Results may differ on x86 (no AMX), NVIDIA GPU, or different Apple Silicon generations.
• CPU only: No GPU benchmarks. At scale, GPU compute (A100: ~312 TFLOPS f32) dwarfs CPU (~1.5 TFLOPS).
• alloc_bytes not comparable: Each language measures allocation differently. Use peak_rss_bytes for cross-language comparison.

Project Structure

rosetta-moe/
├── rust/                 # Rust implementation (core + benches + tests)
├── go/                   # Go implementation (CGO Accelerate bridge)
├── python/               # Python implementation (NumPy + ProcessPoolExecutor)
├── julia/                # Julia implementation (AppleAccelerate.jl + Threads)
├── benchmarks/
│   ├── *.json            # Benchmark summaries per language
│   └── convergence/      # Per-language convergence JSON outputs
├── scripts/
│   ├── summary.py                # Benchmark summary table generator
│   ├── check_docs.py             # Docs-vs-JSON consistency checks
│   ├── convergence_python.py     # Python convergence runner
│   ├── convergence_julia.jl      # Julia convergence runner
│   └── convergence_plots.py      # Sequential run + animated demo generation
├── docs/
│   ├── spec.md                       # Spec: requirements and acceptance criteria
│   ├── bench-results.md              # Full benchmark report (22 scenarios)
│   ├── convergence-revision-history.md # Detailed optimization/change history
│   └── assets/convergence/
│       └── convergence-demo.gif      # README demo animation
├── .github/workflows/     # CI
├── Makefile              # make test / make bench / make convergence / make convergence-plots / make verify
└── Cargo.toml

Each language directory contains its own README with architecture overview, equation-to-code mapping, implementation notes, and performance caveats.

License

Licensed under CC BY-NC-SA 4.0. See LICENSE for details.

⚠️ Usage Restrictions

This content is available for individual learning purposes only.

The following uses require explicit prior permission:

1. Any organizational use (commercial or non-profit)

   • Corporate/institutional training, educational curricula, internal wikis
   • University courses or research institution seminars
   • Non-profit organization training programs
   • Rationale: Organizational use often results in removed attribution links and unauthorized modifications

2. Paid courses, bootcamps, seminars, or information products

   • Distribution at paid events, displaying screenshots in lectures, or creating derivative educational materials

3. LLM/AI model training data

   • Scraping or using this content as a knowledge source for pre-training, fine-tuning, or RAG systems in commercial models

4. Any form of monetization or commercialization without explicit permission

   • Including but not limited to: paid articles, e-books, video courses, subscription platforms (Patreon, note, Medium, etc.), or any derivative works sold for profit

Individual use includes:

• Personal study and research
• Creating personal notes (for personal use only)
• Sharing the original article link with friends

For organizational use, you must:

• Contact the author before use
• Maintain all attribution links
• Notify the author of how the content is being used

Unauthorized use may result in billing for usage fees and public disclosure on social media.