2 commits

Author	SHA1	Message	Date
Daniel Han	1fe8c7a136	Stop padding-free SFT from tripping the TRL >= 1.0.0 max_length guard (#7951) ... * Stop padding-free SFT from tripping the TRL >= 1.0.0 max_length guard TRL 1.0.0 added a guard that refuses to build an SFTTrainer when padding-free is on, packing is off and args.max_length is set. Unsloth auto-enables padding-free whenever padding_free is left at its None default, and the max_length_check codegen always wrote args.max_length from max_seq_length, so every default SFTTrainer(...) raised on TRL >= 1.0.0. rl.py now clears args.max_length for a TRL that carries the guard, and routes the truncation length through max_seq_length instead, which is what Unsloth's own dataset prep reads. An explicit user max_length is honoured there rather than dropped, with a one-line notice when it differs from the resolved length. The block is only emitted when the guard text is present in the TRL source, so older TRLs receive exactly what they did before. rl_replacements.py makes the Zoo's truncation-length seed None-safe, and trainer.py grows a back-off for a TRL that words the same guard differently. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep the resolved length and TRL's guard when padding-free swaps max_length Review follow-ups on the padding-free / max_length swap: Only route the cap into max_seq_length when Unsloth's dataset prep will really run its truncating tokenize pass. It does not for dataset_kwargs= {'skip_prepare_dataset': True}, nor for a dataset that already carries input_ids / labels, which makes the Zoo's sft_prepare_dataset skip tokenization. Clearing args.max_length there told TRL the caller had supplied truncated rows while TRL hands its collator max_length=None under padding-free, so 402-token rows reached training against a 128 cap. Turn padding-free off and keep max_length for those datasets instead, so TRL's own collator truncates. Drop the _unsloth_requested_max_length override. The block above already resolves the length under Unsloth's documented precedence (max_seq_length, capped by the model, beats max_length), so re-reading the raw user max_length inverted it: max_seq_length=4096, max_length=512 truncated at 512 on TRL >= 1.0.0 and at 4096 everywhere else. The resolved length is now used as-is. Tests cover both, and fail on the previous commit. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix padding-free test on CPU-only runners for PR #7951 UNSLOTH_ALLOW_CPU=1 makes `import unsloth` skip both halves of the SFT patch on purpose, so cross-platform CI saw stock trl.SFTTrainer and every construction test failed with "SFT patch did not apply". Request the codegen swap and the __init__ wrapper explicitly, in _gpu_init's order. * Run the padding-free suite in the version-compat jobs that install the deps * Read the yielded row schema, not column_names, before clearing max_length * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Copy the length cap unconditionally before clearing max_length * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trim the comments on the padding-free max_length handshake * Truncate pre-tokenized rows instead of leaving the cap unenforced Disabling padding-free stopped TRL's guard from firing but did not make max_length mean anything. TRL enforces the cap in _prepare_dataset via truncate_dataset, not in the collator: the LM collator it builds for the language-modeling case is constructed without a max_length. The Zoo's _prepare_dataset returns already-tokenized rows untouched, so a 402-token row still reached the model under a 128-token request. Truncate those rows here instead, by the same rule TRL uses (slice every per-row list column, so input_ids, labels, attention_mask and the masks stay aligned). That restores TRL's own contract rather than inventing one, and padding-free can stay on instead of being dropped, so the speed win survives. Written out rather than imported from trl.data_utils: that module pulls in the processor stack, and an ImportError there would silently drop the cap. A failure to truncate now prints rather than passing quietly, since a swallowed error reads exactly like the cap being enforced. Two cases keep the old behaviour on purpose. skip_prepare_dataset means the user asked for the dataset to be left alone and TRL skips truncation there too. A with_transform dataset reports column_names as its backing columns while yielding input_ids, so mapping it would truncate the wrong thing. Tests split accordingly: pre-tokenized rows are truncated, the cap is consumed and padding-free stays on; skip_prepare_dataset and the transformed dataset keep the cap and lose padding-free. The truncated case asserts a collated width of 2 x cap, because padding-free concatenates the batch, and asserting the bare cap there would be asserting padding-free was off. Verified against trl 1.9.2, which has the guard, as well as 0.25.1, which does not. The suite is vacuous without the guard: every assertion sits behind trl_has_guard, so a green run on a guardless TRL proves nothing. * Truncate every eval split, and columns that are not lists Two gaps in the truncation branch, both found in review. TRL accepts eval_dataset as a dict of named splits, and a dict has no .map of its own, so the splits were skipped while max_length was still cleared below. Evaluation then ran at the full length the cap was meant to stop. Each split is mapped now. A dataset with a torch or numpy format hands batched map() tensors rather than lists. raises on a tensor's ambiguous truth value, and the isinstance(list) check would have left the column alone in any case, so the rows stayed long while the cap was cleared. The predicate now asks for a per-row sequence via __len__, excluding str and bytes so a text column is never sliced. Both new tests fail against the previous version: the formatted dataset keeps overlength rows, and the named eval splits are left untouched. * Verify the truncation instead of assuming it worked Three more findings, and together they say the shape was wrong: every round found a new way an unconditional dataset rewrite corrupts something. So the branch now decides per DATASET and then checks its own work. A raw eval split was being sliced as if it held tokens. messages is a per-row sequence too, so a blanket map cut conversation turns off the end and silently corrupted evaluation. The truncatable test now runs against each split on its own, so a raw split stays raw for the tokenizer pass that follows. A with_transform dataset whose BACKING table carries input_ids passed the schema test, but map() writes that table while the transform keeps rebuilding overlength rows on read, and the cap was cleared anyway. Custom-format datasets are now rejected outright. Dataset.map ran single-process regardless of dataset_num_proc, which TRL's own truncate_dataset honours through its map_kwargs. It is forwarded now. The real change is the last one: after rewriting, a row is read back and compared against the cap, and anything still over it restores the original datasets, keeps max_length and drops padding-free. Every predicate above is a guess about what a dataset will do; this is the only step that observes it, so an unforeseen shape degrades to the safe path rather than training uncapped. Confirmed independent: the with_transform case only fails when BOTH the format guard and this read-back are removed. 19 passed on trl 1.9.2 (guard present), 13 passed 6 skipped on 0.25.1. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Refuse a cap nothing can enforce, and cover every split Three Codex items on the padding-free handshake. Dropping padding-free keeps max_length for TRL's collator, and no TRL from 0.22.2 to main truncates there: truncation lives only in _prepare_dataset, which returns pre-tokenized rows untouched. So for a with_transform dataset that already yields input_ids there was no enforcement path left, and construction now succeeded where TRL's own guard used to refuse. Verify a row and raise when one is over. The cap was consumed on the train split's word alone. A tokenized eval split in a shape the truncation cannot rewrite is left as it was, and prep does not re-tokenize rows that carry input_ids, so evaluation ran over the cap. Every split is checked now. The truncation map forwarded args.dataset_num_proc raw. The config layer writes serial as 1 and datasets >= 4.1 builds a Pool(1) for it, so an explicit dataset_num_proc=1 or UNSLOTH_DATASET_NUM_PROC=0 could still fork a tokenizer worker. Resolved through the same helper as every other map site. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Check every row for the cap, not just the first A short row 0 in front of a long row 5000 read as within the cap, and in the fallback branch nothing downstream truncates it. A map-style split is now read in full; a stream cannot be rewound, so a bounded prefix is all there is and the check says so rather than pretending otherwise. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cap every split, and truncate a stream instead of sampling it The truncation block was gated on the train split, so a raw train set beside a pre-tokenized eval set skipped it entirely, consumed `max_length`, and left evaluation uncapped: preparation does not re-tokenize rows that already carry `input_ids`. Each split is now capped and checked on its own. `IterableDataset.map` takes no `num_proc`, so forwarding the auto-sized one raised TypeError, the catch restored the stream, and the run died on "cannot be enforced". The map kwargs are chosen per split now, and the lazy map caps every row a stream will ever yield, which the 1024-row prefix scan could not promise. A stream that cannot be rewritten stays unverified and is refused. Enforcement is no longer confused with observation. A `with_transform` split that happens to sit under the cap rebuilds its rows on every read, so it holds `max_length` and turns padding-free off, the same answer the train side already gave. Splits that were rewritten keep their truncation when another split cannot be enforced; rolling them back put an overlength train set back and turned a healthy run into the same error. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Match TRL's own truncation: keep_end, packed splits, and masked rows Three ways the manual truncation diverged from what TRL's `_prepare_dataset` does, each of them silent because this path consumes `max_length` and stops TRL from doing it. `truncation_mode = 'keep_end'` slices `[-max_length:]` upstream, and callers use it when the completion sits at the tail of a long prompt. Always keeping the prefix trained on the wrong half of every row. A packed split carries `seq_lengths`, which holds document lengths rather than tokens. Slicing it by the cap left it describing the pre-truncation row, so padding-free built position ids for more tokens than `input_ids` still held. TRL skips truncation entirely when packing, so such a split is refused rather than cut, and the remaining columns are matched by row length against `input_ids` so a sidecar list is never mistaken for a token sequence. TRL drops rows left fully masked immediately after truncating, since a prompt that alone fills the cap has every label at -100 and contributes no loss. Without that filter, completion-only datasets fed batches with no supervised tokens. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Three fixes to the pre-tokenized max_length handshake A 0-dimensional tensor has __len__ and raises on it, so under set_format("torch") a scalar id column read as a per-token sequence and the later len() threw TypeError. The outer catch then restored the overlength dataset and a truncatable run died on "cannot be enforced". Probe with len() instead of hasattr. Supervision is carried three ways and only labels was filtered after truncation. A completion-only or assistant-only row whose prompt fills the cap is left with an all-zero completion_mask/assistant_masks, which TRL's collator turns into all -100: no supervised token in the row, and a batch made of them is a NaN loss. skip_prepare_dataset exempted the overlength check, which made it the one route to a silently uncapped run: TRL then neither truncates nor builds its collator with a truncation length, so the oversized rows reach the model with max_length set and ignored. The behavioural tests only run on a TRL that ships the guard, so the three changes are pinned in the emitted source as well. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Gate the mask filters on their loss mode, and fix the test I contradicted A mask is only supervision when its loss mode is on. With completion_only_loss = False TRL ignores completion_mask and trains from the normal labels, so a row whose retained mask is all zero is still a good row; filtering it biased the split or emptied it. completion_only_loss defaults to None, which TRL reads as on for a prompt-completion dataset, so only an explicit False opts out. assistant_only_loss defaults to False and has to be asked for. labels stays unconditional: it IS the supervision. test_unprepared_datasets_keep_their_length_cap passed an OVERLENGTH split with skip_prepare_dataset, which is now the one configuration that raises, so on a TRL carrying the guard it asserted the opposite of test_skip_prepare_dataset_does_not_excuse_an_overlength_row. Rewritten against rows that already fit, which is what the flag's promise is about. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Filter assistant_masks on presence, the way TRL's collator reads it * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cap a pre-tokenized eval split handed to evaluate() after construction * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Resolve the loss mode like TRL, intersect the masks, cap predict too * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make the late eval cap agree with the one at construction, and leave a packed eval split to the packer The evaluate/predict wrapper is the same cap arriving late, so it has to match the construction-time cap on every detail. Four ways it did not. A stream came back uncapped, and silently: dataset[0] does not raise on one, because datasets 4.x reads the 0 as a COLUMN name and returns an IterableColumn whose len() then threw TypeError into the catch that returns the original. So neither this wrapper nor TRL enforced anything. The prefix scan is now skipped for a stream and map() is applied directly, which is lazy and covers every row it will ever yield. truncation_mode was ignored: [:cap] always, while TRL and the constructor slice [-cap:] for keep_end. Callers who set it are the ones whose completion sits at the tail, so eval ran on the wrong half of every long row. Rows left with no supervised token were kept. TRL filters those right after its own truncation, but args.max_length is None on this path so TRL does not run, and a batch of them reports NaN. Same labels-then-mask-intersection rule as the constructor, with completion_only_loss resolved from the dataset shape. A packed split was partially sliced. seq_lengths describes documents, not tokens, so cutting input_ids under a stale one makes padding-free build position ids for tokens the row no longer has. The constructor refuses this shape; so does this now. Separately, eval_packing is resolved apart from packing (packing = args.packing if args.eval_packing is None else args.eval_packing), so packing = False with eval_packing = True reaches a branch gated on not packing. TRL then packs that eval split instead of truncating it, and wrapped packing concatenates the whole token stream before chunking, so truncating each row first evaluates on a truncated corpus. That case now drops the enforcement claim rather than the split: max_length stays and padding-free turns off, which is also what packing itself requires. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Let the eval packer own its split without raising, and make supervision one intersection Sparing an eval split for TRL's packer set _unsloth_capped = False, which drops into the branch that scans every split and raises on an overlength row. That split is overlength on purpose, so the previous commit turned a working eval-packing run into a ValueError and denied wrapped and bfd_split the overflow they exist to handle. The scan now skips the eval splits the packer owns, and keeps scanning the train split, which nothing packs. _unsloth_eval_packing moved one level out so the fallback can read it even when skip_prepare_dataset skips the truncation block, and it no longer requires an eval split at construction: a split handed to evaluate() later would otherwise find max_length already gone. The late evaluate/predict cap makes the same call now, and returns a split the eval packer will take untouched. completion_only_loss is resolved once, from the training sample, because that is what TRL does (dataset_sample = next(iter(train_dataset))). Resolving it per split disagreed with the collator whenever the schemas differ: prompt/completion training data makes the collator apply completion_mask, while a pre-tokenized eval split carrying only input_ids and completion_mask read as full-sequence loss, so rows whose mask truncated to all zeros survived and went all -100. The resolved value is parked on args so the late cap uses it too. labels joined the mask intersection instead of being a filter of its own. The collator applies the masks onto the labels, so a row supervised at one position and masked in at another passed both filters separately and still went out with no supervised token. truncation_mode now has to be keep_start or keep_end. TRL's SFT path never reads that attribute, so nothing downstream would catch a third value and mapping it to the default cuts from the side the caller asked us not to; the enforcement claim is dropped instead. A late split with no column_names is read from a row rather than assumed raw. A custom map-style dataset carries none, and args.max_length is already cleared on that path, so nothing else would have capped it. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cap a late eval split that cannot be rewritten, and give the mask-filter test the mode it asserts The late cap called .map(), which belongs to datasets. A torch.utils.data.Dataset or a plain list has neither map nor filter, so the call raised AttributeError into the broad catch, which returns the original, and the split reached the collator uncapped on a path where max_length is already None. A with_transform dataset is the same problem from the other side: it HAS map, but rebuilds its rows on every read, so mapping writes a backing table nobody reads, and its column_names still reports the backing schema while it yields input_ids, so it did not even read as pre-tokenized. Both are now capped on read through a small wrapper: rows are sliced on __getitem__ and __iter__, which is how the collator reaches them either way, and the surviving indices are resolved once up front since dropping unsupervised rows changes the length. Columns are read from metadata AND from a yielded row, because either alone misses one of these two shapes. Separately, test_rows_whose_mask_is_truncated_away_are_dropped now passes completion_only_loss = True. Its split has no prompt/completion columns, so a None resolves to False from the train sample, the collator ignores completion_mask entirely, and filtering on it would be deleting rows that still carry full-sequence supervision. The test is about the filter, so it has to ask for the mode that makes the mask supervision. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Close four gaps in the late eval cap: streams, predict, the stored split, short rows * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cap the late splits nothing else will, and keep the packer bypass out of it Reconciles the four late-cap gaps into one implementation. Nothing TRL owns runs on a split handed over after construction: _prepare_dataset is called from __init__ and nowhere else, and SFTTrainer overrides neither evaluate nor predict nor get_eval_dataloader. So the eval_packing bypass was wrong on both entry points, not just predict: no packer will ever see a late split, and skipping the cap there hands the collator raw overlength rows with max_length already cleared. The test asserts that premise off the installed TRL source rather than assuming it. A stream capped on read has to still BE an IterableDataset, so the wrapper subclasses one, and the shared base no longer carries the raising __len__ that made a stream look map-style. Its __getattr__ refuses dunders and its own state as well, since a DataLoader worker pickles the split and a __setstate__ arriving before __init__ would recurse on _inner forever. Being under the cap is not the same as being supervised, so the length scan now only decides whether anything needs truncating; the supervision filter runs either way, as it does at construction time. It hands back the caller's own object when it drops nothing. evaluate() with no argument falls back to self.eval_dataset. That split is swapped onto the trainer for the call rather than passed down as an argument, because HF recurses over a dict of splits by NAME when nothing was passed, and passing the dict turns that into an override. Capping is memoized per split object: every eval during training comes through here, and the scan materializes the whole input_ids column. * Make the late-cap wrappers picklable, probe non-destructively, keep predict whole * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep the mode refusal, refuse a one-shot stream, require the seed rewrite Three from the latest review round. `_unsloth_capped` is seeded from the truncation mode and the eval branches combine into it with `and`, but the train branch assigned straight over it. So a `truncation_mode` that is neither keep_start nor keep_end printed "not being enforced here" and was then served as keep_start anyway, with `max_length` cleared and padding-free still on. Now the same `and` shape the eval splits use. The fallback cap scan iterated the split to check it, which for a one-shot stream IS consuming it: the trainer would get an exhausted dataset, or one short by the whole 1024-row prefix. Two `iter()` calls returning the same object marks one -- true for a generator and for a torch IterableDataset with shared iterator state, false for a datasets.IterableDataset, which restarts -- and neither call reads a row. The answer on a hit is the same as for an unfinished prefix: not proven within the cap, so the enforcement claim is dropped and every row stays. The max_length seed rewrite is not the optional worker-count edit `_replace_or_fallback` was written for, whose warning text it was borrowing. The generated trainer clears `args.max_length`, so an unrewritten seed reads that None rather than the 0 that makes the guard fall through, and a raw dataset stops being truncated. `required = True` keeps the two-anchor tolerance and raises only when both miss, which is where the neighbouring _require_replace edits in the same function would raise a line later anyway. Both anchors match the installed unsloth_zoo, so this changes nothing today. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use one one-shot signal everywhere a split is probed The cap scan learned that two `iter()` calls returning the same object marks a single-pass stream; the three other probe sites kept `iterator is dataset`, which is true for a bare generator and false for an `IterableDataset` whose `__iter__` hands back one stored generator. Just as single-pass, and each site read a row off it and moved on. `_column_names` now uses both tests and chains the probed row back for either. A `datasets.IterableDataset` restarts, answers False and is rewound rather than chained, which is what stops row 0 being duplicated. The two probes in the generated block read the first TRAINING example and had nothing to chain it back to, so the run began at row 2. Writing the test turned up a second one beyond the reported site: the prep-truncation probe just above. Both now ask metadata first and fall back to a row only when the stream can spare one; the completion-only probe then resolves to False, which is what TRL answers for a split with neither `prompt` nor `completion`. Not chained back at these two, because rebinding `train_dataset` to an `itertools.chain` inside the constructor would hand TRL an object without the dataset API it goes on to use. `_memo_token` refuses a split whose format type is custom. `_fingerprint` covers the backing table, not the transform, so a transform closing over mutable state yields different rows under an unchanged token and the memo replayed a supervision filter decided against the old ones. One existing assertion moved off a fixed 600-byte window and onto the next anchor: a comment added inside the block had pushed it out of range. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cap the split a string eval key names, and stop trusting a transform's schema Three things the same optimistic reading of dataset metadata was behind. evaluate(eval_dataset = 'validation') is the supported way to pick one split out of a stored dict: get_eval_dataloader resolves it as self.eval_dataset[key], so capping the key was a no-op and the split it named reached the collator over max_seq_length with max_length = None. Cap the stored split in place instead and hand the key straight back. column_names describes the BACKING table, so a with_transform split storing text while yielding input_ids read as raw and the cap was cleared for rows nothing then truncates. Discard the metadata when the split has a custom format and probe the yielded row, which is free there: such a split rebuilds its rows on read. A stream that cannot spare a probe row cannot be ruled tokenized either, and leaving _unsloth_prep_truncates at its optimistic seed cleared the cap on that guess. Refuse instead, which costs padding-free on those streams and keeps max_length for TRL's own collator. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cap the dataloader paths, judge token columns by value, and stop the last destructive probe Three from the same round, all ways the late cap could be walked past. get_eval_dataloader and get_test_dataloader are public API and neither goes through evaluate/predict, so a caller building a dataloader directly reached the padding-free collator with args.max_length already cleared and nothing capping the split. Same wrapper, same argument position. The per-token allow-list also treated every token-shaped NAME as sliceable. An optional column stored as token_type_ids = None made the map raise, and the broad catch around it handed back the uncapped split; a 2-D position_ids sliced on the wrong axis and came out misaligned with the truncated input_ids. Judged by a row now, the way the construction-time truncation already does. And _unsloth_pretokenized still read a row off a one-shot stream, which the schema probe and the cap scan had both stopped doing: read raw it declares the split safe and training starts at row 2, read tokenized it rejects a caller-owned stream it has already mutated. Schema first, a row only when one is free, and an unprobeable stream holds max_length rather than clearing it. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cap a split once, probe a transform, and leave an unknown mode alone The dataloader wrappers made `evaluate()` reach the cap twice: it caps, stores, and the original then calls `get_eval_dataloader`, which is wrapped too. Over a one-shot stream that was destructive, since `_CappedStream` hands out a fresh generator over the same exhausting source rather than rewinding, so the second pass's probes ate rows. `_cap` now stamps what it capped to and hands a matching split straight back. `_unsloth_pretokenized` returned on `column_names`, which for a transform describes the BACKING table, so a split storing `text` and yielding overlength `input_ids` was called raw and had its cap cleared. The transform rule is now one helper both schema probes read. An unknown `truncation_mode` seeded the refusal but still sliced, so the fallback scanned an already-trimmed split and only turned padding-free off: every row silently cut from the start. Both mutation sites are gated now. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Align every column, stay lazy, key the cache, and refuse an unknown mode - `_column_names` read a row and threw it away, so on a one-shot stream `_sliceable_per_token` had no widths to compare and cut `input_ids` alone, leaving `labels` and `attention_mask` overlength. It hands the row back now. - `_CappedRows` built an identity index even with no supervision to filter on, a whole extra pass over a `with_transform` split before the dataloader starts. No supervision means no index. - The eval-cap memo omitted `truncation_mode`, so keep_start then keep_end reused the cached prefixes. - The late cap took an unknown `truncation_mode` as keep_start, cutting from the side the caller ruled out. It refuses, like the construction path. - A retained `max_length` was read as proof the cap is enforced downstream. It is what the block leaves behind when it turns padding-free OFF instead, and `_prepare_dataset` never runs for a late split, so an overlength one reached the model uncapped. This reverses an earlier test whose premise -- that TRL truncates in its own prep -- does not hold for the late path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Refuse the padding-free retry when nothing enforces the cap * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Validate the late cap per row, and bound what it remembers The read-side wrapper chose its sliceable columns from one probed row and then applied that slice to every row, so an optional column that is None or a different width further in raised inside the dataloader. It now measures each row the way the map path already does. A user-defined per-token field was never on the allow-list, so it kept its full length while input_ids was cut and a custom collator saw mismatched rows. Those ride along when the row proves them aligned and flat. The cap mark was trusted forever on the caller's own object, which a set_transform can mutate under it. It now carries the fingerprint the memo already requires, and the memo itself is bounded: every entry pinned both the split and its capped copy for the trainer's lifetime. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cut nested token fields, and stop three probes reading the wrong source A [seq_len, channels] field has its token axis first, so the slice is correct there; the nested test that rejected it was aimed at the channel-major shape, which the length check already catches. Left uncut, it handed a custom collator the old sequence length beside capped tokens. The construction-time batch map classified a column from its first row and then called len() on every value, so an optional field that is None further in raised, the handler restored the overlength split, and a truncatable run died on cannot be enforced. Per value now, as the late cap already does. completion-only resolved from column_names, which under set_format(columns = [...], output_all_columns = False) still lists the backing table while the rows yield only the named columns. The format's own list is what is yielded. The fallback cap scan now resolves eval_packing: TRL's eval packer owns and chunks that overflow, and the generated path already excludes those splits, so scanning them refused a configuration that path accepts. And a source edit whose result is already present is no longer an edit that failed. old is a prefix of new for the max_length seed, so the wide anchor matched the normalized line and appended a second or 0, while a Zoo that spells it differently matched nothing and raised under required = True. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Run the pre-truncation rewrite under the rank window, and read the mode TRL resolved Every rank reaches the inserted map and filter before TRL's _prepare_dataset, and TRL runs its own preparation maps under main_process_first. Without the same window, eight ranks each start num_proc workers against one Arrow cache: 64 processes doing identical work and writing the same files at once. The body moved into _unsloth_cap_one so the map AND the filter sit inside it; a single process gets a no-op context manager. The late cap read completion-only from args, which only the generated block sets. That block does not always run -- a TRL whose guard did not match, or padding-free off from the start -- and the fallback then read the LATE split's schema, answering False for one carrying only input_ids and completion_mask. Rows whose completion was cut away entirely survived and the collator turned them into all -100, i.e. a NaN eval loss. TRL's own resolved value is what the collator uses, so that is what is read now. And the idempotence check compares quote-normalized text. The narrow regex already accepts either style, so a Zoo carrying the replacement single-quoted matched neither anchor and raised under required = True. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Honour eval_packing on a late split when TRL packs one TRL 1.7.0 gave SFTTrainer its own `evaluate`, which prepares a split passed straight to it and packs it under `packing = args.packing if args.eval_packing is None else args.eval_packing`. The late cap assumed the opposite -- that nothing TRL owns ever runs on a late split -- so on 1.7.0 and newer it cut each row at the cap before the packer could redistribute the overflow. `_trl_prepares_late_evals` reads that off the class rather than off a version number, and `packs_late` is passed per entry point because `predict` and the two dataloader builders stay the base Trainer's on every TRL. It joins the memo key, since `evaluate` and `get_eval_dataloader` see the same object in one call and only the first may skip the cut. A split the supervision filter emptied now says so. Every TRL 1.x reads `next(iter(train_dataset))` in `__init__` to resolve `completion_only_loss` and `_is_vision_dataset`, so an emptied split came back out as a bare `StopIteration` naming nothing the caller could act on. The tests that asserted the old premise are rewritten to cover both sides of 1.7.0 with stub trainers, so they pin the wrapper's logic rather than whichever TRL is installed, and the version fact is pinned separately. Three more expected a transformed or unprobeable overlength split to construct quietly; that shape keeps `max_length` and reports the rows instead, which is the behaviour the block already had, so they now assert the report and a within-cap companion shows it is about the rows and not the transform. 147 tests pass on trl 0.24.0, 1.0.0, 1.5.1, 1.6.0, 1.7.0, 1.8.0 and 1.9.2. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Do not mark a split capped when only the packer was meant to own it * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Lift the quote-normalising helper into the anchor-helper test namespace * Give the drift fixture the Zoo's real max_length seed line * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: danielhanchen <unslothai@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen <unslothshared@gmail.com> Co-authored-by: danielhanchen <danielhanchen@users.noreply.github.com>	2026-08-09 04:04:58 -07:00
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