ci(mlx): real MLX training + inference smoke test on Mac M1

Add tests/studio/run_real_mlx_smoke.py and wire it into the macos-14
job as the final step. The script trains unsloth/gemma-3-270m-it
for 7 deterministic LoRA steps on an in-memory dataset of the SAME
row repeated:

    "<<HELLO!!>> My name is Unsloth!"

then prompts the trained model with "<<HELLO!!>> My name is " and
asserts the completion contains "Unsloth". Captures and asserts:

- per-step training loss (via MLXTrainer.add_step_callback);
- pre- and post-training loss + gradient norm (computed manually via
  mx.nn.value_and_grad over the training row, since MLXTrainer does
  not currently expose per-step grad norms);
- losses are finite, do not diverge, and post-train loss < pre-train;
- grad norms are finite and positive;
- the inference output contains "Unsloth".

Determinism: seeds python random, numpy, and mlx.core.random; passes
random_state=SEED to FastMLXModel.from_pretrained and
get_peft_model (both invoke _seed_mlx_random_state internally) and
seed=SEED to MLXTrainingConfig (drives batch shuffling). Uses fp16
+ no quant (gemma-3-270m is small enough to skip 4-bit) and LoRA
r=8 on the four attention projections.

This is the only place in CI that exercises a real MLX backward
pass + optimizer step + mlx_lm.generate call.
This commit is contained in:
Daniel Han 2026-05-07 03:39:55 +00:00
commit c7e3989d72
2 changed files with 289 additions and 0 deletions

View file

@ -28,6 +28,13 @@
# environment (the test fixture installs a MetaPathFinder that
# blocks `import mlx.core` for "no-mlx" profiles, faithfully
# simulating a Mac without mlx even when mlx IS installed).
# 4. End-to-end MLX training + inference smoke test:
# run_real_mlx_smoke.py trains unsloth/gemma-3-270m-it for 7
# deterministic LoRA steps on a single repeated text row, then
# verifies the trained model can complete the prompt and that
# losses + grad norms are finite and well-behaved. This is the
# only place in CI that exercises a real MLX backward pass +
# optimizer step + inference call.
#
# Three dispatch test files documented in tests/studio/README.md:
# - test_hardware_dispatch_matrix.py parametrized 7-profile matrix
@ -60,6 +67,7 @@ on:
- 'tests/studio/test_hardware_dispatch_matrix.py'
- 'tests/studio/test_is_mlx_dispatch_gate.py'
- 'tests/studio/test_mlx_training_worker_behaviors.py'
- 'tests/studio/run_real_mlx_smoke.py'
- 'tests/conftest.py'
- '.github/workflows/mlx-ci.yml'
push:
@ -194,3 +202,18 @@ jobs:
tests/studio/test_hardware_dispatch_matrix.py \
tests/studio/test_is_mlx_dispatch_gate.py \
tests/studio/test_mlx_training_worker_behaviors.py
# Real MLX training + inference smoke test. Trains
# unsloth/gemma-3-270m-it for 7 deterministic LoRA steps on a
# single repeated row ("<<HELLO!!>> My name is Unsloth!"),
# captures per-step losses and pre/post-training grad norms,
# then completes "<<HELLO!!>> My name is " and asserts the
# generation contains "Unsloth". This is the only place in CI
# that exercises the real MLX backward pass + optimizer step +
# inference path end to end.
- name: Real MLX training + inference smoke test
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
UNSLOTH_COMPILE_DISABLE: '1'
run: |
python tests/studio/run_real_mlx_smoke.py