Skip to content

Forward SFT gradient accumulation to MLX-LM#17

Open
kargarisaac wants to merge 1 commit into
ARahim3:mainfrom
kargarisaac:fix-sft-gradient-accumulation-forwarding
Open

Forward SFT gradient accumulation to MLX-LM#17
kargarisaac wants to merge 1 commit into
ARahim3:mainfrom
kargarisaac:fix-sft-gradient-accumulation-forwarding

Conversation

@kargarisaac

Copy link
Copy Markdown

Summary

  • forward SFTConfig.gradient_accumulation_steps into MLX-LM TrainingArgs
  • print the effective native-training gradient accumulation value in the SFT config log
  • add a mock-based regression test that verifies native SFT training passes the configured value through

Why

SFTTrainer already reads and stores gradient_accumulation_steps, but the native MLX training path did not pass it into mlx_lm.tuner.trainer.TrainingArgs, so MLX-LM used its default of 1.

This makes the Unsloth-compatible SFTConfig(gradient_accumulation_steps=...) behave as expected for native SFT training.

Tests

.venv/bin/python -m pytest tests/test_trainers.py -q
# 16 passed

.venv/bin/python -m py_compile mlx_tune/sft_trainer.py tests/test_trainers.py
git diff --check

Notes:

  • uv run pytest tests/test_trainers.py -q currently fails during dependency resolution before tests run because pyproject.toml allows Python >=3.9, while mlx>=0.31.0 requires Python >=3.10.
  • ruff check mlx_tune/sft_trainer.py tests/test_trainers.py reports pre-existing style issues in these files, so I did not broaden this bugfix PR into a formatting cleanup.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant