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:
parent
50f257d1a5
commit
c7e3989d72
2 changed files with 289 additions and 0 deletions
23
.github/workflows/mlx-ci.yml
vendored
23
.github/workflows/mlx-ci.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue