Small models, sparse experts

Problem

A language model spends most of its weights in feed-forward layers, and a dense model runs every one of them for every token. A mixture-of-experts (MoE) model replaces each feed-forward layer with many smaller ones, the experts, plus a router that sends each token to only a few. The model can hold far more parameters than it uses on any one token.

MoE models are usually trained on racks of datacenter GPUs. I wanted to know whether a multi-billion-parameter sparse model can be pre-trained on one desk machine, a Mac Studio, and what decides whether that is feasible.

How it works

1. Start from a small dense model. Gemma 3 270M from Google: 18 layers, hidden size 640, and feed-forward layers 2,048 units wide.
2. Drop-upcycle every feed-forward layer. Following the Marco-MoE method from Alibaba AIDC-AI, I split each layer into four slices of 512 units, copied each slice 64 times, scaled the weights by λ = 64^1/3 ≈ 4.0, and re-drew a random 10% of the weights in every copy except the first. That gives 256 experts per layer that start almost identical but not quite, so they can drift apart in training.
3. Add a router to each layer. It scores all 256 experts for each token, keeps the top four, and weights their outputs with a softmax over those four scores. Routers start near zero, drawn from N(0, 0.01), and learn from there.
4. Pre-train with MLX. The experts run through SwitchGLU from mlx-lm, which uses MLX’s fused gather_mm kernel. Training reads FineWeb-EDU, 512-token sequences, 16 to a batch, with AdamW and a cosine learning-rate schedule.

The cost moves from compute to memory. The upcycled model is about 9.2 GB on disk, and it has to sit in the Mac’s unified memory together with gradients and optimiser state while it trains.

Dispatch decides everything

The maths of an MoE layer is simple. The hard part is the plumbing: each token in a batch goes to a different set of four experts out of 256, and the GPU wants large, regular matrix multiplies. I benchmarked five ways of getting tokens to their experts on Apple Silicon, all on the same 18-layer, 256-expert shape.

• 6,874training tokens/s with gather_mm and SwitchGLU
• <3tokens/s with naive per-token multiplies
• 2.68 GBone projection’s copied weights, one layer, 2,048 tokens
• 93+ sGPU unresponsive before macOS rebooted

The naive ways crashed the Mac. Gathering each token’s expert weights before multiplying creates a tensor of shape [2,048, 512, 640] for one projection in one layer: 2.68 GB. There are three projections per layer and 18 layers. The GPU stopped answering for more than 93 seconds, the macOS watchdog raised a kernel panic, and the machine rebooted. It happened twice. The panic write-up is in the repo.

Computing every expert works, but wastes 64 times the arithmetic. The dense-expert approach runs all 256 experts on every token as one large multiply and masks out all but the top four. It avoids the giant copies and reaches 211 tokens/s.

Sorting helps, fusing wins. Sorting tokens by expert and padding each group to the same length turns dispatch into one batched multiply, which got PyTorch on MPS to 2,753 tokens/s. MLX goes further: mx.gather_mm reads each token’s expert weights in place inside one fused Metal kernel, with no large intermediate tensor. SwitchGLU from mlx-lm wraps it and sorts tokens by expert first. PyTorch’s MPS backend has no equivalent primitive, which is why the fastest path was MLX.

Tuning the training loop

The dispatch benchmark times the 18 MoE layers on their own. The full model adds attention and a 262,144-token vocabulary, so the real training loop started lower, at 2,777 tokens/s. Four changes more than doubled it.

• Compile the whole step. mx.compile over forward, backward and the optimiser update together, passing the model and optimiser state as inputs and outputs so MLX can track them.
• Gradient checkpointing. mx.checkpoint on each transformer layer recomputes activations in the backward pass instead of storing them. That freed the memory for bigger batches.
• Asynchronous evaluation. mx.async_eval lets Python build the next step’s graph while the GPU is still running the current one.
• Batch size. 8 to 16 sequences was the sweet spot on 96 GB. At 32, memory pressure halved throughput.

The run

The pre-training run read 52M tokens of FineWeb-EDU: 6,362 steps of 16 sequences of 512 tokens. It took about 2.6 hours, including one resume from a checkpoint.

• 11.74 → 3.56training loss over 52M FineWeb-EDU tokens
• 2.16×training loop speed-up, 2,777 → ~6,000 tokens/s
• ~2.6 hfor the whole run on one Mac Studio
• 5.7%of the 4.73B parameters active per token

Before training, the upcycled model produced gibberish. After it, the output was coherent English. That is the extent of the evaluation: I ran no downstream benchmarks, and the repo does not include samples.

At about 6,000 tokens/s, here is what larger budgets would cost on the same machine, from the README’s feasibility table:

Limits

• 52M tokens is about 1% of the roughly 5.4B that would be compute-optimal for 270M active parameters. This was a feasibility run, not a model to rely on.
• No downstream benchmarks were run on the trained model, and “coherent English” is a qualitative judgement, not a score.
• The training loss is plain cross-entropy with no load-balancing term, and I did not measure how evenly tokens spread across experts. Some experts may never be picked.
• The dispatch benchmarks time the MoE layers alone on random inputs, and their batch sizes are not identical (the naive and dense-expert runs used smaller batches to fit). Read the gaps as orders of magnitude, not exact ratios.
• One machine, one run. Only six points of the loss curve are published.
• The upcycling method is Marco-MoE’s. My work is the Apple Silicon implementation, the dispatch benchmarks and the training notes. This is an engineering write-up, not a peer-reviewed paper.

Credits

Drop-upcycling follows Marco-MoE (Alibaba AIDC-AI). Built on MLX and mlx-lm (Apple), Gemma 3 (Google) and FineWeb-EDU (Hugging Face). I have no affiliation with any of these organisations.

Read more

• GitHub repository: the MLX model, the upcycling script, the training script and all six dispatch benchmarks. MIT licence.
• GPU watchdog kernel panic: what the naive dispatch did to macOS, and why.
• Marco-MoE: the drop-upcycling method this follows.
• Writing Fast MLX: notes on compile, async evaluation and other MLX performance habits.