Long Yixing d918245834 Add MLX-aware public Unsloth trainer API (#6462) ... * feat: add mlx public trainer api * test: cover mlx public trainer api * fix: preserve mlx epoch trainer configs * fix: pass mlx warmup ratio through config * fix: align mlx trainer dataset order * fix: keep mlx chat templates import-light * fix: infer mlx trainer context length * fix: mirror cuda mlx context defaults * fix: align mlx notebook trainer defaults * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: keep mlx public helpers import-light * refactor: reuse mlx optimizer normalization * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: address mlx review feedback * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: tighten mlx training argument parity * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: align mlx trainer eos default * Fix MLX trainer to accept DataCollatorForSeq2Seq and handle TokenizerWrapper in get_chat_template * Trim redundant docstrings on internal MLX helpers * MLX review fixes: Studio optimizer import-safe on non-MLX hosts, preserve explicit max_length, skip MLX tests before import * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * MLX review round 2: defer max_length to model context, optimizer alias fallback for older zoo, skip non-MLX test on missing GPU deps * MLX review round 3: keep chat_templates importable without torch on MLX * fix: preserve MLX trainer notebook shims * fix: ignore CUDA tokenizer moves on MLX * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: harden MLX trainer shims * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: unwrap MLX scheduler enum args * fix: coerce integral MLX epoch counts * fix: spoof CUDA compatibility APIs on MLX * fix: harden MLX notebook compatibility shims * MLX: add torch.cuda.mem_get_info to the compatibility shim Notebook memory cells call torch.cuda.mem_get_info()[0] directly (not gated by is_available), so on MLX it raises without a shim. Return (free, total) bytes from the MLX device stats, consistent with the other torch.cuda compat helpers, and add a matching assertion to the compat-API test. * MLX: use active memory for mem_get_info; fix BatchEncoding.to keyword device Address review on the MLX compatibility shim: - torch.cuda.mem_get_info() now derives free bytes from current active MLX memory instead of the peak high-water mark, so a capacity check stays accurate after a transient spike or a prior run. - BatchEncoding.to(device=...) passed by keyword no longer forwards a positional None alongside the keyword (which raised "multiple values for 'device'"), so non-CUDA keyword moves like .to(device="cpu") delegate correctly. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * MLX: accept preserve_dataset_order; stub RL trainers with a clear error Two fixes so unmigrated notebooks behave predictably on MLX (torch present): - preserve_dataset_order is a real MLXTrainingConfig field but was missing from the extra-argument allowlist, so passing it (as a config or trainer kwarg) could be rejected as unknown on a zoo without the field. Add it to _MLX_IMPLEMENTED_EXTRA_ARGUMENTS so the documented no-shuffle path is reachable. - GRPO/DPO/ORPO (and KTO/PPO/Reward) have no MLX trainer yet. Retarget the ones the installed trl exposes to a stub that raises a clear 'not supported on MLX' error instead of importing the real torch/CUDA trainer and crashing deep inside it. Only existing trainers are retargeted (no invented attributes), idempotent across re-imports. * MLX: make RL-trainer stubbing import-safe; back current-memory APIs with active memory Address review on the MLX shims: - The RL-trainer stub loop probed trl with getattr(_trl, name), which triggers trl's lazy trainer import and pulls torch -- that can crash import unsloth on a torch-free MLX install just to check existence. Decide what to stub from trl.__all__ + already-materialized attrs (vars) instead; never resolve the real trainer. All trl trainer names are in __all__, so they are still stubbed (even torch-free), and the probe no longer imports torch. - torch.cuda.memory_reserved / memory_allocated (the current, non-max APIs) were aliased to peak max_memory_reserved. Back them with current active MLX memory so cleanup / capacity checks see live usage; max_* keep the peak high-water mark. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * MLX: keep TRL's SFTConfig epoch default under the trl.SFTConfig alias Unmigrated notebooks import SFTConfig from trl, which the MLX build aliases to the public training-args class. TRL/HF SFTConfig defaults to num_train_epochs=3 (max_steps=-1); the native MLX config defaults to max_steps=60. So an SFTConfig built without an explicit length silently ran 60 MLX steps instead of TRL's 3 epochs under the alias. Alias trl.SFTConfig to a thin subclass that seeds the TRL epoch default only when neither max_steps nor num_train_epochs is given; explicit lengths pass through untouched, and the native public args class keeps its MLX default. Epoch mode is supported by the MLX trainer. * MLX CI: keep the GGUF reload smoke under the job timeout The RELOAD-GGUF-via-llama-cli step timed out at 300s. BF16 GGUF decode is CPU-bound on the macOS runner (~10s+/token), so generating 24 tokens landed right on the 300s cliff and killed the process. This step is a save/reload integrity smoke (it only needs a few chars of output), so the token count is incidental: generate 8 tokens with explicit threads and a small headroom on the subprocess timeout, all env-tunable (UNSLOTH_GGUF_RELOAD_N / _THREADS / _TIMEOUT). Cuts the reload well under the 25 minute job budget. * MLX: broaden trainer stubs, real peak-memory reset, fix shim tests Address review on the MLX public API: - The SFTConfig identity tests asserted trl.SFTConfig is UnslothTrainingArguments, but the alias now points at the _MLXSFTConfig subclass that preserves TRL's epoch default, so the MLX suite failed before testing the shim. Assert issubclass instead. - torch.cuda.reset_peak_memory_stats was a no-op, so max_memory_reserved kept earlier model-load peaks across a scoped run. Wire it to mx.reset_peak_memory with the same core/metal fallback used for the reads. - The unsupported-trainer stubs were a fixed list, so trainers outside it (a newer RLOOTrainer) still routed to the real torch trainer. Derive the set from trl.__all__ (every non-SFT *Trainer) so all non-SFT surfaces fail with a clear MLX message; names come from __all__ so trl is never resolved. - The non-MLX export smoke skipped only on missing bitsandbytes/triton; other absent GPU deps (numpy/torch/unsloth-zoo, or _gpu_init re-raising ImportError) made it fail on CPU hosts. Skip on any ImportError. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: keep MLX notebook compatibility minimal * MLX CI: force CPU + small context for the GGUF reload smoke The RELOAD-GGUF-via-llama-cli step timed out even at 8 tokens (>420s), so it is a fixed hang, not per-token cost: on the paravirtual macOS runner GPU llama.cpp's Metal backend stalls, and the gemma3 GGUF advertises a 32768 context that llama-cli would otherwise fully allocate. Run llama-cli CPU-only (-ngl 0) with a small context (-c 256); keep generation short. All env-tunable (UNSLOTH_GGUF_RELOAD_NGL / _CTX / _N / _THREADS / _TIMEOUT). Also print llama.cpp's partial stdout/stderr on timeout so a future hang is diagnosable instead of an opaque TimeoutExpired. * MLX CI: export the reload-smoke GGUF as q8_0, not bf16 The GGUF reload via llama-cli timed out on the runner even CPU-only with a tiny context and 8 tokens. Root cause is the format, not the flags: the smoke exported quantization_method='not_quantized', which maps to a bf16 GGUF, and llama.cpp's bf16 CPU decode is unusably slow on the paravirtual macOS runner. Export q8_0 (fast_quantized, the exporter default and what users deploy) instead -- llama.cpp has optimized q8_0 CPU kernels, so the fresh-process reload loads and generates in seconds. The reload stays CPU-only (-ngl 0) with a small context. * test: clear TRL shim before availability check --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: imagineer99 <samleejackson0@gmail.com>		2026-07-02 23:02:26 +01:00
..
__init__.py	Consolidate dual venvs and separate install from update (#4530)	2026-03-25 05:24:21 -07:00
conftest.py	Formatting: ruff line-length 100, kwarg-spacing passes, drop blank after short local imports (#6079)	2026-06-08 04:24:13 -07:00
test_change_system_message.py	Reduce comments across recent fixes (#6776)	2026-06-30 23:13:36 -07:00
test_construct_chat_template_validation.py	Fix construct_chat_template leaking {INPUT}/{OUTPUT} sentinel into the chat template (#6531)	2026-06-24 00:40:53 -07:00
test_cpo_processor_text_tokenizer.py	Fix CPOTrainer crash with multimodal processors (Gemma 4) (#6522)	2026-06-23 06:29:22 -07:00
test_cross_platform_parity.py	Reduce and tighten comments and docstrings across the test suite (#6429)	2026-06-18 01:07:09 -07:00
test_dpo_vision_processor_passthrough.py	Formatting: ruff line-length 100, kwarg-spacing passes, drop blank after short local imports (#6079)	2026-06-08 04:24:13 -07:00
test_e2e_no_torch_sandbox.py	Studio: remove OpenEnv and other unused packages (#6585)	2026-06-23 07:20:47 -07:00
test_fast_language_model_text_only.py	Reduce and tighten comments and docstrings across the test suite (#6429)	2026-06-18 01:07:09 -07:00
test_fast_model_config_passthrough.py	Fix FastModel config passthrough for sequence classification (#6203)	2026-06-12 11:15:37 +03:00
test_fast_sentence_transformer_redirect_lifecycle.py	Reduce and tighten comments and docstrings across the test suite (#6429)	2026-06-18 01:07:09 -07:00
test_flash_attn_install_python_stack.py	Formatting: ruff line-length 100, kwarg-spacing passes, drop blank after short local imports (#6079)	2026-06-08 04:24:13 -07:00
test_gpu_init_ldconfig_guard.py	Formatting: ruff line-length 100, kwarg-spacing passes, drop blank after short local imports (#6079)	2026-06-08 04:24:13 -07:00
test_grpo_ddp_model_config.py	Fix GRPO logit scaling when model is wrapped by DDP (#5955)	2026-07-02 00:01:15 -03:00
test_install_python_stack.py	Make _uv_safe_path space-safe on macOS/Linux (#6503) (#6534)	2026-06-24 04:02:24 -07:00
test_install_uv_override_space.py	Installer: make UV_OVERRIDE space-safe on Apple Silicon (#6503) (#6639)	2026-06-24 17:34:18 -07:00
test_mlx_public_trainer_api.py	Add MLX-aware public Unsloth trainer API (#6462)	2026-07-02 23:02:26 +01:00
test_no_torch_filtering.py	Studio: remove OpenEnv and other unused packages (#6585)	2026-06-23 07:20:47 -07:00
test_orpo_processor_text_tokenizer.py	Reduce and tighten comments and docstrings across the test suite (#6429)	2026-06-18 01:07:09 -07:00
test_pad_token_fix.py	Keep pad-named pad_tokens (e.g. <|vision_pad|>); fix Qwen3-Base load crash (#6652)	2026-06-25 04:41:09 -07:00
test_patch_trl_rl_trainers_defensive.py	Reduce and tighten comments and docstrings across the test suite (#6429)	2026-06-18 01:07:09 -07:00
test_studio_import_no_torch.py	Add Hugging Face dataset streaming mode to Studio (#4946)	2026-06-22 17:48:18 +03:00
test_tokenizers_and_torch_constraint.py	Reduce and tighten comments and docstrings across the test suite (#6429)	2026-06-18 01:07:09 -07:00
test_unsloth_run_tool_policy_resolver.py	Keep server-side tools enabled under --secure (#6403)	2026-06-18 05:52:40 -07:00
test_v100_fullft_precision.py	Fix full finetuning precision on V100 / no-bf16 GPUs (#5880)	2026-06-29 18:35:23 -03:00
test_vision_lora_targeting.py	MLX Training updates (#5656)	2026-06-14 04:58:50 -07:00