3 commits

Author	SHA1	Message	Date
Daniel Han	cd3aef8a70	Bound dataset_num_proc by memory, and stop treating 1 as "no multiprocessing" (#7831) ... * Bound dataset_num_proc by memory, and stop treating 1 as "no multiprocessing" Training intermittently died with "One of the subprocesses has abruptly died during map operation", then succeeded on the next run (#2693, and a fresh Studio report). Measuring the tokenization map on an 8000-row dataset with a fast tokenizer found the mechanism: - Each pool task dill-pickles the tokenizer closure over a pipe, 5,369,755 bytes, once per worker per map. This happens under fork too: datasets does `from multiprocess import Pool`, and multiprocess/queues.py pickles every task regardless of start method. - Each worker peaks around 680 MB RSS. The old auto count was min(max(cpu_count + 4, 2), 64), so a large host forked up to 64 workers for roughly 43 GB resident. On a smaller box the OOM killer takes one and the parent reports only the generic message above, because datasets/utils/py_utils.py compares pool PIDs and never reads the child's exit status. Killing a worker with SIGKILL or SIGSEGV reproduces the error character for character, at any num_proc including 1. Two further defects made it worse: - The guard asked stdlib multiprocessing for the start method while datasets uses multiprocess, which keeps an independent default context. It was reading the wrong module. - num_proc=1 was used as the "no multiprocessing" sentinel. On datasets 4.3.0 (the Studio pin) map() takes the pool branch for any num_proc >= 1, so 1 still builds a Pool(1). Measured, num_proc=1 is 51% slower than None while buying no parallelism. Only None is in-process on every supported release. Changes: - New unsloth/utils/dataset_num_proc.py, one policy instead of four drifted copies. It asks multiprocess about the start method, caps the auto count at 8, and bounds any count, explicit ones included, by available memory at roughly 1 GB per worker over half of free RAM. Studio's explicit cpu_count // 4 previously bypassed every bound, which is how a 192-core host reached 48 workers. UNSLOTH_DATASET_NUM_PROC remains an uncapped escape hatch. - The config layer records intent and the map() call site makes it safe. These cannot be collapsed: unsloth_zoo reads a config None as "auto-size me", so writing None for a user who asked for 1 would inflate it. - worker.py no longer forces stdlib multiprocessing onto fork. It never reached Dataset.map, and Linux already defaults to fork. - A dead worker now raises with the start method, the worker count, the approximate memory cost and the escape hatch, chained from the original. No CUDA guard: 300 forced-fork map() runs on an initialized CUDA context produced no failures, and the child only runs the tokenizer. Since detect_hardware() always initializes CUDA, such a guard would cost every CUDA run its tokenization parallelism for no measured benefit. Known gap: the 1 -> None normalisation reaches SFT only, since that is the path sft_prepare_dataset owns. DPO, KTO, CPO, ORPO, Reward, PRM, PPO and BCO read args.dataset_num_proc in their own _prepare_dataset, so an explicit 1 there still builds a Pool(1). The memory bound does apply to all of them, so the OOM mechanism is covered everywhere. Reported by Eyera, who traced it to the commit and the call chain. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Stop UNSLOTH_DATASET_NUM_PROC=0 from inflating the worker count The env override returned before the serial encoding, so at the config layer it wrote None. unsloth_zoo.sft_prepare_dataset reads a config None as "auto-size me" and re-derives with its own uncapped min(max(cpu+4,2),64). So a user hitting a dead worker, following the escape hatch the new diagnostics message tells them to use, could get 64 workers. The hatch did the opposite of what it advertises, in exactly the OOM scenario it exists for. Affected 0, none, null, false, "" and 1: 14 of 42 config-layer cells. Measured end to end with the num_proc anchor skipped: before: requested auto / explicit 64 / Studio cpu//4 -> config None -> up to 64 after: all three -> config 1 -> bounded Also make the unsloth_zoo num_proc anchor non-required. It is the one anchor whose absence is harmless: the Zoo then reads args.dataset_num_proc, which the config layer has already bounded, so the memory ceiling still holds. unsloth_zoo is a floor dependency rather than a pin, and this is the block whose policy this branch changes, so it is likelier than the others to drift upstream. Hard-failing every install to recover an optimisation is the wrong trade. This is also what made the bug above reachable, so the two belong together. Verified across 630 cells: Python 3.10-3.14 x datasets 3.4.1-4.3.0 x fork/spawn/forkserver x four memory levels. Policy identical throughout. Zero config-layer None cells remain. 52 tests pass; reverting either fix fails them. Worth recording from that matrix: on Python 3.14 stdlib multiprocessing defaults to forkserver while multiprocess still defaults to fork, so the two disagree. Reading multiprocess, as this branch now does, is what matches what datasets will actually do; the old stdlib read would have silently disabled multiprocessing that was available. * Bound train_on_responses_only, the most-travelled map() in the library unsloth_zoo.dataset_utils.train_on_responses_only is a third copy of the same heuristic, and nothing in this branch touched it. unsloth re-exports it verbatim from chat_templates, it appears in essentially every Unsloth SFT notebook, and it is how Studio's apply_completion_masking reaches a map(). So the up-to-64-workers exposure sat on the most-travelled path while the branch fixed the quieter ones. Measured on a 192-core host with datasets 4.3.0, auto over a large split: 64 workers before, 8 after. Two constraints shaped the wrapper, both discovered before writing it: None cannot mean "in-process" at this boundary. The zoo's first act is `_num_proc_was_auto = num_proc is None or ...`, so a None arriving from outside reads as "size it for me" and triggers the very heuristic being bounded. Passing None for a caller who asked for 1 would inflate to 64, the same class of bug as the UNSLOTH_DATASET_NUM_PROC=0 one fixed earlier in this branch. So 1, not None, is the serial value here. The consequence is that an explicit 1 still builds a Pool(1) on datasets >= 4.0, unchanged from the raw zoo, so this is not a regression but the branch's "1 -> in-process" claim does not extend here. An explicit count also disables the zoo's 5000-row guard, which it applies to train_dataset and eval_dataset independently and only when it chose the count itself. Substituting unconditionally would hand workers to a small eval split that never had them. So substitute only when some split is actually at or above the threshold, which keeps the guard wherever it was doing work. A drift canary reads the zoo's constant off disk so the duplicated 5000 cannot diverge silently. 17 tests. Pool spy on a real masking run confirms a small auto dataset still builds no pool. Mutations: threshold drift 1 fail, explicit-serial returning None 2 fails, auto ignoring split size 6 fails. Known hole: when any split has no length (IterableDataset), this returns None and lets the zoo decide, so a huge train split beside a streaming eval split stays unbounded. Correct per the zoo's own rule, but it is a hole. * Keep the config layer serial on spawn, and stop an unsized split hiding a sized one Two genuine bugs from review, both of them regressions this branch introduced. The config sentinel leaked onto spawn platforms. On a non-fork start method the config layer wrote 1, and only SFT's map site rewrites that back to None. DPO, KTO, CPO, ORPO, Reward and PRM hand args.dataset_num_proc straight to Dataset.map, where datasets >= 4.1 builds a Pool(1) whose spawned child re-executes the user's __main__ (the Windows spawn loop, #3211/#3397). origin/main carried None there for the auto path, so this was a regression. The sentinel exists only to stop a downstream auto-sizer re-inflating serial, and no auto-sizer can do that when forking is unavailable, so it is now conditioned on fork. An unsized split masked a sized sibling. _largest_split_rows returned None the moment any split had no length, so a trainer with a large sized split next to a streaming one took the shortcut and returned a bare None past the env check. The Zoo reads that None as "auto" and picked 64 workers on this host: the exact inflation this branch exists to remove, and it happened with no env var set at all. Unsized splits are now skipped rather than allowed to veto, and an explicit UNSLOTH_DATASET_NUM_PROC wins on the shortcut too. Codex suggested resolving the override before the shortcut; that would return 1 for a small split and build a Pool(1) on datasets >= 4.1, so the fix is split in two instead. Also corrects the datasets boundary throughout: it is 4.1.0, not 4.0. huggingface/datasets#7702 flipped `num_proc > 1` to `>= 1`; 4.0.0 still ran num_proc=1 in-process. Verified against the 3.6.0, 4.0.0, 4.1.0 and 4.3.0 tags. One studio test asserted on 4.0.0 and would have failed on exactly that release. New file headers switched from LGPL to Apache 2.0, byte-identical to unsloth/dataprep/raw_text.py and tests/utils/data_utils.py, matching the repo LICENSE. 18 new tests. 87 pass; reverting the helper alone fails 16. * Keep macOS in-process by policy, not by a wrong start-method probe The probe read multiprocess.get_all_start_methods()[0]. multiprocess copies that function from the stdlib verbatim, darwin branch included, but not the darwin default that goes with it: its _default_context is still fork, carrying a literal '#FIXME: spawn'. So on macOS the probe said spawn while Dataset.map actually forks, and the dead-worker diagnostics printed the wrong method. Read the default context's own name instead, which fixes the report, and add _workers_unusable_reason() so the macOS refusal survives the corrected probe. Forking on macOS is what CPython itself declared unsafe when it moved the default to spawn in 3.8 (bpo-33725), and this parent has already loaded Torch and a threaded BLAS, so macOS stays in-process -- now as a stated policy rather than as a side effect of a misreport. Also run both num_proc suites in CI. They were never on the consolidated workflow's tests/utils allowlist, so all 87 guards were dead weight. * Bound the worker count Studio computes for itself A simulation across the platform x start-method x cpu x memory x request x env product found the one path that still reached Dataset.map unbounded. Studio's numbers are backend heuristics: trainer.py asks for cpu_count // 4 and safe_num_proc's own auto path is cpu_count // 3. By the time this module sees them they are explicit ints, which it reads as deliberate user intent and clamps by free memory only -- so a large host with RAM to spare kept every one of them. Measured end to end: 64 cores gave 16 workers, 96 gave 24, 192 gave 48 at ~1GB each, against a cap of 8 that the auto path has obeyed all along. The benchmark in dataset_num_proc.py has 32 workers at 14.2s versus 6.3s in-process, so those counts were slower as well as heavier, and Studio on a big machine is the configuration issue #2693 was reported from. Cap in safe_num_proc, which every Studio map() site routes through, and before the multi-GPU cap so the tighter of the two still wins. The constant is duplicated rather than imported, because importing it would pull unsloth's whole __init__ into hardware detection; a canary asserts the two stay equal, the same arrangement the Zoo's row threshold already uses in the other direction. UNSLOTH_DATASET_NUM_PROC is unaffected: it is read downstream and bypasses this. The simulation is scripts/matrix_numproc_policy.py. After the fix all 17280 cells hold every invariant: no workers on a start method that cannot support them, never 1 at a map() call site, never None at the config layer while forking works, never more workers than memory covers, never over the cap on the auto path, the env var obeyed verbatim, and deterministic throughout. * Say what UNSLOTH_DATASET_NUM_PROC=0 actually does The dead-worker message told the reader to tokenize in-process with UNSLOTH_DATASET_NUM_PROC=0. That is true almost everywhere and false in the one case the message is most likely to be read: train_on_responses_only on fork, with a split at or over the Zoo's 5000-row threshold, resolves to 1 rather than None, and datasets >= 4.1 turns 1 into a Pool(1). So the recovery advice offered for a large-dataset worker death did not remove the workers. The value is still right. A bare None there is read by the Zoo as 'size it for me' and would inflate to its uncapped count, and unsloth_zoo's _effective_num_proc returns num_proc unchanged when it is None or 1, so no value expresses in-process on fork for a large split without changing the Zoo. What was wrong was the sentence, so the sentence is now specific: fewest workers this path can use, in-process everywhere except that case, one worker there. Two tests. One reads the rendered message and requires it to name the exception, the path and the row threshold. The other drives resolve_responses_only_num_proc on both sides of the threshold and asserts 1 and None, so the message cannot claim a behaviour the resolver does not have. * Make the studio num_proc tests runnable off Linux and without torch The cross-platform staging legs failed all three, and the file's own docstring claimed it ran on any host, so both halves of that were wrong. dataset_map_num_proc returns None outright on win32 and darwin, so every assertion expecting a worker count was really an assertion about Linux and failed on the macOS and Windows runners. An autouse fixture pins the platform; the parametrised spawn-platform test sets its own value afterwards and still wins. _patch_runtime imported torch directly, which is a hard failure on a runner that has none. Worse than the error: dataset_map_num_proc treats an ImportError as "runtime not touched yet", so a torch-less host turns the XPU guard into a no-op and the test asserting None would have been passing for the wrong reason wherever it did not outright fail. It now falls back to a stub module in sys.modules, which a real "import torch" finds. Verified by reproducing both runner conditions locally rather than waiting on CI: 9 passed with torch and 9 passed with it removed. That harness needed correcting too -- it first blocked __import__ unconditionally, which is stricter than any real runner, since real Python consults sys.modules first and that is exactly what the stub relies on. * Do not trust a start method the host does not offer The cross-platform legs found a real bug in the probe, not just in its tests. On a Windows runner the private default-context chain answered "fork" while get_all_start_methods() was ["spawn"]. Those attributes are private and not consistent across builds, and a start method the platform does not offer cannot be the one in use. Believing it read Windows as forkable, so _workers_unusable_reason() returned None and workers were allowed through - the spawn re-import loop of #3211 / #3397 that this module exists to prevent. The probe now cross-checks its answer against the available methods and falls back to the documented list, with a regression test that reproduces the exact shape: spawn-only host, private chain saying fork, result None at both layers. The test failures around it were mine too, and they share a cause: the macOS policy added earlier made sys.platform load-bearing in get_dataset_num_proc, so a batch of tests that assert a worker count became platform-dependent. They passed on the Linux runner and failed on macOS. The module fixture pins the platform; the tests that are about the platform set their own value afterwards. The two studio tests that build real worker processes now skip when the host cannot fork. Under spawn inside pytest the pool fails for reasons that have nothing to do with the claim being made (WinError 10038 closing a handle, os.WNOHANG missing), and the version split they check is also asserted without processes in tests/utils. * Import multiprocess before the tests spoof the platform The Windows leg of staging CI failed inside a real worker pool with AttributeError: module 'os' has no attribute 'WNOHANG'. multiprocess picks its concrete contexts at import time from sys.platform, and both test files spoof that to linux, so the first import under the spoof handed a Windows runner the POSIX fork contexts. get_all_start_methods() then reported fork, the skip guard did not fire, and the pool tried to reap a child the way only POSIX can. Import multiprocess at module scope, before any fixture runs, and read the real platform there too so the two real-pool tests skip on Windows even if something later lies about it. * Tighten the comments added by this PR * Import the num_proc policy from the zoo, not back into unsloth The trainer source rl.py generates ran `from unsloth.utils.dataset_num_proc import ...`. unsloth/__init__.py is what generates that source, so the import reaches back into the package mid-flight, and it also drags unsloth/utils/__init__.py -> packing -> attention_dispatch -> models._utils, which means a module whose only imports are contextlib, os, sys and typing arrives through torch and the whole model stack. Nothing circular in practice: the injected imports are function-body imports that run at config construction and dataset prep, and tripping them mid-import at rl, attention_dispatch, packing, llama and chat_templates all resolved. But the coupling is real. Cold, in a process that imports only the compiled trainer cache, it costs a 9.7s `import unsloth`, and it inherits any unrelated failure in that import: on a box with a torchao/torch mismatch the stdlib-only helper failed to import along with everything else. The policy now lives in unsloth_zoo.dataset_num_proc (unslothai/unsloth-zoo#984), which unsloth already depends on and which never imports unsloth. Every call site tries the zoo first and falls back to the copy here, so upgrading unsloth alone still fixes the bug on an older zoo, and a new zoo takes over with no further change. test_the_two_copies_have_not_drifted compares the two, with docstrings stripped, whenever both are importable, so they cannot silently disagree about a worker count. Verified both directions end to end: with the zoo module present the generated config imports unsloth_zoo.dataset_num_proc and never touches the unsloth copy, and with it absent the fallback runs and the config still comes out bounded. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten the comments this PR adds * Mirror the zoo's container and start-method-split fixes unslothai/unsloth-zoo#984 review found three holes in this policy, and the copy here is a code twin of that module, so it takes the same changes. psutil reports the HOST inside a cgroup, so a 2GB container on a large box read as having room for the full worker set, and a one-core pinned job auto-sized workers that contended for that core. Memory is now the smaller of the host reading and the cgroup limit less its current usage, and the CPU count the smallest of the host, the affinity mask and any cgroup quota. resolve_responses_only_num_proc handed the zoo a bare None to mean serial, but the zoo's own veto reads stdlib multiprocessing. Where multiprocess is on spawn while stdlib is on fork, that None is read as "size it for me". It re-encodes as 1 when the two disagree. The CPU-count test patches move to the resolved count: patching psutil alone would let a 4-vCPU runner override a test that asks for 128. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Mirror the cgroup usage-path correction from the zoo /sys/fs/cgroup/memory.current is the whole machine's usage at the root, so subtracting it from a systemd unit's own MemoryMax left every run with nothing free. Usage now comes from the directories the limit was resolved from. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Re-run CI Every check on the previous head was cancelled at 16:33 by an Actions-level event that hit both repos and many unrelated branches, main included. * Bound the paths the review found still unbounded Five findings, all about a value reaching Dataset.map without the policy. The config sentinel was written as 1 for every patched trainer. Only SFT has a downstream auto-sizer to defend it against: DPO, KTO, CPO, ORPO, Reward and PRM hand args.dataset_num_proc straight to Dataset.map, where nothing can inflate a None but a 1 is a Pool(1) on datasets >= 4.1 -- one worker holding its own tokenizer copy, on the low-memory host that had just refused workers. The codegen now picks the encoding per trainer. train_on_responses_only with UNSLOTH_DATASET_NUM_PROC=0 and an explicit count returned 1, which bypasses the small-split guard and builds that Pool(1) even on a 100-row split. Under the threshold the guard is in-process, so None is what expresses the request exactly, and that is what it now returns. Studio's dataset_map_num_proc handed its own count straight to callers in format_conversion.py and chat_templates.py with no memory ceiling and no environment override, though the cap's log line advertised one. It now runs the count through the shared policy when unsloth_zoo has it, so those paths get the memory and cgroup clamp and the escape hatch, and the log line no longer names a variable that path never read. The fallback copy moves to unsloth/dataset_num_proc.py. Under unsloth/utils it sat behind an __init__ that imports .packing (torch) and .attention_dispatch (unsloth.models._utils), so a torch-free MLX host with an older zoo raised before train_on_responses_only could delegate. Also found while running the wider suite: the tokenizing map() anchor was required, so a Zoo release moving that line would hard-fail every SFT run over a diagnostic wrapper. It is optional now, like the selection anchor above it, and the drift canary is what reports it. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Mirror the zoo cgroup fix into the fallback copy unsloth_zoo.dataset_num_proc is the source of truth; this copy exists only so upgrading unsloth alone still fixes the bug, and test_the_two_copies_have_not_drifted holds the two together. Also moves the three new files onto the AGPL-3.0 header the repo now uses for new sources. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep serial requests serial across the config boundary and the drifting anchor Three review points. Studio computes dataset_num_proc for a config, not for a map() call, and the shared policy was applied with the map-site sentinel: the audio and CUDA-audio paths ask for 1, that became None, and SFTConfig read None as "auto-size me" and returned 8. Measured on the generated class, not simulated. The XPU leg was worse: unlike win32 and darwin, forking still works there, so a config None was auto-sized back up and forked the Level-Zero context the guard protects. dataset_map_num_proc now takes serial_as_none, default True so the seven map-site callers are untouched, and the trainer passes False. The spawn platforms keep None at both layers, where nothing can inflate it and a 1 would reach Dataset.map from DPO and friends as a Pool(1). The sft_prepare_dataset num_proc anchor was optional on the grounds that its absence was harmless. It was not: the config layer encodes serial as 1 for that rewrite to turn back into None, and an un-rewritten zoo hands the 1 to Dataset.map, which pools for any count from datasets 4.1. It now falls back to the assignment the block ends with, unchanged in the zoo since Aug 2025 while the block around it was rewritten three times in 2026, and warns only when both anchors miss. Hard-failing instead would break every install on a newer zoo. Two cgroup tests read the host tree once the fallback reader stopped needing unsloth_zoo, so they passed on a laptop and failed in a limited container. They are isolated now, and the six unaided-reader tests from the zoo copy came with them. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Skip the map-site half of the new tests when the policy is absent Two of the new config-boundary tests asserted dataset_map_num_proc(1) is None without guarding on the policy import, so they failed on all three cross-platform runners: unsloth_zoo there has no dataset_num_proc module yet (it ships in unslothai/unsloth-zoo#984), and _bounded_by_the_shared_policy returns the count unchanged in that case by design. Same pytest.importorskip guard the three memory and env tests beside them already use. With an unsloth_zoo that lacks the module the file is 15 passed, 6 skipped instead of 2 failed. * Let the shared policy see the request Studio was actually given Three points on the hardware.py side. safe_num_proc materialized an auto request before the policy could see it, so the policy never ran its own auto path: it reads this process's CPU affinity and cgroup quota, while safe_num_proc reads the host os.cpu_count(). A 2-core container on a 64-core box asked for cpu_count // 3 workers and was bounded only by memory. The request now passes through as written, and Studio's caps are applied to whatever the policy chose, since the multi-GPU fork-deadlock cap is knowledge the policy does not have. For an explicit count the two orders are equivalent, both being min(studio cap, request, affordable). The escape hatch is unvetoed by contract, but the win32/darwin return fired before the policy could read it, so UNSLOTH_DATASET_NUM_PROC was silently ignored on the platforms whose dead-worker message recommends it. It is now checked before that veto, and Studio's caps never apply to it. The older-zoo path returned the Studio count unchanged rather than trying unsloth.dataset_num_proc, the byte-identical fallback every other call site uses. It is used now, but only when unsloth is already imported: importing it from here would make hardware detection patch torch and pull in the model stack. The torch-less XPU branch routes through the policy too, having been the one path that ignored both the ceiling and the hatch. Seven new tests, 28 total, 17 passed and 11 skipped against an unsloth_zoo without the module. Reverting each of the three fails its own test. * Mirror the zoo test isolation and prose The fallback copy tracks unslothai/unsloth-zoo#984: the dnp fixture pins the memory ceiling at its sources, so a memory-limited runner cannot turn a start-method test into a clamp test, and the dead-worker advice now says that the single-worker exception applies to a Zoo older than the one that reads 1 as in-process. * Honour the hatch on XPU, leave the ordinary case to the policy, ignore typos Three follow-ups to the previous round, all of the same shape as fixes already made one line away. The XPU-initialized return bypassed the policy the way the spawn platforms did before this, so UNSLOTH_DATASET_NUM_PROC was ignored there too. It takes the same route now: the guard exists because fork corrupts the Level-Zero context, but a user who set the variable has accepted that, and unset the veto stands at both layers. The trainer's non-audio branch passed max(1, os.cpu_count() // 4), which the policy reads as an explicit request and so skips its own auto path, the only one that consults this process's affinity mask and cgroup quota. It passes None now, and Studio's caps still apply to whatever the policy chooses. The override probe treated any non-empty value as active, but the policy warns about and ignores an unparseable or negative one, so a typo skipped the multi-GPU cap while contributing nothing. It reads the parsed result through the zoo's new environment_override(), falling back to presence on a copy that predates it. Four new tests, 32 total. Reverting the XPU check or the override probe fails three of them; the trainer's None is pinned by an AST guard, since dropping it changes only the worker count on a container. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Name the encoding on the cgroup reads tests/test_runtime_text_encoding.py caught both call sites the unaided cgroup reader added: a locale-dependent text read crashes or produces mojibake on a Windows console codepage, and the gate is absolute even for ASCII kernel files. Mirrors unslothai/unsloth-zoo#984. * Neutralise the zoo cgroup readers by name in the num_proc fixture Pinning hf_xet_tuning.CGROUP_ROOT only works against a zoo that has that global. An older one still exposes the private dir helpers the policy prefers, so monkeypatch finds nothing to pin and the readers walk the runner's real cgroup: under a 2GB memory.max that turns a test about the start method into a test of the clamp. * Patch the zoo cgroup readers through the cache the policy actually uses unsloth_zoo/__init__ imports hf_xet_tuning near the top and only raises "Please install Unsloth" at the end, so a failed package import drops unsloth_zoo from sys.modules and leaves unsloth_zoo.hf_xet_tuning behind. The policy reaches the submodule through that surviving cache entry, so treating the failure as absence left the real readers live on the runner's own /sys/fs/cgroup and every sizing assertion silently became a test of the container's memory limit. * Make the Studio half of this PR actually run, and three tests mean what they say studio-backend-ci is the only job that executes studio/backend/tests, and it installs studio.txt, which carries no unsloth_zoo: 14 of the 32 cases importorskip away there, and with no policy installed the survivors fall back to the pre-PR safe_num_proc, so they would pass with the whole wiring deleted. Run the file in the hard-gate step instead, which has an editable unsloth_zoo, and pin the memory ceiling so the counts are not really assertions about the runner's free RAM. test_env_override_is_uncapped never exercised the exemption it is named for: the fixture leaves room for 512 workers, so asking for 100 was never near the clamp. test_unrelated_errors_pass_through_untouched held under 'except Exception' too, since the guard re-raises the same object; it now also passes a non-RuntimeError carrying the dead-worker text. And the codegen tests supplied their own copy of rl.py's serial_as_none rule, which made them self-fulfilling -- they now read it out of rl.py's AST, so flipping SFT to True fails the behavioural test and not only the literal match. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Reach the policy without importing unsloth, so the gate is not inert The CI step I added runs the Studio file in a job with an editable unsloth_zoo, but the policy module is not on unsloth_zoo main -- it is in the companion PR -- and the job clones main, so all 14 cases that reach the policy still skipped and the survivors still exercised the pre-PR path. The claim in that step's comment was wrong. The same gap is live in production: _shared_policy only fell back to the in-repo copy when unsloth was already imported, and no Studio backend module imports it, so the API process reaching format conversion got no policy at all on every install whose zoo predates the module -- the 2GB container with eight cores this PR exists to fix. It now loads the file off disk when the package is not imported, which is safe because the module is stdlib-only by design, and memoises through sys.modules so the warn-once state and the cgroup reads are not redone per map() call. The tests ask _shared_policy for the same object, so they patch what production uses: 32 pass with the zoo copy blocked, where 18 passed and 14 skipped before. Also parenthesise the source segment _rl_serial_as_none evals, matching its sibling: a formatter reflowing that ternary in rl.py turned all eight codegen tests into an IndentationError. And mirror the two cgroup and escape-hatch tests just added on the zoo side. * Count pools at the class, not at a module attribute datasets moved The hard gate I added surfaced this the first time the file ran against HF=latest: datasets 3.x and 4.x do 'from multiprocess import Pool', so datasets.arrow_dataset.Pool exists, but 5.x calls mp.Pool() and a spawn context instead and the attribute is simply gone, so the spy raised AttributeError on four Python versions. Patching multiprocess.pool.Pool's __init__ catches every route. Verified against a real datasets 5.0.1: num_proc=None builds no pool, num_proc=1 builds one, so the claim the test makes about 4.1+ still holds there. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>	2026-08-05 05:31:27 -07:00
Long Yixing	934f879043	feat(mlx): route trainer callbacks (#6929)	2026-07-08 03:25:50 -07:00
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