Mlx-smolvla: SmolVLA checkpoint execution and validation on Apple Silicon

Hmm… I don’t have a Mac, so I can’t test Metal directly, but for now:

On the validation split: yes, the strict CPU / production Metal separation makes sense to me, as long as the guarantees of the two lanes stay explicit and separate.

I would think of it as a small validation ladder rather than one global definition of “same”:

1. the serialized checkpoint/processor contract is coherent;
2. MLX CPU matches the PyTorch reference under controlled inputs/noise;
3. Metal’s direct numerical divergence is measured separately;
4. an offline functional/non-regression metric checks whether that divergence actually degrades the policy output;
5. eventually, closed-loop robot behavior is a separate layer again.

That also avoids making bitwise CPU↔accelerator equality the production requirement. PyTorch makes the same general warning in its numerical accuracy notes: mathematically equivalent floating-point computations are not guaranteed to be bitwise identical across devices/backends.

So I would keep the strict negative Metal results rather than “fixing” them by widening the tolerance. A failed strict gate is still useful evidence; the production lane can answer a different question.

There is also one concrete fine-tuned-checkpoint edge that may be useful for the other question you asked.

I found a public SmolVLA fine-tune, jackvial/so101_smolvla_pickplaceorangecube_e100, whose preprocessor contains an additional step before normalization.

Its policy_preprocessor.json therefore declares the normalizer state as:

policy_preprocessor_step_6_normalizer_processor.safetensors

and that file actually exists. The processor configuration also says STATE/ACTION use MEAN_STD.

This matters because LeRobot does not appear to treat “normalizer == step 5” as the semantic contract. In the current processor/factory.py, policies explicitly compose processor steps in their own order, and the source describes that step order as a Hub-serialized contract. Checkpoints likewise contain policy_preprocessor_step_*.safetensors / policy_postprocessor_step_*.safetensors, rather than promising one universal normalizer index; see the current checkpoint layout docs.

The serialized processor JSON already carries the actual state_file, and LeRobot’s processor/pipeline.py uses that serialized state information when reconstructing the pipeline.

By contrast, in the mlx-smolvla code path I checked, processor state discovery is still based on the canonical fixed filenames around:

with the usual step-5 preprocessor normalizer / step-0 postprocessor unnormalizer layout.

I tested this part on Linux/CPU only, using a metadata fixture matching the published step-6 layout. The result was:

| layout | current fixed lookup | lookup via serialized state_file | | canonical step-5 normalizer | mean_std | mean_std | | published-style step-6 normalizer | identity | mean_std |

So this is not a Metal repro, and I also would not call it a proven full-checkpoint failure yet. The narrower finding is:

a checkpoint can look acceptable at the core-config level while a valid serialized state normalizer lives at a different processor step, causing the current fixed filename lookup to treat state normalization as identity.

That looks like a useful compatibility edge because it can be silent rather than producing an obvious load exception.

My default route here would be to make the serialized processor contract authoritative for current-format checkpoints:

• read policy_preprocessor.json /policy_postprocessor.json ;
• identify the relevant normalizer_processor /unnormalizer_processor step;
• follow that step’s declared state_file ;
• include that declared file in the Hub download set;
• then keep the existing tensor/header/feature-key checks after resolving the file.

I would not just change the magic number from step 5 to step 6, since the whole point is that the index depends on the composed processor pipeline.

A useful failure split might be:

• processor JSON + declared state file present → use the declared file;
• legacy checkpoint without processor JSON → use a clearly separate legacy/migration path;
• processor JSON declares state, but that state file is missing → explicit diagnostic rather than guessing;
• state file is present, but feature/stat keys do not match → report that separately as a schema/stats mismatch.

That last distinction may be worth preserving because LeRobot itself currently warns that processor/normalization mismatches can run without raising an error while silently damaging results; see Adding a Policy. There is also a recent concrete example in LeRobot #4415, where normalization/unnormalization can be silently skipped because the serialized stats keys do not match the runtime feature key.

So for “compatible checkpoint” reports, I think the processor artifacts are part of the compatibility contract, not just ancillary files next to model.safetensors.

A compact triage order might therefore be:

1. architecture/config supported;
2. required weight tensors/names/shapes supported;
3. processor JSON and its declared state files available;
4. camera/state/action schema and rename mapping compatible;
5. normalization mode + stats keys/shapes compatible;
6. only then investigate MLX CPU/Metal numerical behavior.

That should make future “this fine-tune doesn’t work” reports cheaper to classify.

#

Small public-checkpoint scan

I also did a metadata-only CPU scan across a convenience sample of public Hub repositories; no model weights or Apple execution were involved.

The scan saw:

- 86 candidate repositories;
- 72 with `config.json` identifying`type == smolvla` ;
- 68 accepted by the dependency-light `mlx-smolvla` config parser;
- 0 scan errors;

• 1 strong case where the core config was accepted, an active serialized normalizer state existed at another step, and the fixed lookup resolved state normalization to identity .

That strong case was the jackvial/... step-6 checkpoint above.

I would not interpret 1/86 as a prevalence estimate: the Hub search was not a census, and many of the candidate repositories were outside the native port’s full compatibility surface for unrelated reasons. The useful result is just that the edge exists in a real published SmolVLA checkpoint rather than only as a synthetic possibility.

I would keep this lane, but I think its meaning is slightly different from direct backend parity.

As I read scripts/statistical_check.py, it compares each backend’s first predicted action with the dataset action and then compares the resulting MAEs.

So approximately it answers:

“Does production MLX make this offline imitation-error proxy worse than the PyTorch reference?”

That is useful, but it is distinct from:

“How far did the MLX prediction itself move from the PyTorch prediction?”

I think keeping both concepts separate actually strengthens the validation story.

A cheap addition, if useful, would be to reuse the same deterministic 50-frame inputs/noise and record a direct MLX↔Torch action-difference distribution too — for example median / p95 / max absolute delta, perhaps per action dimension or over the whole action chunk.

I would not replace the existing 8-case direct gate or invent a new pass threshold from that distribution. It would mainly make the distinction visible:

- **direct numerical disagreement** ;
- **dataset-relative offline non-regression** .

Then if Metal differs numerically but the offline metric remains stable, that is a more informative result than either number alone.

#

Why I would keep strict numerical and downstream checks separate on M5

There are some recent MLX examples that make this distinction useful, although I would not assume they have the same root cause as mlx-smolvla.

In MLX #3897, batched/masked attention on M5 differs enough from the single-sequence path to fail a strict equivalence test, while the reported selected token remains unchanged. The same discrepancy is much smaller on M3 Max.

That is a good example of:

strict numerical mismatch → observed downstream choice still unchanged

But the opposite case exists too. MLX #3953 isolated a real float32 correctness problem in a broadcasted matmul path. A very simple invariant made the distinction possible: in the reported L=1 attention case, the output mathematically has to equal V, but it did not.

So “strict test failed” alone does not tell us whether a difference is harmless rounding, a different-but-acceptable backend path, or an actual correctness bug.

That is why the current separation into a reference lane and a production-behavior lane seems useful to me.

If you ever want one very cheap M5-side discriminator, MLX_ENABLE_TF32=0 might also be worth a single A/B run using the same fixed inputs/noise. This is only a diagnostic hypothesis, not a proposed root cause. A different MLX project reported M5 equivalence tests that failed normally but passed with MLX_ENABLE_TF32=0 in mlx-swift-lm #357. MLX has also had discussion around the flag and fp32 precision behavior, e.g. MLX #3860.

So I would use it only like this:

• default environment → record the existing Vision/action deltas + latency;
• MLX_ENABLE_TF32=0 before MLX initialization → repeat exactly;
• if the delta changes substantially, that narrows the kernel/precision-path investigation;
• if it does not, that hypothesis loses weight.

I would not infer from another project’s result that the current SmolVLA Metal delta is “a TF32 bug”.

#

Where offline parity eventually stops

One final boundary that I think your current wording already handles sensibly: offline action parity is not yet sustained robot task success.

That distinction matters especially for SmolVLA because deployment is chunked/closed-loop rather than a sequence of independent one-step predictions.

LeRobot’s current asynchronous inference documentation makes this explicit: actions_per_chunk, chunk_size_threshold, overlap aggregation, inference latency, and the state of the action queue can all change behavior. It even notes that larger action horizons can accumulate prediction error.

So if the project eventually wants a stronger production claim, I would see the next validation layer as paired/repeated rollout behavior under controlled initial conditions and the actual serving mode — not simply tightening the numerical tolerance further.

I would not make that a requirement for the current milestone; it is just the natural endpoint of the validation ladder.

So if I had to pick the lowest-cost next steps, I would probably keep the current strict/production split, preserve the existing 50-frame offline metric, optionally add a direct 50-frame MLX↔Torch delta summary, and make processor-state resolution follow the serialized state_file for current-format LeRobot checkpoints.

The processor-state case seems especially useful because it is independent of Apple hardware and gives a concrete compatibility condition that future checkpoint reports can test.

source & further reading

discuss.huggingface.co — original article Rakarrack-0.6.1 port making progress! ( AI assisted ) Cloud Storage Poll Welcome to Haiku basic(Haiku Docs, Haiku slide and Haiku sheets)