[mitmul]

Summary

mlx-integration-test failed on main after #38 (run 42134257749). The test reported only Timed out waiting for the MCP server to become ready. after 15 minutes. The server had in fact died within a minute of starting, while loading the model:

plamo3_runtime.py:91  mx.eval(self.model.parameters())
RuntimeError: [METAL] Command buffer execution failed: Caused GPU Timeout Error (00000002:kIOGPUCommandBufferCallbackErrorTimeout).

This PR makes such failures visible at once, and keeps the evaluation that failed off long GPU command buffers.

1. The test hid the crash

The integration tests started the server with stdout=PIPE, stderr=PIPE and never read either pipe. They also never checked whether the server was alive, so they polled a dead port until the timeout.

The tests now:

• write the server's output to a file;
• fail as soon as the server exits, showing its output;
• on a timeout, abort the server with PYTHONFAULTHANDLER=1, so the failure shows where each thread was stuck.

The traceback above was captured this way.

2. Metal aborted the load's largest command buffer

The CI runners are virtual 3-core M1s with 7.5 GB of memory, and they compress memory heavily while the model loads. The CLI also sets MLX_MAX_MB_PER_BUFFER=1000000, so that decoding does not end a command buffer after every matrix. During loading, this lets one evaluation become one command buffer.

The evaluation after the layer loop converted both the embedding and lm_head, each 107,520 × 2,560 (550 MB in BF16). That is about 2.4 GB of memory in one command buffer, five times a layer's.

The crash is rare: it happened in 2 of about 140 natural server starts on the runners. Load experiments on the runners isolated the cause:

Load experiment on the runners	GPU timeouts
Whole model in one evaluation, CLI buffer limits	10 / 92
Whole model in one evaluation, MLX's default limits (40 MB)	0 / 52
A layer per evaluation (~0.4 GB), all experiments	0 in over 3,000
Embedding and lm_head evaluated apart, on the GPU	1 / 10 starts

Ruled out: the GPU waiting for disk reads. MLX loads safetensors on the CPU, and a 43 s GPU wait on CPU work passed on the runners.

Two other designs made things worse, so they are not used here:

• casting every layer on the CPU inside the GPU evaluation;
• async_eval per parameter.

The fix

• Embedding and lm_head casts: load_model casts both on the CPU. BF16 → FP16 gives bit-identical values there (checked on all 243 tensors), takes 0.08 s, and has no Metal time limit.
• Order of evaluation: Plamo3Translator evaluates these casts on the CPU alone, after the layers.
• lm_head quantization: it is quantized on the GPU in blocks of 13,440 rows, a third of a layer each, then joined on the CPU. Quantization groups never span rows, so the result is identical. CPU quantization was rejected: it differs from the GPU's and is 10× slower.
• Unchanged: the layer loop, which never failed, and decoding.

Verification

• Parameters: all loaded parameters are bit-identical to main for 4bit, mixed, 8bit and bf16 (437 / 437 / 437 / 243 tensors, compared by MD5).

• Translations: identical for 4bit, mixed and bf16.

• Memory: peak is unchanged (2.57 GB for 4bit, 6.0 GB for bf16). Active memory after loading equals the parameter bytes.

• Load time: unchanged within noise (4bit: main 4.6–5.4 s, this PR 4.6–5.0 s on M1 Max).

• Tests: mock tests pass (121), and so do the real-model MLX tests (38).

• On the CI runners:

  • The real-model tests passed in 10 runs with this code.
  • Alternating with main for 80 server starts each: no failures for either.
  • Under 4 GB of incompressible memory pressure, 30 loads each: no timeouts for either.

  main did not fail in those runs either, so they show no regression rather than a measured improvement. The failing evaluation is gone by construction: no load command buffer is larger than a layer's.

Not in this PR

• Main-only job: mlx-integration-test runs only on pushes to main, so this failure surfaced after merging. Running it on PRs that touch servers/mlx/ would catch such problems earlier.
• PLaMo 2: the PLaMo 2 loader still converts and quantizes the whole model in one evaluation. No failure has been seen there, and CI does not exercise it.
• Hub request: each server start still asks the Hugging Face Hub for config.json (is_plamo3_checkpoint), which logs an unauthenticated-request warning.