mirror of
https://github.com/unslothai/unsloth.git
synced 2026-08-20 22:34:00 +00:00
1fe8c7a136
* 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>
1680 lines
70 KiB
Python
1680 lines
70 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
|
|
|
"""Tests for unsloth.dataset_num_proc, the fallback copy of the policy.
|
|
|
|
unsloth_zoo.dataset_num_proc owns it; this copy only runs on a zoo that predates
|
|
the module, and ``test_the_two_copies_have_not_drifted`` holds them together.
|
|
|
|
The module is stdlib-only by design and is loaded straight off disk rather than
|
|
via ``import unsloth``, so these assertions stay meaningful on a host whose
|
|
torch/unsloth_zoo pair cannot import. The ``test_rl_codegen_*`` cases tie it back
|
|
to the import path ``unsloth/models/rl.py`` generates, catching a rename.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
import importlib.util
|
|
import re
|
|
import sys
|
|
import textwrap
|
|
import types
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
try:
|
|
# Import before the dnp fixture spoofs sys.platform: multiprocess picks its
|
|
# contexts at import time, so a Windows runner would get POSIX fork contexts.
|
|
import multiprocess # noqa: F401
|
|
except ImportError:
|
|
pass
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
MODULE_PATH = REPO_ROOT / "unsloth" / "dataset_num_proc.py"
|
|
RL_PATH = REPO_ROOT / "unsloth" / "models" / "rl.py"
|
|
|
|
# The dotted paths unsloth/models/rl.py bakes into every generated trainer: the
|
|
# zoo first, so generated source does not import back into unsloth, then this
|
|
# package as the fallback for a zoo that predates the module.
|
|
GENERATED_IMPORT_MODULE = "unsloth_zoo.dataset_num_proc"
|
|
GENERATED_FALLBACK_MODULE = "unsloth.dataset_num_proc"
|
|
GENERATED_IMPORT_NAME = "get_dataset_num_proc"
|
|
|
|
|
|
def _load_module():
|
|
spec = importlib.util.spec_from_file_location(
|
|
"unsloth_dataset_num_proc_under_test", MODULE_PATH
|
|
)
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
@pytest.fixture
|
|
def dnp(monkeypatch):
|
|
module = _load_module()
|
|
module.reset_warning_state()
|
|
monkeypatch.delenv(module.NUM_PROC_ENV_VAR, raising = False)
|
|
# Pin the platform: macOS is refused by policy whatever the start method says,
|
|
# so every assertion expecting a worker count is really about a forking
|
|
# platform. Platform tests set their own value afterwards, which wins.
|
|
monkeypatch.setattr(module.sys, "platform", "linux")
|
|
# Pin the memory ceiling too, at its two sources rather than at the reader,
|
|
# so the memory tests can still patch either and win. Every count this
|
|
# module returns is clamped by free RAM and by the cgroup budget, so on a
|
|
# memory-limited runner a test about the start method or about an explicit
|
|
# value silently becomes a test of the clamp instead:
|
|
# `get_dataset_num_proc(6) == 6` returns 4 in a small container.
|
|
try:
|
|
import psutil
|
|
monkeypatch.setattr(
|
|
psutil, "virtual_memory", lambda: type("m", (), {"available": 1024 * 1024**3})()
|
|
)
|
|
except ImportError:
|
|
pass
|
|
# Point this module's cgroup reader at a path that does not exist rather than
|
|
# stubbing the reader itself, so the tests that are about it can still
|
|
# install a fixture tree and win.
|
|
monkeypatch.setattr(module, "CGROUP_ROOT", "/nonexistent-cgroup-root-for-tests")
|
|
# unsloth_zoo's readers are neutralised by name instead, and without
|
|
# requiring the name to be there: pinning hf_xet_tuning.CGROUP_ROOT alone
|
|
# only works on a zoo that has that global. An older zoo still exposes the
|
|
# private dir helpers this module prefers, monkeypatch finds no attribute to
|
|
# pin, and the policy reads the runner's real cgroup -- which in a
|
|
# memory-limited container turns a test about the start method into a test
|
|
# of the clamp. The tests that are about these readers install their own
|
|
# unsloth_zoo.hf_xet_tuning in sys.modules, so they are unaffected.
|
|
try:
|
|
from unsloth_zoo import hf_xet_tuning
|
|
except Exception:
|
|
# Importing the package can fail long after it has imported this
|
|
# submodule (__init__ pulls in hf_xet_tuning near the top and only
|
|
# raises "Please install Unsloth" at the end), and the failure removes
|
|
# unsloth_zoo from sys.modules while leaving unsloth_zoo.hf_xet_tuning
|
|
# behind. The module under test reaches it through exactly that cache
|
|
# entry, so treating the failure as absence would leave the real readers
|
|
# live on the runner's own /sys/fs/cgroup.
|
|
hf_xet_tuning = sys.modules.get("unsloth_zoo.hf_xet_tuning")
|
|
if hf_xet_tuning is not None:
|
|
for name, neutral in (
|
|
("CGROUP_ROOT", Path("/nonexistent-cgroup-root-for-tests")),
|
|
("_cgroup_v2_dirs", lambda: []),
|
|
("_cgroup_v1_dirs", lambda controller: []),
|
|
("cgroup_memory_limit", lambda: None),
|
|
("cgroup_cpu_limit", lambda: None),
|
|
):
|
|
monkeypatch.setattr(hf_xet_tuning, name, neutral, raising = False)
|
|
return module
|
|
|
|
|
|
def _force_start_method(monkeypatch, dnp, method):
|
|
monkeypatch.setattr(dnp, "multiprocessing_start_method", lambda: method)
|
|
|
|
|
|
def _force_cpus(monkeypatch, dnp, count):
|
|
"""Pin the CPU count the auto path sizes from.
|
|
|
|
Patching psutil alone is not enough: the count is the smallest of the host's
|
|
CPUs, this process's affinity mask and any cgroup quota, so on a 4-vCPU runner
|
|
a "128 CPU host" would still come out as 4.
|
|
"""
|
|
monkeypatch.setattr(dnp, "_usable_cpus", lambda: count)
|
|
|
|
|
|
# ---------- start-method veto ----------
|
|
|
|
|
|
@pytest.mark.parametrize("method", ["spawn", "forkserver", None])
|
|
def test_non_fork_start_method_disables_multiprocessing(monkeypatch, dnp, method):
|
|
# Under spawn/forkserver the child must re-import the dynamically generated
|
|
# trainer module, which has no importable name, so workers cannot run.
|
|
_force_start_method(monkeypatch, dnp, method)
|
|
assert dnp.get_dataset_num_proc(8) is None
|
|
assert dnp.get_dataset_num_proc(None) is None
|
|
|
|
|
|
def test_non_fork_start_method_warns_once(monkeypatch, dnp, capsys):
|
|
# Regression for eeffa4c065: an explicit value used to sail through the
|
|
# guard. It must be vetoed, and the veto must be visible.
|
|
_force_start_method(monkeypatch, dnp, "spawn")
|
|
dnp.get_dataset_num_proc(8)
|
|
dnp.get_dataset_num_proc(8)
|
|
out = capsys.readouterr().out
|
|
assert out.count("uses the 'spawn' start method") == 1
|
|
assert "dataset_num_proc = 8" in out
|
|
|
|
|
|
def test_fork_start_method_honours_explicit_value(monkeypatch, dnp):
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
assert dnp.get_dataset_num_proc(6) == 6
|
|
|
|
|
|
# ---------- the 1 -> None normalisation ----------
|
|
|
|
|
|
@pytest.mark.parametrize("value", [1, 0, -4])
|
|
def test_non_positive_and_one_normalise_to_none(monkeypatch, dnp, value):
|
|
# `1` is a trap: callers mean "serial", datasets >= 4.0 gives a Pool(1).
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
assert dnp.get_dataset_num_proc(value) is None
|
|
|
|
|
|
def test_serial_as_none_false_preserves_an_explicit_one(monkeypatch, dnp):
|
|
"""The config layer must not collapse 1 to None.
|
|
|
|
unsloth_zoo.sft_prepare_dataset reads a config ``None`` as "auto-size me",
|
|
so writing None back for a user who asked for 1 would inflate it.
|
|
"""
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
assert dnp.get_dataset_num_proc(1, serial_as_none = False) == 1
|
|
# 0 and negatives are incoherent requests but still mean "not parallel", so
|
|
# they land on the config serial sentinel (1), not on None.
|
|
assert dnp.get_dataset_num_proc(0, serial_as_none = False) == 1
|
|
assert dnp.get_dataset_num_proc(-4, serial_as_none = False) == 1
|
|
|
|
|
|
def test_config_layer_never_returns_none_while_forking_is_available(monkeypatch, dnp):
|
|
"""On a fork host no path may write None back to a config.
|
|
|
|
None means "auto-size me" downstream, so any route to it -- memory clamp,
|
|
explicit serial -- would re-inflate.
|
|
"""
|
|
psutil = pytest.importorskip("psutil")
|
|
_force_cpus(monkeypatch, dnp, 64)
|
|
|
|
# memory clamp all the way down to serial
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
monkeypatch.setattr(
|
|
psutil, "virtual_memory", lambda: type("m", (), {"available": 1 * 1024**3})()
|
|
)
|
|
assert dnp.get_dataset_num_proc(16, serial_as_none = False) == 1
|
|
assert dnp.get_dataset_num_proc(None, serial_as_none = False) == 1
|
|
|
|
|
|
@pytest.mark.parametrize("method", ["spawn", "forkserver", None])
|
|
@pytest.mark.parametrize("desired", [None, 1, 16])
|
|
def test_config_layer_is_none_not_one_on_a_non_fork_start_method(monkeypatch, dnp, method, desired):
|
|
"""Regression: the config sentinel 1 reached unpatched TRL map() call sites.
|
|
|
|
Only SFT gets its map site rewritten (rl_replacements.py); DPO, KTO, CPO,
|
|
ORPO, Reward and PRM pass ``args.dataset_num_proc`` straight into
|
|
``Dataset.map``, where a ``1`` builds a ``Pool(1)`` whose spawned child
|
|
re-imports the user's ``__main__`` (#3211 / #3397). None is safe here
|
|
precisely because forking is unavailable, so every auto-sizer vetoes too.
|
|
"""
|
|
_force_start_method(monkeypatch, dnp, method)
|
|
assert dnp.get_dataset_num_proc(desired, serial_as_none = False) is None
|
|
|
|
|
|
def test_config_layer_env_forced_serial_is_none_on_a_non_fork_start_method(monkeypatch, dnp):
|
|
"""UNSLOTH_DATASET_NUM_PROC=0 must not build a Pool(1) either."""
|
|
_force_start_method(monkeypatch, dnp, "spawn")
|
|
monkeypatch.setenv(dnp.NUM_PROC_ENV_VAR, "0")
|
|
assert dnp.get_dataset_num_proc(None, serial_as_none = False) is None
|
|
|
|
|
|
def test_layering_config_then_map_site_is_correct(monkeypatch, dnp):
|
|
"""Composing the two layers must land on the right value for each intent."""
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
psutil = pytest.importorskip("psutil")
|
|
_force_cpus(monkeypatch, dnp, 32)
|
|
monkeypatch.setattr(
|
|
psutil,
|
|
"virtual_memory",
|
|
lambda: type("m", (), {"available": 256 * 1024**3})(),
|
|
)
|
|
cfg = lambda v: dnp.get_dataset_num_proc(v, serial_as_none = False) # noqa: E731
|
|
site = dnp.get_dataset_num_proc
|
|
|
|
# user asked for serial -> stays serial, never auto-inflated
|
|
assert site(cfg(1)) is None
|
|
# user asked for a specific count -> honoured end to end
|
|
assert site(cfg(6)) == 6
|
|
# user asked for nothing -> capped auto, and re-applying is idempotent
|
|
assert cfg(None) == dnp.AUTO_NUM_PROC_CAP
|
|
assert site(cfg(None)) == dnp.AUTO_NUM_PROC_CAP
|
|
|
|
|
|
def test_low_memory_auto_path_returns_none_not_one(monkeypatch, dnp):
|
|
# The old heuristic returned 1 here, which still forked a Pool(1).
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
psutil = pytest.importorskip("psutil")
|
|
_force_cpus(monkeypatch, dnp, 32)
|
|
monkeypatch.setattr(
|
|
psutil,
|
|
"virtual_memory",
|
|
lambda: type("m", (), {"available": 1 * 1024**3})(),
|
|
)
|
|
assert dnp.get_dataset_num_proc(None) is None
|
|
|
|
|
|
# ---------- auto sizing ----------
|
|
|
|
|
|
def test_auto_value_is_capped(monkeypatch, dnp):
|
|
# Was min(max(cpu_count + 4, 2), 64) -- up to 64 forked workers.
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
psutil = pytest.importorskip("psutil")
|
|
_force_cpus(monkeypatch, dnp, 128)
|
|
monkeypatch.setattr(
|
|
psutil,
|
|
"virtual_memory",
|
|
lambda: type("m", (), {"available": 512 * 1024**3})(),
|
|
)
|
|
assert dnp.get_dataset_num_proc(None) == dnp.AUTO_NUM_PROC_CAP
|
|
assert dnp.AUTO_NUM_PROC_CAP < 64
|
|
|
|
|
|
def test_auto_value_clamped_by_available_memory(monkeypatch, dnp):
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
psutil = pytest.importorskip("psutil")
|
|
_force_cpus(monkeypatch, dnp, 64)
|
|
monkeypatch.setattr(
|
|
psutil,
|
|
"virtual_memory",
|
|
lambda: type("m", (), {"available": 10 * 1024**3})(),
|
|
)
|
|
# 10 GB free, half of it budgeted, ~1 GB per worker -> 5.
|
|
assert dnp.get_dataset_num_proc(None) == 5
|
|
|
|
|
|
def test_explicit_value_is_clamped_by_memory(monkeypatch, dnp, capsys):
|
|
"""The gap that caused issue #2693.
|
|
|
|
Studio passes an explicit ``max(1, cpu_count // 4)``, dozens of workers at
|
|
~680 MB each on a big-core machine, and the old heuristic bounded only the
|
|
auto path, so that sailed through however little RAM there was.
|
|
"""
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
psutil = pytest.importorskip("psutil")
|
|
monkeypatch.setattr(
|
|
psutil,
|
|
"virtual_memory",
|
|
lambda: type("m", (), {"available": 16 * 1024**3})(),
|
|
)
|
|
assert dnp.get_dataset_num_proc(48) == 8
|
|
assert "reducing dataset_num_proc 48 -> 8" in capsys.readouterr().out
|
|
|
|
|
|
def test_explicit_value_is_not_capped_by_the_auto_cap(monkeypatch, dnp):
|
|
# AUTO_NUM_PROC_CAP bounds auto-sizing only.
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
psutil = pytest.importorskip("psutil")
|
|
monkeypatch.setattr(
|
|
psutil,
|
|
"virtual_memory",
|
|
lambda: type("m", (), {"available": 512 * 1024**3})(),
|
|
)
|
|
assert dnp.get_dataset_num_proc(32) == 32
|
|
assert 32 > dnp.AUTO_NUM_PROC_CAP
|
|
|
|
|
|
def test_memory_clamp_is_skipped_without_psutil(monkeypatch, dnp):
|
|
# No psutil means no memory reading, so honour the request.
|
|
monkeypatch.setattr(dnp, "_affordable_workers", lambda: None)
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
assert dnp.get_dataset_num_proc(32) == 32
|
|
|
|
|
|
def test_bool_is_not_treated_as_an_int(monkeypatch, dnp):
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
psutil = pytest.importorskip("psutil")
|
|
_force_cpus(monkeypatch, dnp, 8)
|
|
monkeypatch.setattr(
|
|
psutil,
|
|
"virtual_memory",
|
|
lambda: type("m", (), {"available": 64 * 1024**3})(),
|
|
)
|
|
assert dnp.get_dataset_num_proc(True) == 4
|
|
|
|
|
|
# ---------- environment escape hatch ----------
|
|
|
|
|
|
def test_env_override_beats_start_method_veto(monkeypatch, dnp):
|
|
# A user who knows their workload is fork-safe is never downgraded.
|
|
_force_start_method(monkeypatch, dnp, "spawn")
|
|
monkeypatch.setenv(dnp.NUM_PROC_ENV_VAR, "24")
|
|
assert dnp.get_dataset_num_proc(None) == 24
|
|
|
|
|
|
@pytest.mark.parametrize("raw", ["0", "none", "None", "false", ""])
|
|
def test_env_override_can_force_in_process(monkeypatch, dnp, raw):
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
monkeypatch.setenv(dnp.NUM_PROC_ENV_VAR, raw)
|
|
assert dnp.get_dataset_num_proc(16) is None
|
|
|
|
|
|
@pytest.mark.parametrize("raw", ["0", "none", "None", "false", "", "1"])
|
|
def test_env_override_in_process_is_encoded_for_the_config_layer(monkeypatch, dnp, raw):
|
|
# Regression: the env override used to return before _serial(), writing None
|
|
# into the *config*, which unsloth_zoo.sft_prepare_dataset reads as
|
|
# "auto-size me" -- so the hatch the dead-worker message recommends raised
|
|
# the worker count instead of removing it. Config serial is 1, never None.
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
monkeypatch.setenv(dnp.NUM_PROC_ENV_VAR, raw)
|
|
assert dnp.get_dataset_num_proc(16, serial_as_none = False) == 1
|
|
|
|
|
|
def test_env_override_is_uncapped(monkeypatch, dnp):
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
monkeypatch.setenv(dnp.NUM_PROC_ENV_VAR, "100")
|
|
assert dnp.get_dataset_num_proc(None) == 100
|
|
# Above the auto cap is the easy half. The memory clamp is the one that
|
|
# matters: the fixture leaves room for 512 workers, so without pinning it
|
|
# this asserts nothing about the exemption. map_failure_diagnostics points
|
|
# users at this hatch on exactly the host where the clamp would bite.
|
|
monkeypatch.setattr(dnp, "_affordable_workers", lambda: 2)
|
|
assert dnp.get_dataset_num_proc(None) == 100
|
|
assert dnp.get_dataset_num_proc(4) == 100
|
|
|
|
|
|
def test_invalid_env_override_is_ignored_with_a_warning(monkeypatch, dnp, capsys):
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
monkeypatch.setenv(dnp.NUM_PROC_ENV_VAR, "banana")
|
|
assert dnp.get_dataset_num_proc(4) == 4
|
|
assert "is not an integer" in capsys.readouterr().out
|
|
|
|
|
|
# ---------- start-method probing must not mutate global state ----------
|
|
|
|
|
|
def test_start_method_probe_prefers_multiprocess_and_has_no_side_effects(dnp):
|
|
"""datasets does `from multiprocess import Pool`, so `multiprocess` -- not
|
|
stdlib multiprocessing -- decides how map() spawns. Reading it must also not
|
|
pin the context, which would make a later set_start_method() raise."""
|
|
multiprocess = pytest.importorskip("multiprocess")
|
|
import multiprocessing
|
|
|
|
before_mp = multiprocess.get_start_method(allow_none = True)
|
|
before_std = multiprocessing.get_start_method(allow_none = True)
|
|
|
|
method = dnp.multiprocessing_start_method()
|
|
|
|
# Must name a method this host offers. The private default-context chain has
|
|
# answered "fork" on Windows, which offers only spawn; believing it would
|
|
# read Windows as forkable and let workers through.
|
|
assert method in multiprocess.get_all_start_methods()
|
|
assert multiprocess.get_start_method(allow_none = True) == before_mp
|
|
assert multiprocessing.get_start_method(allow_none = True) == before_std
|
|
|
|
|
|
def test_start_method_probe_reports_an_explicit_setting(monkeypatch, dnp):
|
|
import sys as _sys
|
|
import types
|
|
|
|
fake = types.ModuleType("multiprocess")
|
|
fake.get_start_method = lambda allow_none = False: "forkserver"
|
|
fake.get_all_start_methods = lambda: ["fork", "spawn", "forkserver"]
|
|
monkeypatch.setitem(_sys.modules, "multiprocess", fake)
|
|
assert dnp.multiprocessing_start_method() == "forkserver"
|
|
|
|
|
|
def _fake_multiprocess(listed, default_name):
|
|
"""A multiprocess stand-in with nothing pinned yet."""
|
|
import types
|
|
|
|
fake = types.ModuleType("multiprocess")
|
|
fake.get_start_method = lambda allow_none = False: None
|
|
fake.get_all_start_methods = lambda: list(listed)
|
|
if default_name is not None:
|
|
context = types.ModuleType("multiprocess.context")
|
|
context._default_context = types.SimpleNamespace(
|
|
_default_context = types.SimpleNamespace(_name = default_name),
|
|
_actual_context = None,
|
|
)
|
|
fake.context = context
|
|
return fake
|
|
|
|
|
|
def test_start_method_probe_prefers_the_real_default_over_list_order(monkeypatch, dnp):
|
|
"""macOS: multiprocess lists 'spawn' first but its default context is fork.
|
|
|
|
It copies stdlib ``get_all_start_methods()`` verbatim, darwin branch
|
|
included, while keeping ``fork`` as its ``_default_context`` on every POSIX
|
|
platform (``#FIXME: spawn`` in multiprocess/context.py). Since ``datasets``
|
|
pools come from ``multiprocess``, trusting the list would veto every worker
|
|
and misreport the start method in the dead-worker diagnostics.
|
|
"""
|
|
import sys as _sys
|
|
|
|
darwin_order = ["spawn", "fork", "forkserver"]
|
|
fake = _fake_multiprocess(darwin_order, "fork")
|
|
monkeypatch.setitem(_sys.modules, "multiprocess", fake)
|
|
assert dnp.multiprocessing_start_method() == "fork"
|
|
|
|
|
|
def test_macos_stays_in_process_even_though_multiprocess_forks(monkeypatch, dnp, capsys):
|
|
"""The probe reports fork on macOS; policy still refuses to use it.
|
|
|
|
``multiprocess`` really does fork on darwin, so the diagnostics must say so,
|
|
but CPython moved the macOS default to spawn in 3.8 (bpo-33725) because
|
|
forking there "can lead to crashes of the subprocess as macOS system
|
|
libraries may start threads" -- and this parent holds Torch and a threaded
|
|
BLAS. Fixing the probe without the guard would take macOS from always-serial
|
|
to AUTO_NUM_PROC_CAP forked workers.
|
|
"""
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
# Pin memory so the contrast is about policy, not the runner's free RAM.
|
|
monkeypatch.setattr(dnp, "_affordable_workers", lambda: 1000)
|
|
monkeypatch.setattr(dnp.sys, "platform", "darwin")
|
|
|
|
# Serial at a map() call site, None -- not 1 -- at the config layer, so no
|
|
# Pool is built on datasets >= 4.1 either way.
|
|
assert dnp.get_dataset_num_proc(8) is None
|
|
assert dnp.get_dataset_num_proc(8, serial_as_none = False) is None
|
|
assert dnp.get_dataset_num_proc(None) is None
|
|
assert "macOS" in capsys.readouterr().out
|
|
|
|
# The escape hatch still overrides it, and Linux is unaffected.
|
|
monkeypatch.setenv(dnp.NUM_PROC_ENV_VAR, "4")
|
|
assert dnp.get_dataset_num_proc(8) == 4
|
|
monkeypatch.delenv(dnp.NUM_PROC_ENV_VAR)
|
|
monkeypatch.setattr(dnp.sys, "platform", "linux")
|
|
assert dnp.get_dataset_num_proc(8) == 8
|
|
|
|
|
|
def test_start_method_probe_falls_back_to_list_order(monkeypatch, dnp):
|
|
"""The default context is private, so an unreadable one must not raise."""
|
|
import sys as _sys
|
|
|
|
fake = _fake_multiprocess(["spawn", "fork"], None)
|
|
monkeypatch.setitem(_sys.modules, "multiprocess", fake)
|
|
assert dnp.multiprocessing_start_method() == "spawn"
|
|
|
|
|
|
def test_start_method_probe_matches_the_pool_multiprocess_would_build(dnp):
|
|
"""The probe must agree with multiprocess's own default on this host."""
|
|
multiprocess = pytest.importorskip("multiprocess")
|
|
if multiprocess.get_start_method(allow_none = True) is not None:
|
|
pytest.skip("a start method is already pinned in this process")
|
|
assert (
|
|
dnp.multiprocessing_start_method()
|
|
== multiprocess.context._default_context._default_context._name
|
|
)
|
|
|
|
|
|
# ---------- the generated trainer's import must match the real module ----------
|
|
|
|
|
|
def _rl_serial_as_none(tree, source, trainer_file):
|
|
"""Evaluate rl.py's own serial_as_none rule rather than restating it.
|
|
|
|
Restating it made every codegen case below self-fulfilling: they passed
|
|
whatever rl.py decided, so flipping SFT to True -- losing the config
|
|
sentinel these modules exist to protect -- was caught only by the one test
|
|
that matches the emitted literal, which is the line a developer would edit
|
|
in the same change.
|
|
"""
|
|
for node in ast.walk(tree):
|
|
if isinstance(node, ast.Assign) and getattr(node.targets[0], "id", "") == "_serial_as_none":
|
|
# Parenthesised for the same reason as the sibling below: a
|
|
# formatter may reflow the ternary, and the segment's continuation
|
|
# lines are an IndentationError on their own.
|
|
return eval( # noqa: S307
|
|
"(" + ast.get_source_segment(source, node.value) + ")",
|
|
{"trainer_file": trainer_file},
|
|
)
|
|
raise AssertionError("_serial_as_none assignment not found in unsloth/models/rl.py")
|
|
|
|
|
|
def _rl_num_proc_snippet(trainer_file = "sft_trainer"):
|
|
"""The snippet rl.py would splice into that trainer's generated config.
|
|
|
|
The expression is evaluated rather than literal_eval'd because it is no
|
|
longer one literal: the serial encoding depends on the trainer being
|
|
patched, and the point of these tests is what each one ends up with.
|
|
"""
|
|
source = RL_PATH.read_text(encoding = "utf-8")
|
|
tree = ast.parse(source)
|
|
for node in ast.walk(tree):
|
|
if isinstance(node, ast.Assign) and getattr(node.targets[0], "id", "") == "num_proc_check":
|
|
# Parenthesised: the segment starts at the first literal, so its
|
|
# continuation lines are an IndentationError on their own.
|
|
expression = "(" + ast.get_source_segment(source, node.value) + ")"
|
|
return eval( # noqa: S307
|
|
expression, {"_serial_as_none": _rl_serial_as_none(tree, source, trainer_file)}
|
|
)
|
|
raise AssertionError("num_proc_check literal not found in unsloth/models/rl.py")
|
|
|
|
|
|
def test_rl_codegen_writes_back_without_collapsing_serial():
|
|
# SFT's config is read by an auto-sizer, so serial has to survive as 1.
|
|
assert "serial_as_none = False" in _rl_num_proc_snippet("sft_trainer")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"trainer_file",
|
|
["dpo_trainer", "kto_trainer", "cpo_trainer", "orpo_trainer", "reward_trainer", "prm_trainer"],
|
|
)
|
|
def test_rl_codegen_keeps_serial_as_none_where_the_config_reaches_map(trainer_file):
|
|
"""Only SFT has a downstream auto-sizer to defend the sentinel against.
|
|
|
|
These trainers hand args.dataset_num_proc straight to Dataset.map, so a
|
|
config `1` is a Pool(1) on datasets >= 4.1 -- one worker holding its own
|
|
tokenizer copy, on the low-memory host that just refused workers. Nothing
|
|
inflates a None there, so None is what serial must be.
|
|
"""
|
|
assert "serial_as_none = True" in _rl_num_proc_snippet(trainer_file)
|
|
|
|
|
|
def test_rl_codegen_only_sft_gets_the_config_sentinel():
|
|
# Pin the discriminator itself: a rename of the trainer file would silently
|
|
# give every trainer the None encoding, including the one that must not have it.
|
|
source = RL_PATH.read_text(encoding = "utf-8")
|
|
assert '_serial_as_none = "False" if trainer_file == "sft_trainer" else "True"' in source
|
|
|
|
|
|
# ---------- the map-site rewrite and the anchors it hangs on ----------
|
|
|
|
# The tag rl_replacements.py gives the num_proc edit; both anchors share it, so a
|
|
# rename shows up as a missing call rather than a silently unpinned anchor.
|
|
NUM_PROC_WHERE = "sft_prepare_dataset dataset_num_proc selection"
|
|
|
|
# The helpers that decide whether a source edit lands. Named here so the tests
|
|
# below fail loudly if they are renamed rather than quietly testing nothing.
|
|
ANCHOR_HELPERS = ("_require_replace", "_replace_or_fallback", "_same_source")
|
|
|
|
# The module-level regex the narrow fallback anchor uses.
|
|
NARROW_ANCHOR_NAME = "_ZOO_MAP_NUM_PROC_ASSIGNMENT"
|
|
|
|
|
|
def _zoo_dataset_utils_source():
|
|
# find_spec resolves the package path without executing its __init__, so these
|
|
# canaries still run where torch/unsloth_zoo cannot import.
|
|
spec = importlib.util.find_spec("unsloth_zoo")
|
|
if spec is None or not spec.submodule_search_locations:
|
|
pytest.skip("unsloth_zoo not installed")
|
|
zoo_file = Path(list(spec.submodule_search_locations)[0]) / "dataset_utils.py"
|
|
if not zoo_file.is_file():
|
|
pytest.skip("unsloth_zoo.dataset_utils not found")
|
|
return zoo_file.read_text(encoding = "utf-8")
|
|
|
|
|
|
def _rl_replacements_tree():
|
|
return ast.parse(RL_PATH.with_name("rl_replacements.py").read_text(encoding = "utf-8"))
|
|
|
|
|
|
def _anchor_calls(where):
|
|
return [
|
|
node
|
|
for node in ast.walk(_rl_replacements_tree())
|
|
if isinstance(node, ast.Call)
|
|
and getattr(node.func, "id", "") in ANCHOR_HELPERS
|
|
and any(k.arg == "where" and ast.literal_eval(k.value) == where for k in node.keywords)
|
|
]
|
|
|
|
|
|
def _anchor_and_count(where):
|
|
"""The (anchor, expected occurrences) of the source edit tagged ``where``."""
|
|
found = _anchor_calls(where)
|
|
assert len(found) == 1, f"expected exactly one anchored edit for {where!r}"
|
|
node = found[0]
|
|
count = next((ast.literal_eval(k.value) for k in node.keywords if k.arg == "count"), 1)
|
|
return ast.literal_eval(node.args[1]), count
|
|
|
|
|
|
def _keyword(where, name):
|
|
node = _anchor_calls(where)[0]
|
|
value = next(k.value for k in node.keywords if k.arg == name)
|
|
return value
|
|
|
|
|
|
def _narrow_num_proc_pattern():
|
|
"""The compiled fallback regex, read out of rl_replacements.py by name.
|
|
|
|
Read from the module-level ``re.compile`` rather than re-typed here: a test
|
|
holding its own copy of the pattern would keep passing after the real one
|
|
drifted away from the Zoo.
|
|
"""
|
|
name = _keyword(NUM_PROC_WHERE, "fallback_pattern")
|
|
assert (
|
|
getattr(name, "id", "") == NARROW_ANCHOR_NAME
|
|
), f"the num_proc fallback no longer uses {NARROW_ANCHOR_NAME}"
|
|
for node in ast.walk(_rl_replacements_tree()):
|
|
if (
|
|
isinstance(node, ast.Assign)
|
|
and any(getattr(t, "id", None) == NARROW_ANCHOR_NAME for t in node.targets)
|
|
and isinstance(node.value, ast.Call)
|
|
):
|
|
args = [ast.literal_eval(a) for a in node.value.args]
|
|
flags = next((k for k in node.value.keywords if k.arg == "flags"), None)
|
|
assert flags is not None and "MULTILINE" in ast.dump(
|
|
flags.value
|
|
), "the narrow anchor must be MULTILINE to match a line in a block"
|
|
return re.compile(args[0], flags = re.MULTILINE)
|
|
raise AssertionError(f"{NARROW_ANCHOR_NAME} not found in rl_replacements.py")
|
|
|
|
|
|
def _narrow_num_proc_replacement():
|
|
return ast.literal_eval(_keyword(NUM_PROC_WHERE, "fallback_new"))
|
|
|
|
|
|
def test_zoo_sft_prepare_dataset_anchor_has_not_drifted():
|
|
"""unsloth/models/rl_replacements.py rewrites unsloth_zoo's
|
|
sft_prepare_dataset by exact string match. A Zoo release that touches those
|
|
lines makes _require_replace raise at import time, so catch drift here."""
|
|
source = _zoo_dataset_utils_source()
|
|
|
|
# _require_replace raises on a missing anchor but cannot notice a count = 2
|
|
# anchor dropping to one occurrence, so assert the counts here.
|
|
for where in (
|
|
NUM_PROC_WHERE,
|
|
"sft_prepare_dataset tokenizing map() calls",
|
|
):
|
|
anchor, count = _anchor_and_count(where)
|
|
assert source.count(anchor) == count, (
|
|
f"unsloth_zoo.dataset_utils has {source.count(anchor)} occurrences of "
|
|
f"the {where!r} anchor, expected {count}; update rl_replacements.py"
|
|
)
|
|
|
|
|
|
def test_the_narrow_num_proc_anchor_still_matches_the_installed_zoo():
|
|
"""The num_proc site has a second, narrower anchor; it needs its own canary.
|
|
|
|
The block anchor drifting is survivable precisely because this regex catches
|
|
the assignment the block ends on. If the Zoo ever renames or restructures
|
|
that line too, both anchors are gone and nothing rewrites the worker count --
|
|
so pin it here, in CI, rather than discovering it from a user's Pool(1).
|
|
"""
|
|
source = _zoo_dataset_utils_source()
|
|
pattern = _narrow_num_proc_pattern()
|
|
|
|
matches = pattern.findall(source)
|
|
assert len(matches) == 1, (
|
|
f"unsloth_zoo.dataset_utils has {len(matches)} lines matching the narrow "
|
|
f"num_proc anchor {pattern.pattern!r}, expected 1; update rl_replacements.py"
|
|
)
|
|
|
|
# The block anchor and the narrow anchor have to describe the same site, or
|
|
# the fallback would rewrite something the primary edit never touched.
|
|
block_anchor, _ = _anchor_and_count(NUM_PROC_WHERE)
|
|
assert pattern.search(block_anchor) is not None
|
|
|
|
# And the fallback has to leave the file parseable at the Zoo's indentation.
|
|
rewritten = pattern.sub(_narrow_num_proc_replacement(), source)
|
|
assert rewritten != source
|
|
ast.parse(rewritten)
|
|
|
|
|
|
# ---------- what the layered anchor actually does to a drifted Zoo ----------
|
|
|
|
|
|
def _load_anchor_helpers():
|
|
"""Exec just the anchor helpers out of rl_replacements.py.
|
|
|
|
That module imports torch and trl at its top and this file is torch-free by
|
|
design, so lift the definitions the anchor logic needs instead of copying
|
|
them: a test carrying its own copy would keep passing after the real helper
|
|
changed. Returns the namespace plus the list the fake logger warns into.
|
|
"""
|
|
tree = _rl_replacements_tree()
|
|
wanted = set(ANCHOR_HELPERS) | {"_warn_once", "_WARNED_MISSING_ANCHORS", NARROW_ANCHOR_NAME}
|
|
kept = [
|
|
node
|
|
for node in tree.body
|
|
if (isinstance(node, ast.FunctionDef) and node.name in wanted)
|
|
or (
|
|
isinstance(node, ast.Assign)
|
|
and any(getattr(t, "id", None) in wanted for t in node.targets)
|
|
)
|
|
]
|
|
assert {node.name for node in kept if isinstance(node, ast.FunctionDef)} >= set(
|
|
ANCHOR_HELPERS
|
|
), "the anchor helpers were renamed"
|
|
|
|
warnings = []
|
|
namespace = {
|
|
"re": re,
|
|
"logger": types.SimpleNamespace(warning = warnings.append),
|
|
}
|
|
module = ast.fix_missing_locations(ast.Module(body = kept, type_ignores = []))
|
|
exec(compile(module, str(RL_PATH.with_name("rl_replacements.py")), "exec"), namespace) # noqa: S102
|
|
return namespace, warnings
|
|
|
|
|
|
def _apply_num_proc_edit(source):
|
|
"""Run the real layered edit for NUM_PROC_WHERE over ``source``."""
|
|
namespace, warnings = _load_anchor_helpers()
|
|
node = _anchor_calls(NUM_PROC_WHERE)[0]
|
|
assert (
|
|
getattr(node.func, "id", "") == "_replace_or_fallback"
|
|
), "the num_proc edit lost its fallback and is a plain replace again"
|
|
result = namespace["_replace_or_fallback"](
|
|
source,
|
|
ast.literal_eval(node.args[1]),
|
|
ast.literal_eval(node.args[2]),
|
|
fallback_pattern = namespace[NARROW_ANCHOR_NAME],
|
|
fallback_new = _narrow_num_proc_replacement(),
|
|
where = NUM_PROC_WHERE,
|
|
)
|
|
return result, warnings
|
|
|
|
|
|
def test_the_block_anchor_is_used_when_the_zoo_has_not_moved():
|
|
source = _zoo_dataset_utils_source()
|
|
result, warnings = _apply_num_proc_edit(source)
|
|
assert "_unsloth_get_dataset_num_proc" in result
|
|
assert 'map_kwargs["num_proc"] = dataset_num_proc' not in result
|
|
# The block replacement takes the Zoo's own sizing with it; the fallback,
|
|
# which only rewrites the assignment, leaves it standing.
|
|
block_anchor, _ = _anchor_and_count(NUM_PROC_WHERE)
|
|
assert block_anchor not in result
|
|
assert warnings == [], f"a matching anchor must not warn: {warnings}"
|
|
ast.parse(result)
|
|
|
|
|
|
def test_the_narrow_anchor_takes_over_when_the_block_drifts():
|
|
"""A Zoo that rewrites the block but keeps the assignment must still be fixed.
|
|
|
|
This is the case `required = False` used to swallow: the edit no-ops, the Zoo
|
|
reads the config `1` the config layer writes for "serial", and datasets >= 4.1
|
|
turns that into a Pool(1) on the host that asked for no workers.
|
|
"""
|
|
source = _zoo_dataset_utils_source()
|
|
# Drift the block without touching the line the fallback keys on.
|
|
drifted = source.replace(
|
|
" import multiprocessing as _mp\n",
|
|
" import multiprocessing as _mp # zoo refactor\n",
|
|
1,
|
|
)
|
|
assert drifted != source
|
|
assert _narrow_num_proc_pattern().findall(drifted)
|
|
|
|
result, warnings = _apply_num_proc_edit(drifted)
|
|
assert "_unsloth_get_dataset_num_proc" in result, "the fallback anchor did not apply"
|
|
assert 'map_kwargs["num_proc"] = dataset_num_proc' not in result
|
|
# Only the assignment was rewritten, so the Zoo's own sizing is still there,
|
|
# computing a value nothing reads. That is the price of the narrow anchor.
|
|
assert "if _mp.get_start_method() != 'fork':" in result
|
|
assert len(warnings) == 1 and "moved in this unsloth_zoo" in warnings[0]
|
|
ast.parse(result)
|
|
|
|
|
|
def test_neither_anchor_matching_only_warns():
|
|
"""The remaining escape hatch stays a warning, not a raise.
|
|
|
|
Hard-failing here would break every SFT run on a Zoo whose text merely moved,
|
|
which is a far bigger blast radius than the one worker at stake.
|
|
"""
|
|
source = (
|
|
_zoo_dataset_utils_source()
|
|
.replace(
|
|
' map_kwargs["num_proc"] = dataset_num_proc\n',
|
|
' map_kwargs.update({"num_proc": dataset_num_proc})\n',
|
|
1,
|
|
)
|
|
.replace(
|
|
" import multiprocessing as _mp\n",
|
|
" import multiprocessing as _mp # zoo refactor\n",
|
|
1,
|
|
)
|
|
)
|
|
assert not _narrow_num_proc_pattern().findall(source)
|
|
|
|
result, warnings = _apply_num_proc_edit(source)
|
|
assert result == source, "nothing should be rewritten when both anchors miss"
|
|
assert len(warnings) == 1 and "anchor not found" in warnings[0]
|
|
|
|
|
|
def test_the_narrow_anchor_keeps_indentation_and_yields_none(monkeypatch):
|
|
"""The injected fallback has to run, at the Zoo's indentation, and give None.
|
|
|
|
Indentation comes from the match, so a Zoo that nests the assignment deeper
|
|
still gets valid source; and the value it computes for a config `1` -- what
|
|
the config layer writes for "serial" -- has to be None, never 1.
|
|
"""
|
|
module = _load_module()
|
|
module.reset_warning_state()
|
|
monkeypatch.delenv(module.NUM_PROC_ENV_VAR, raising = False)
|
|
# The injected snippet imports the zoo copy first; point that name at the
|
|
# copy under test so this stays a torch-free, offline assertion.
|
|
if "unsloth_zoo" not in sys.modules:
|
|
monkeypatch.setitem(sys.modules, "unsloth_zoo", types.ModuleType("unsloth_zoo"))
|
|
monkeypatch.setitem(sys.modules, "unsloth_zoo.dataset_num_proc", module)
|
|
|
|
pattern = _narrow_num_proc_pattern()
|
|
replacement = _narrow_num_proc_replacement()
|
|
|
|
for indent in (" ", " "):
|
|
snippet = pattern.sub(replacement, f'{indent}map_kwargs["num_proc"] = dataset_num_proc')
|
|
for line in snippet.split("\n"):
|
|
assert line.startswith(indent), f"lost the {len(indent)}-space indent: {line!r}"
|
|
|
|
namespace = {
|
|
"map_kwargs": {},
|
|
"args": types.SimpleNamespace(dataset_num_proc = 1),
|
|
}
|
|
exec(compile(textwrap.dedent(snippet), "<fallback>", "exec"), namespace) # noqa: S102
|
|
assert (
|
|
namespace["map_kwargs"]["num_proc"] is None
|
|
), "a config 1 has to become None at the map site; 1 is a Pool(1) on datasets >= 4.1"
|
|
|
|
|
|
def test_rl_codegen_imports_the_module_that_exists():
|
|
# The snippet is spliced into generated source as text, so a rename would
|
|
# otherwise surface only at trainer-construction time in production.
|
|
snippet = _rl_num_proc_snippet()
|
|
assert f"from {GENERATED_IMPORT_MODULE} import {GENERATED_IMPORT_NAME}" in snippet
|
|
assert f"from {GENERATED_FALLBACK_MODULE} import {GENERATED_IMPORT_NAME}" in snippet
|
|
assert snippet.index(GENERATED_IMPORT_MODULE) < snippet.index(
|
|
GENERATED_FALLBACK_MODULE
|
|
), "the zoo has to be tried first; the unsloth copy is only the fallback"
|
|
assert MODULE_PATH.is_file()
|
|
module = _load_module()
|
|
assert callable(getattr(module, GENERATED_IMPORT_NAME))
|
|
|
|
|
|
def test_generated_source_reaches_for_the_zoo_before_unsloth():
|
|
"""Generated trainer source must not import back into unsloth.
|
|
|
|
unsloth/__init__.py generates that source, so a `from unsloth...` there is an
|
|
import of the package mid-flight, and it drags unsloth/utils/__init__.py ->
|
|
packing -> attention_dispatch -> models._utils (torch and the model stack)
|
|
into a module needing none of it. Both call sites must try the zoo first.
|
|
"""
|
|
tree = _rl_replacements_tree()
|
|
injected = []
|
|
for node in ast.walk(tree):
|
|
if not isinstance(node, ast.Call) or getattr(node.func, "id", "") not in ANCHOR_HELPERS:
|
|
continue
|
|
if len(node.args) >= 3 and isinstance(node.args[2], ast.Constant):
|
|
injected.append(ast.literal_eval(node.args[2]))
|
|
# The narrow fallback injects source too, as an re.sub template.
|
|
injected += [
|
|
ast.literal_eval(k.value)
|
|
for k in node.keywords
|
|
if k.arg == "fallback_new" and isinstance(k.value, ast.Constant)
|
|
]
|
|
reaching = [text for text in injected if "dataset_num_proc" in text]
|
|
assert (
|
|
len(reaching) == 3
|
|
), "expected the num_proc selection, its narrow fallback and the map() wrapper"
|
|
for text in reaching:
|
|
assert f"from {GENERATED_IMPORT_MODULE} import" in text
|
|
assert text.index(GENERATED_IMPORT_MODULE) < text.index(
|
|
GENERATED_FALLBACK_MODULE
|
|
), f"this injection imports unsloth before the zoo:\n{text}"
|
|
|
|
|
|
def test_the_two_copies_have_not_drifted():
|
|
"""The zoo owns the policy; this package keeps a fallback copy.
|
|
|
|
Two copies can disagree, and the failure would be silent: whichever import
|
|
wins decides the worker count. Compare them with docstrings stripped, so
|
|
prose may differ per repo but no branch, constant or message may.
|
|
"""
|
|
spec = importlib.util.find_spec("unsloth_zoo")
|
|
if spec is None or not spec.submodule_search_locations:
|
|
pytest.skip("unsloth_zoo not installed")
|
|
zoo_file = Path(list(spec.submodule_search_locations)[0]) / "dataset_num_proc.py"
|
|
if not zoo_file.is_file():
|
|
pytest.skip("this unsloth_zoo predates dataset_num_proc")
|
|
|
|
def _shape(path):
|
|
tree = ast.parse(path.read_text(encoding = "utf-8"))
|
|
for node in ast.walk(tree):
|
|
if not isinstance(
|
|
node, (ast.Module, ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)
|
|
):
|
|
continue
|
|
body = node.body
|
|
if (
|
|
body
|
|
and isinstance(body[0], ast.Expr)
|
|
and isinstance(body[0].value, ast.Constant)
|
|
and isinstance(body[0].value.value, str)
|
|
):
|
|
node.body = body[1:] or [ast.Pass()]
|
|
return ast.dump(ast.parse(ast.unparse(tree)))
|
|
|
|
assert _shape(MODULE_PATH) == _shape(zoo_file), (
|
|
"unsloth/dataset_num_proc.py and unsloth_zoo/dataset_num_proc.py "
|
|
"have diverged; the zoo copy is the source of truth"
|
|
)
|
|
|
|
|
|
def test_rl_codegen_snippet_is_valid_python_at_method_indent():
|
|
# rl.py re-indents extra_args to 8 spaces and drops it into __init__.
|
|
snippet = _rl_num_proc_snippet()
|
|
body = "\n".join(" " * 8 + line for line in snippet.split("\n"))
|
|
source = (
|
|
"class C:\n def __init__(self, dataset_num_proc = None):\n" + body + "\n pass\n"
|
|
)
|
|
ast.parse(source)
|
|
|
|
|
|
def test_rl_codegen_snippet_survives_an_unimportable_helper():
|
|
# A generated file can outlive an unsloth downgrade: constructing a config
|
|
# must still work, just leaving the caller's value alone.
|
|
snippet = _rl_num_proc_snippet()
|
|
namespace = {"dataset_num_proc": 7}
|
|
import builtins
|
|
|
|
real_import = builtins.__import__
|
|
|
|
def _blocked(name, *args, **kwargs):
|
|
if name.startswith("unsloth"):
|
|
raise ImportError("simulated downgrade")
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
builtins.__import__ = _blocked
|
|
try:
|
|
exec(snippet, namespace)
|
|
finally:
|
|
builtins.__import__ = real_import
|
|
assert namespace["dataset_num_proc"] == 7
|
|
|
|
|
|
# ---------- worker-death diagnostics ----------
|
|
|
|
|
|
_DATASETS_MESSAGE = (
|
|
"One of the subprocesses has abruptly died during map operation."
|
|
"To debug the error, disable multiprocessing."
|
|
)
|
|
|
|
|
|
def test_worker_death_is_reraised_with_context(dnp):
|
|
# datasets discards the child's exit status, so the original message cannot
|
|
# distinguish an OOM kill from anything else.
|
|
with pytest.raises(RuntimeError) as caught:
|
|
with dnp.map_failure_diagnostics(8):
|
|
raise RuntimeError(_DATASETS_MESSAGE)
|
|
|
|
message = str(caught.value)
|
|
assert "dataset_num_proc = 8" in message
|
|
assert "8 workers" in message
|
|
assert "8GB" in message, "should estimate what those workers cost"
|
|
assert dnp.NUM_PROC_ENV_VAR in message, "must name the escape hatch"
|
|
assert "out-of-memory" in message
|
|
# The child's traceback must survive for anyone who wants it.
|
|
assert isinstance(caught.value.__cause__, RuntimeError)
|
|
assert _DATASETS_MESSAGE in str(caught.value.__cause__)
|
|
|
|
|
|
def test_worker_death_diagnostics_handles_in_process_runs(dnp):
|
|
# num_proc=None still reaches the wrapper; it must not divide by a None.
|
|
with pytest.raises(RuntimeError) as caught:
|
|
with dnp.map_failure_diagnostics(None):
|
|
raise RuntimeError(_DATASETS_MESSAGE)
|
|
assert "dataset_num_proc = None" in str(caught.value)
|
|
assert "1 worker," in str(caught.value)
|
|
|
|
|
|
def test_unrelated_errors_pass_through_untouched(dnp):
|
|
# Only the dead-worker message is rewritten, and non-RuntimeError types are
|
|
# not caught at all.
|
|
original = RuntimeError("CUDA out of memory")
|
|
with pytest.raises(RuntimeError) as caught:
|
|
with dnp.map_failure_diagnostics(4):
|
|
raise original
|
|
assert caught.value is original
|
|
|
|
key = KeyError("text")
|
|
with pytest.raises(KeyError) as caught_key:
|
|
with dnp.map_failure_diagnostics(4):
|
|
raise key
|
|
assert caught_key.value is key
|
|
|
|
# The identity assertions above hold under `except Exception` as well, since
|
|
# the guard re-raises the same object. This one does not: it carries the
|
|
# dead-worker text, so a widened clause would rewrite it into a RuntimeError.
|
|
lookalike = ValueError("One of the subprocesses has abruptly died during map operation.")
|
|
with pytest.raises(ValueError) as caught_other:
|
|
with dnp.map_failure_diagnostics(4):
|
|
raise lookalike
|
|
assert caught_other.value is lookalike
|
|
|
|
|
|
def test_successful_map_is_not_disturbed(dnp):
|
|
with dnp.map_failure_diagnostics(4):
|
|
result = "tokenized"
|
|
assert result == "tokenized"
|
|
|
|
|
|
def test_studio_num_proc_cap_has_not_drifted(dnp):
|
|
"""Studio duplicates AUTO_NUM_PROC_CAP; the two must stay equal.
|
|
|
|
hardware.py cannot import this module without pulling unsloth's whole
|
|
__init__ into hardware detection, so it carries its own copy. Read it out of
|
|
the source, since importing studio needs a backend environment this lacks.
|
|
"""
|
|
hardware = REPO_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
|
|
if not hardware.is_file():
|
|
pytest.skip("studio backend not present")
|
|
|
|
tree = ast.parse(hardware.read_text(encoding = "utf-8"))
|
|
found = [
|
|
node.value.value
|
|
for node in ast.walk(tree)
|
|
if isinstance(node, ast.Assign)
|
|
and any(getattr(t, "id", "") == "_STUDIO_NUM_PROC_CAP" for t in node.targets)
|
|
and isinstance(node.value, ast.Constant)
|
|
]
|
|
assert len(found) == 1, "expected exactly one _STUDIO_NUM_PROC_CAP assignment"
|
|
assert found[0] == dnp.AUTO_NUM_PROC_CAP, (
|
|
f"studio caps dataset workers at {found[0]} while this module caps the "
|
|
f"auto path at {dnp.AUTO_NUM_PROC_CAP}; they must agree"
|
|
)
|
|
|
|
|
|
def test_studio_bounds_its_own_computed_worker_count(dnp):
|
|
"""A backend heuristic must not outrank the cap.
|
|
|
|
trainer.py asks for cpu_count // 4 and safe_num_proc's auto path is
|
|
cpu_count // 3, so a 64-core host produced 16 workers and a 192-core one 48.
|
|
Those arrive downstream as explicit ints, only clamped by free memory, so a
|
|
roomy machine kept all 48.
|
|
"""
|
|
hardware = REPO_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
|
|
if not hardware.is_file():
|
|
pytest.skip("studio backend not present")
|
|
source = hardware.read_text(encoding = "utf-8")
|
|
|
|
tree = ast.parse(source)
|
|
safe = next(
|
|
(n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) and n.name == "safe_num_proc"),
|
|
None,
|
|
)
|
|
assert safe is not None, "safe_num_proc is where every studio map() count is decided"
|
|
body = ast.dump(safe)
|
|
assert "_STUDIO_NUM_PROC_CAP" in body, (
|
|
"safe_num_proc no longer bounds its result; studio would send "
|
|
"cpu_count // 3 workers straight to Dataset.map"
|
|
)
|
|
|
|
|
|
def test_the_recovery_advice_does_not_promise_more_than_it_delivers(dnp):
|
|
"""UNSLOTH_DATASET_NUM_PROC=0 is not in-process on every installation.
|
|
|
|
The dead-worker message is the one place a user is told what to do next, so
|
|
it has to be true. On fork, train_on_responses_only over the Zoo's threshold
|
|
gets ``1`` -- a bare None would read as "size it for me". The Zoo release
|
|
that ships with this reads that ``1`` as in-process, but an older one turns
|
|
it into a Pool(1), and generated code runs against whichever Zoo is
|
|
installed. Saying "tokenize in-process" flatly was wrong for exactly the
|
|
large-dataset runs that die.
|
|
"""
|
|
with pytest.raises(RuntimeError) as excinfo:
|
|
with dnp.map_failure_diagnostics(8):
|
|
raise RuntimeError("One of the subprocesses has abruptly died during map operation.")
|
|
message = str(excinfo.value)
|
|
assert f"{dnp.NUM_PROC_ENV_VAR}=0" in message
|
|
assert "single worker" in message, "the exception to in-process has to be stated"
|
|
assert "train_on_responses_only" in message, "and which path it applies to"
|
|
assert f"{dnp.ZOO_MIN_ROWS_FOR_MULTIPROC:,}" in message, "and above which size"
|
|
|
|
|
|
def test_the_advice_matches_what_the_resolver_actually_returns(dnp, monkeypatch):
|
|
"""Executed, not just read: drive both branches of the claim above.
|
|
|
|
A message asserting a behaviour the resolver does not have would be worse
|
|
than the vague one it replaced.
|
|
"""
|
|
|
|
class _Split:
|
|
def __init__(self, n):
|
|
self.n = n
|
|
|
|
def __len__(self):
|
|
return self.n
|
|
|
|
class _Trainer:
|
|
def __init__(self, split):
|
|
self.train_dataset = split
|
|
self.eval_dataset = None
|
|
|
|
monkeypatch.setattr(dnp, "multiprocessing_start_method", lambda: "fork")
|
|
monkeypatch.setattr(dnp, "_affordable_workers", lambda: 1000)
|
|
monkeypatch.setattr(dnp.sys, "platform", "linux")
|
|
monkeypatch.setenv(dnp.NUM_PROC_ENV_VAR, "0")
|
|
|
|
over = dnp.resolve_responses_only_num_proc(
|
|
_Trainer(_Split(dnp.ZOO_MIN_ROWS_FOR_MULTIPROC + 1)), None
|
|
)
|
|
under = dnp.resolve_responses_only_num_proc(
|
|
_Trainer(_Split(dnp.ZOO_MIN_ROWS_FOR_MULTIPROC - 1)), None
|
|
)
|
|
assert over == 1, "over the threshold the best expressible request is one worker"
|
|
assert under is None, "under it the Zoo's own guard already goes in-process"
|
|
|
|
|
|
def test_probe_rejects_a_start_method_the_host_does_not_offer(monkeypatch, dnp):
|
|
"""The private default-context chain is not trustworthy on its own.
|
|
|
|
On a Windows runner it answered "fork" while get_all_start_methods() was
|
|
["spawn"]. Believing it reads Windows as forkable, so
|
|
_workers_unusable_reason() returns None and workers get through -- the spawn
|
|
re-import loop of #3211 / #3397.
|
|
"""
|
|
import sys as _sys
|
|
import types
|
|
|
|
fake = types.ModuleType("multiprocess")
|
|
fake.get_start_method = lambda allow_none = False: None
|
|
fake.get_all_start_methods = lambda: ["spawn"]
|
|
context = types.ModuleType("multiprocess.context")
|
|
context._default_context = types.SimpleNamespace(
|
|
_default_context = types.SimpleNamespace(_name = "fork"),
|
|
)
|
|
fake.context = context
|
|
monkeypatch.setitem(_sys.modules, "multiprocess", fake)
|
|
|
|
assert dnp.multiprocessing_start_method() == "spawn"
|
|
|
|
monkeypatch.setattr(dnp.sys, "platform", "win32")
|
|
assert dnp.get_dataset_num_proc(8) is None
|
|
assert dnp.get_dataset_num_proc(8, serial_as_none = False) is None
|
|
|
|
|
|
class _Split:
|
|
"""Minimal sized stand-in for a datasets.Dataset."""
|
|
|
|
def __init__(self, n):
|
|
self.n = n
|
|
|
|
def __len__(self):
|
|
return self.n
|
|
|
|
|
|
# ---------- containers: the host is not what this process may use ----------
|
|
|
|
|
|
def test_memory_budget_follows_the_cgroup_not_the_host(monkeypatch, dnp):
|
|
"""psutil reports the HOST inside a container.
|
|
|
|
A 2GB pod on a 512GB box read as having room for the full worker set and got
|
|
OOM-killed, which is the failure the memory ceiling exists to prevent.
|
|
"""
|
|
psutil = pytest.importorskip("psutil")
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
_force_cpus(monkeypatch, dnp, 64)
|
|
monkeypatch.setattr(
|
|
psutil, "virtual_memory", lambda: type("m", (), {"available": 512 * 1024**3})()
|
|
)
|
|
monkeypatch.setattr(dnp, "_cgroup_free_bytes", lambda: None)
|
|
assert dnp.get_dataset_num_proc(None) == dnp.AUTO_NUM_PROC_CAP
|
|
|
|
monkeypatch.setattr(dnp, "_cgroup_free_bytes", lambda: 2 * 1024**3)
|
|
assert dnp.get_dataset_num_proc(None) is None, "a 2GB container has no room for workers"
|
|
|
|
monkeypatch.setattr(dnp, "_cgroup_free_bytes", lambda: 8 * 1024**3)
|
|
assert dnp.get_dataset_num_proc(None) == 4
|
|
|
|
|
|
def test_memory_already_spent_in_the_container_is_not_counted_as_free(monkeypatch, dnp):
|
|
# Otherwise a container that has already spent most of its limit still reads
|
|
# as having the whole thing available.
|
|
psutil = pytest.importorskip("psutil")
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
_force_cpus(monkeypatch, dnp, 64)
|
|
monkeypatch.setattr(
|
|
psutil, "virtual_memory", lambda: type("m", (), {"available": 512 * 1024**3})()
|
|
)
|
|
monkeypatch.setattr(dnp, "_cgroup_free_bytes", lambda: 32 * 1024**3)
|
|
assert dnp.get_dataset_num_proc(None) == dnp.AUTO_NUM_PROC_CAP
|
|
|
|
# 30 of the 32GB already spent leaves 2, which is not enough for workers.
|
|
monkeypatch.setattr(dnp, "_cgroup_free_bytes", lambda: 2 * 1024**3)
|
|
assert dnp.get_dataset_num_proc(None) is None
|
|
|
|
|
|
def test_cpu_count_follows_the_affinity_mask(monkeypatch, dnp):
|
|
# Under taskset or Slurm pinning the host count is not what this process can
|
|
# run on, and workers would only contend for the cores it does have.
|
|
psutil = pytest.importorskip("psutil")
|
|
monkeypatch.setattr(psutil, "cpu_count", lambda *a, **k: 128)
|
|
monkeypatch.setattr(dnp, "_cgroup_cpu_quota", lambda: None)
|
|
monkeypatch.setattr(dnp.os, "sched_getaffinity", lambda pid: set(range(4)), raising = False)
|
|
assert dnp._usable_cpus() == 4
|
|
|
|
|
|
def test_cpu_count_follows_a_fractional_cgroup_quota(monkeypatch, dnp):
|
|
# Kubernetes "cpu: 500m" is cpu.max "50000 100000" = 0.5 cores. Requiring a
|
|
# whole core would fall back to the host count, so a half-core pod would size
|
|
# workers from every core on the machine.
|
|
psutil = pytest.importorskip("psutil")
|
|
monkeypatch.setattr(psutil, "cpu_count", lambda *a, **k: 128)
|
|
monkeypatch.setattr(dnp.os, "sched_getaffinity", lambda pid: set(range(128)), raising = False)
|
|
monkeypatch.setattr(dnp, "_cgroup_cpu_quota", lambda: 0.5)
|
|
assert dnp._usable_cpus() == 1
|
|
|
|
|
|
def test_a_single_usable_cpu_tokenizes_in_process(monkeypatch, dnp):
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
_force_cpus(monkeypatch, dnp, 1)
|
|
assert dnp.get_dataset_num_proc(None) is None
|
|
|
|
|
|
def test_the_cgroup_readers_never_raise(dnp):
|
|
# They run on every auto-sizing call, on hosts with no cgroup at all.
|
|
free = dnp._cgroup_free_bytes()
|
|
assert free is None or (isinstance(free, int) and free >= 0)
|
|
quota = dnp._cgroup_cpu_quota()
|
|
assert quota is None or isinstance(quota, float)
|
|
|
|
|
|
# ---------- the zoo reads the other module ----------
|
|
|
|
|
|
def _force_stdlib_start_method(monkeypatch, dnp, method):
|
|
real = dnp._module_start_method
|
|
monkeypatch.setattr(
|
|
dnp,
|
|
"_module_start_method",
|
|
lambda name: method if name == "multiprocessing" else real(name),
|
|
)
|
|
|
|
|
|
def test_serial_is_one_when_the_two_modules_disagree(monkeypatch, dnp, capsys):
|
|
"""A None handed to train_on_responses_only is "size it for me" to it.
|
|
|
|
Its auto path asks stdlib multiprocessing. Where that says fork while
|
|
multiprocess says spawn, None is not serial: it picks cpu_count + 4 workers,
|
|
and datasets then builds that pool on the spawn context -- the #3211 / #3397
|
|
re-import loop, multiplied.
|
|
"""
|
|
_force_start_method(monkeypatch, dnp, "spawn")
|
|
_force_stdlib_start_method(monkeypatch, dnp, "fork")
|
|
|
|
trainer = type("t", (), {"train_dataset": _Split(dnp.ZOO_MIN_ROWS_FOR_MULTIPROC * 2)})()
|
|
assert dnp.resolve_responses_only_num_proc(trainer, None) == 1
|
|
assert dnp.resolve_responses_only_num_proc(trainer, 16) == 1
|
|
assert "disagree about the start method" in capsys.readouterr().out
|
|
|
|
|
|
def test_serial_stays_none_when_the_zoo_would_refuse_workers_too(monkeypatch, dnp):
|
|
# macOS: multiprocess forks, stdlib spawns. Its own veto fires, so None is
|
|
# genuinely in-process there and 1 would be a pool it did not need.
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
monkeypatch.setattr(dnp.sys, "platform", "darwin")
|
|
_force_stdlib_start_method(monkeypatch, dnp, "spawn")
|
|
|
|
trainer = type("t", (), {"train_dataset": _Split(dnp.ZOO_MIN_ROWS_FOR_MULTIPROC * 2)})()
|
|
assert dnp.resolve_responses_only_num_proc(trainer, None) is None
|
|
assert dnp.resolve_responses_only_num_proc(trainer, 16) is None
|
|
|
|
|
|
def test_agreeing_modules_are_left_alone(monkeypatch, dnp):
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
_force_stdlib_start_method(monkeypatch, dnp, "fork")
|
|
_force_cpus(monkeypatch, dnp, 32)
|
|
psutil = pytest.importorskip("psutil")
|
|
monkeypatch.setattr(
|
|
psutil, "virtual_memory", lambda: type("m", (), {"available": 256 * 1024**3})()
|
|
)
|
|
monkeypatch.setattr(dnp, "_cgroup_free_bytes", lambda: None)
|
|
|
|
trainer = type("t", (), {"train_dataset": _Split(dnp.ZOO_MIN_ROWS_FOR_MULTIPROC * 2)})()
|
|
assert dnp.resolve_responses_only_num_proc(trainer, None) == dnp.AUTO_NUM_PROC_CAP
|
|
assert dnp.resolve_responses_only_num_proc(trainer, 1) == 1
|
|
|
|
|
|
def _fake_cgroup_module(
|
|
monkeypatch,
|
|
v2_dirs = (),
|
|
v1_dirs = (),
|
|
):
|
|
import types
|
|
|
|
def _read_first_line(path):
|
|
return path.read_text() if path.is_file() else None
|
|
|
|
def _parse_limit(raw):
|
|
if not raw or raw.strip() == "max":
|
|
return None
|
|
try:
|
|
return int(raw.strip())
|
|
except ValueError:
|
|
return None
|
|
|
|
fake = types.ModuleType("unsloth_zoo.hf_xet_tuning")
|
|
fake._cgroup_v2_dirs = lambda: list(v2_dirs)
|
|
fake._cgroup_v1_dirs = lambda controller: list(v1_dirs)
|
|
fake._read_first_line = _read_first_line
|
|
fake._parse_limit = _parse_limit
|
|
fake.cgroup_memory_limit = lambda: None
|
|
fake.cgroup_cpu_limit = lambda: None
|
|
monkeypatch.setitem(sys.modules, "unsloth_zoo.hf_xet_tuning", fake)
|
|
return fake
|
|
|
|
|
|
def test_free_memory_pairs_each_limit_with_its_own_usage(monkeypatch, dnp, tmp_path):
|
|
"""The binding limit is often an ancestor's, and so is the usage that fills it.
|
|
|
|
A leaf's usage against a slice's limit reports memory that siblings have
|
|
already spent as free. The other direction is worse: the root
|
|
memory.current is the whole machine, and against a unit's own MemoryMax it
|
|
leaves every run with nothing.
|
|
"""
|
|
slice_dir = tmp_path / "user.slice"
|
|
leaf = slice_dir / "session.scope"
|
|
leaf.mkdir(parents = True)
|
|
|
|
# The slice caps 32GB and 30 of them are spent, mostly by a sibling; this
|
|
# leaf has a 16GB cap of its own and has spent 1.
|
|
(slice_dir / "memory.max").write_text("34359738368\n")
|
|
(slice_dir / "memory.current").write_text("32212254720\n")
|
|
(leaf / "memory.max").write_text("17179869184\n")
|
|
(leaf / "memory.current").write_text("1073741824\n")
|
|
|
|
_fake_cgroup_module(monkeypatch, v2_dirs = [leaf, slice_dir])
|
|
# 34 - 32 = 2GB from the slice, 16 - 1 = 15GB from the leaf. The slice binds.
|
|
assert dnp._cgroup_free_bytes() == 2 * 1024**3
|
|
|
|
|
|
def test_free_memory_is_never_negative(monkeypatch, dnp, tmp_path):
|
|
# An over-committed cgroup reports more usage than its limit under pressure.
|
|
leaf = tmp_path / "scope"
|
|
leaf.mkdir()
|
|
(leaf / "memory.max").write_text("1073741824\n")
|
|
(leaf / "memory.current").write_text("2147483648\n")
|
|
_fake_cgroup_module(monkeypatch, v2_dirs = [leaf])
|
|
assert dnp._cgroup_free_bytes() == 0
|
|
|
|
|
|
def test_an_unlimited_cgroup_is_not_a_ceiling(monkeypatch, dnp, tmp_path):
|
|
leaf = tmp_path / "scope"
|
|
leaf.mkdir()
|
|
(leaf / "memory.max").write_text("max\n")
|
|
(leaf / "memory.current").write_text("1073741824\n")
|
|
_fake_cgroup_module(monkeypatch, v2_dirs = [leaf])
|
|
assert dnp._cgroup_free_bytes() is None
|
|
|
|
|
|
def test_a_readable_limit_with_no_readable_usage_still_binds(monkeypatch, dnp, tmp_path):
|
|
leaf = tmp_path / "scope"
|
|
leaf.mkdir()
|
|
(leaf / "memory.max").write_text("2147483648\n")
|
|
_fake_cgroup_module(monkeypatch, v2_dirs = [leaf])
|
|
assert dnp._cgroup_free_bytes() == 2 * 1024**3
|
|
|
|
|
|
def _old_zoo(monkeypatch, memory_limit = None):
|
|
"""An unsloth_zoo predating the private cgroup helpers, with only the public reader."""
|
|
import types
|
|
|
|
fake = types.ModuleType("unsloth_zoo.hf_xet_tuning")
|
|
fake.cgroup_memory_limit = lambda: memory_limit
|
|
fake.cgroup_cpu_limit = lambda: None
|
|
monkeypatch.setitem(sys.modules, "unsloth_zoo.hf_xet_tuning", fake)
|
|
return fake
|
|
|
|
|
|
def test_the_unaided_reader_subtracts_usage_too(monkeypatch, dnp, tmp_path):
|
|
"""An older unsloth_zoo must not turn the ceiling back into the raw limit.
|
|
|
|
cgroup_memory_limit() alone reports an 8GB cgroup holding 6GB as 8GB free,
|
|
which is the one case the ceiling exists for: sizing workers off memory that
|
|
is already spent is how #2693's map() children get OOM-killed.
|
|
"""
|
|
_old_zoo(monkeypatch, memory_limit = 8 * 1024**3)
|
|
leaf = tmp_path / "kubepods" / "podabc"
|
|
leaf.mkdir(parents = True)
|
|
(leaf / "memory.max").write_text("8589934592\n")
|
|
(leaf / "memory.current").write_text("6442450944\n")
|
|
|
|
monkeypatch.setattr(dnp, "CGROUP_ROOT", str(tmp_path))
|
|
monkeypatch.setattr(dnp, "_proc_self_cgroup", lambda: ["0::/kubepods/podabc"])
|
|
assert dnp._cgroup_free_bytes() == 2 * 1024**3
|
|
|
|
|
|
def test_the_unaided_reader_walks_to_the_binding_ancestor(monkeypatch, dnp, tmp_path):
|
|
"""Same pairing rule as the helper-backed path: the slice's limit binds, with the slice's usage."""
|
|
_old_zoo(monkeypatch)
|
|
slice_dir = tmp_path / "user.slice"
|
|
leaf = slice_dir / "session.scope"
|
|
leaf.mkdir(parents = True)
|
|
(slice_dir / "memory.max").write_text("34359738368\n")
|
|
(slice_dir / "memory.current").write_text("32212254720\n")
|
|
(leaf / "memory.max").write_text("17179869184\n")
|
|
(leaf / "memory.current").write_text("1073741824\n")
|
|
|
|
monkeypatch.setattr(dnp, "CGROUP_ROOT", str(tmp_path))
|
|
monkeypatch.setattr(dnp, "_proc_self_cgroup", lambda: ["0::/user.slice/session.scope"])
|
|
assert dnp._cgroup_free_bytes() == 2 * 1024**3
|
|
|
|
|
|
def test_the_unaided_reader_handles_cgroup_v1(monkeypatch, dnp, tmp_path):
|
|
_old_zoo(monkeypatch)
|
|
leaf = tmp_path / "memory" / "slurm" / "job_1"
|
|
leaf.mkdir(parents = True)
|
|
(leaf / "memory.limit_in_bytes").write_text("4294967296\n")
|
|
(leaf / "memory.usage_in_bytes").write_text("3221225472\n")
|
|
|
|
monkeypatch.setattr(dnp, "CGROUP_ROOT", str(tmp_path))
|
|
monkeypatch.setattr(
|
|
dnp,
|
|
"_proc_self_cgroup",
|
|
lambda: ["7:memory,blkio:/slurm/job_1"],
|
|
)
|
|
assert dnp._cgroup_free_bytes() == 1024**3
|
|
|
|
|
|
def test_the_unaided_reader_ignores_the_unlimited_sentinels(monkeypatch, dnp, tmp_path):
|
|
_old_zoo(monkeypatch)
|
|
v2 = tmp_path / "scope"
|
|
v2.mkdir()
|
|
(v2 / "memory.max").write_text("max\n")
|
|
(v2 / "memory.current").write_text("1073741824\n")
|
|
v1 = tmp_path / "memory"
|
|
v1.mkdir()
|
|
# v1's "unlimited": a near-2^63 sentinel, not a 8-exabyte ceiling.
|
|
(v1 / "memory.limit_in_bytes").write_text("9223372036854771712\n")
|
|
(v1 / "memory.usage_in_bytes").write_text("1073741824\n")
|
|
|
|
monkeypatch.setattr(dnp, "CGROUP_ROOT", str(tmp_path))
|
|
monkeypatch.setattr(dnp, "_proc_self_cgroup", lambda: ["0::/scope", "7:memory:/"])
|
|
assert dnp._cgroup_free_bytes() is None
|
|
|
|
|
|
def test_the_unaided_reader_is_never_negative(monkeypatch, dnp, tmp_path):
|
|
_old_zoo(monkeypatch)
|
|
leaf = tmp_path / "scope"
|
|
leaf.mkdir()
|
|
(leaf / "memory.max").write_text("1073741824\n")
|
|
(leaf / "memory.current").write_text("2147483648\n")
|
|
monkeypatch.setattr(dnp, "CGROUP_ROOT", str(tmp_path))
|
|
monkeypatch.setattr(dnp, "_proc_self_cgroup", lambda: ["0::/scope"])
|
|
assert dnp._cgroup_free_bytes() == 0
|
|
|
|
|
|
def test_the_unaided_reader_keeps_the_public_limit_as_a_last_resort(monkeypatch, dnp, tmp_path):
|
|
"""No readable cgroup tree here, but an older zoo may still find one its own way.
|
|
|
|
A limit with no usage beside it is still a ceiling, and is a tighter one than
|
|
psutil's host-wide view inside a container.
|
|
"""
|
|
_old_zoo(monkeypatch, memory_limit = 4 * 1024**3)
|
|
monkeypatch.setattr(dnp, "CGROUP_ROOT", str(tmp_path / "absent"))
|
|
monkeypatch.setattr(dnp, "_proc_self_cgroup", lambda: [])
|
|
assert dnp._cgroup_free_bytes() == 4 * 1024**3
|
|
|
|
|
|
def test_the_unaided_reader_never_raises(monkeypatch, dnp):
|
|
"""It runs on every auto-sizing call under an older zoo, on hosts with no cgroup at all."""
|
|
_old_zoo(monkeypatch)
|
|
free = dnp._cgroup_free_bytes_unaided()
|
|
assert free is None or (isinstance(free, int) and free >= 0)
|
|
|
|
|
|
def _no_hf_xet_tuning(monkeypatch):
|
|
"""Neither the private helpers nor the public readers: no unsloth_zoo at all."""
|
|
import builtins
|
|
|
|
real_import = builtins.__import__
|
|
|
|
def _blocked(name, *args, **kwargs):
|
|
if name == "unsloth_zoo.hf_xet_tuning":
|
|
raise ImportError("older unsloth_zoo")
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(builtins, "__import__", _blocked)
|
|
|
|
|
|
def test_no_unsloth_zoo_and_no_cgroup_is_not_a_ceiling(monkeypatch, dnp, tmp_path):
|
|
# The cgroup tree is pointed away from the host's on purpose: the unaided
|
|
# reader needs no unsloth_zoo, so leaving it on the real /sys/fs/cgroup would
|
|
# make the assertion depend on whether the runner is itself in a limited
|
|
# container, which is how this test would fail on CI and pass on a laptop.
|
|
_no_hf_xet_tuning(monkeypatch)
|
|
monkeypatch.setattr(dnp, "CGROUP_ROOT", str(tmp_path / "absent"))
|
|
monkeypatch.setattr(dnp, "_proc_self_cgroup", lambda: [])
|
|
assert dnp._cgroup_free_bytes() is None
|
|
assert dnp._cgroup_cpu_quota() is None
|
|
|
|
|
|
def test_no_unsloth_zoo_still_reads_the_cgroup(monkeypatch, dnp, tmp_path):
|
|
"""The point of the unaided reader: the ceiling survives having no zoo to ask."""
|
|
_no_hf_xet_tuning(monkeypatch)
|
|
leaf = tmp_path / "scope"
|
|
leaf.mkdir()
|
|
(leaf / "memory.max").write_text("8589934592\n")
|
|
(leaf / "memory.current").write_text("6442450944\n")
|
|
monkeypatch.setattr(dnp, "CGROUP_ROOT", str(tmp_path))
|
|
monkeypatch.setattr(dnp, "_proc_self_cgroup", lambda: ["0::/scope"])
|
|
assert dnp._cgroup_free_bytes() == 2 * 1024**3
|
|
# The CPU quota reader has no unaided path, so it stays silent.
|
|
assert dnp._cgroup_cpu_quota() is None
|
|
|
|
|
|
def test_env_forced_serial_is_in_process_on_a_small_split(monkeypatch, dnp):
|
|
"""The documented recovery has to actually recover.
|
|
|
|
UNSLOTH_DATASET_NUM_PROC=0 with the config sentinel 1 arriving as an explicit
|
|
count used to return 1, which bypasses the small-split guard and builds a
|
|
Pool(1) on datasets >= 4.1. Under the threshold the guard is in-process, so
|
|
None is what expresses the request exactly.
|
|
"""
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
_force_stdlib_start_method(monkeypatch, dnp, "fork")
|
|
monkeypatch.setenv(dnp.NUM_PROC_ENV_VAR, "0")
|
|
|
|
small = type("t", (), {"train_dataset": _Split(100)})()
|
|
assert dnp.resolve_responses_only_num_proc(small, 1) is None
|
|
assert dnp.resolve_responses_only_num_proc(small, None) is None
|
|
|
|
# Over the threshold the guard is gone, and 1 is the least it can be given.
|
|
big = type("t", (), {"train_dataset": _Split(dnp.ZOO_MIN_ROWS_FOR_MULTIPROC * 2)})()
|
|
assert dnp.resolve_responses_only_num_proc(big, 1) == 1
|
|
|
|
|
|
def test_a_memory_starved_explicit_count_is_in_process_on_a_small_split(monkeypatch, dnp):
|
|
# Same shape without the env var: the memory clamp resolves to serial, and
|
|
# under the threshold that has an exact encoding.
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
_force_stdlib_start_method(monkeypatch, dnp, "fork")
|
|
monkeypatch.setattr(dnp, "_affordable_workers", lambda: 0)
|
|
|
|
small = type("t", (), {"train_dataset": _Split(100)})()
|
|
assert dnp.resolve_responses_only_num_proc(small, 16) is None
|
|
|
|
|
|
def test_an_explicit_count_the_host_can_afford_is_untouched_by_the_row_guard(monkeypatch, dnp):
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
_force_stdlib_start_method(monkeypatch, dnp, "fork")
|
|
monkeypatch.setattr(dnp, "_affordable_workers", lambda: 1000)
|
|
|
|
small = type("t", (), {"train_dataset": _Split(100)})()
|
|
assert dnp.resolve_responses_only_num_proc(small, 4) == 4
|
|
|
|
|
|
def test_the_fallback_does_not_sit_behind_a_torch_import():
|
|
"""MLX hosts have no torch, and reach this module through chat_templates.
|
|
|
|
unsloth/utils/__init__.py imports .packing, which imports torch, and
|
|
.attention_dispatch, which imports unsloth.models._utils on top. A fallback
|
|
under that package would raise on the torch-free host it exists to serve, so
|
|
it lives at the top level and every reference to it has to stay there.
|
|
"""
|
|
assert (
|
|
MODULE_PATH.parent.name == "unsloth"
|
|
), "the fallback moved back under a package whose __init__ imports torch"
|
|
|
|
utils_init = REPO_ROOT / "unsloth" / "utils" / "__init__.py"
|
|
reached = {
|
|
node.module.split(".")[0]
|
|
if isinstance(node, ast.ImportFrom) and node.level == 0
|
|
else (node.module or "").lstrip(".")
|
|
for node in ast.parse(utils_init.read_text(encoding = "utf-8")).body
|
|
if isinstance(node, ast.ImportFrom)
|
|
}
|
|
assert reached, "unsloth/utils/__init__.py stopped importing anything; re-check the premise"
|
|
|
|
for path in (
|
|
REPO_ROOT / "unsloth" / "chat_templates.py",
|
|
REPO_ROOT / "unsloth" / "models" / "rl.py",
|
|
REPO_ROOT / "unsloth" / "models" / "rl_replacements.py",
|
|
):
|
|
source = path.read_text(encoding = "utf-8")
|
|
assert (
|
|
"unsloth.utils.dataset_num_proc" not in source
|
|
), f"{path.name} reaches it via unsloth.utils"
|
|
assert ".utils.dataset_num_proc" not in source, f"{path.name} reaches it via unsloth.utils"
|
|
|
|
|
|
def test_the_fixture_really_neutralises_the_zoo_readers(dnp):
|
|
"""The fixture's patching must survive a package __init__ that raises.
|
|
|
|
unsloth_zoo/__init__ imports hf_xet_tuning near the top and can raise at the
|
|
end, which drops the package from sys.modules but leaves the submodule
|
|
cached -- and that cache entry is what the policy imports. Patching only on
|
|
a clean `from unsloth_zoo import hf_xet_tuning` leaves the real readers live
|
|
on the runner's own /sys/fs/cgroup, so every sizing assertion silently
|
|
becomes a test of the container's memory limit.
|
|
"""
|
|
module = sys.modules.get("unsloth_zoo.hf_xet_tuning")
|
|
if module is None:
|
|
pytest.skip("unsloth_zoo.hf_xet_tuning is not reachable here")
|
|
assert str(module.CGROUP_ROOT).startswith("/nonexistent"), module.CGROUP_ROOT
|
|
assert module._cgroup_v2_dirs() == []
|
|
assert module._cgroup_v1_dirs("memory") == []
|
|
|
|
|
|
def test_the_unaided_reader_picks_its_own_v1_line_too(monkeypatch, dnp, tmp_path):
|
|
"""The other half of the scan: a v1 line that is not the first line.
|
|
|
|
The hybrid case above puts the memory controller first, so taking line 0
|
|
would still land on it. Here a pids line comes first and the memory
|
|
controller sits at a different path, which is what a Slurm step looks like:
|
|
reading line 0 walks a directory that does not exist and loses the ceiling
|
|
the step was given.
|
|
"""
|
|
_no_hf_xet_tuning(monkeypatch)
|
|
v2_leaf = tmp_path / "user.slice" / "app.scope"
|
|
v2_leaf.mkdir(parents = True)
|
|
(v2_leaf / "memory.max").write_text("8589934592\n")
|
|
(v2_leaf / "memory.current").write_text("6442450944\n")
|
|
# 4GB capped, 3 spent: 1GB free, which is less than the v2 side's 2GB.
|
|
v1_leaf = tmp_path / "memory" / "slurm" / "job_1"
|
|
v1_leaf.mkdir(parents = True)
|
|
(v1_leaf / "memory.limit_in_bytes").write_text("4294967296\n")
|
|
(v1_leaf / "memory.usage_in_bytes").write_text("3221225472\n")
|
|
|
|
monkeypatch.setattr(dnp, "CGROUP_ROOT", str(tmp_path))
|
|
monkeypatch.setattr(
|
|
dnp,
|
|
"_proc_self_cgroup",
|
|
lambda: [
|
|
"11:pids:/user.slice/user-1000.slice/session-3.scope",
|
|
"10:memory:/slurm/job_1",
|
|
"0::/user.slice/app.scope",
|
|
],
|
|
)
|
|
assert dnp._cgroup_free_bytes() == 1024**3
|
|
|
|
|
|
def test_the_hatch_wins_on_a_small_split_too(monkeypatch, dnp):
|
|
"""A split under the threshold is where the resolver would otherwise stand
|
|
down, and standing down there would silently discard the count a user set
|
|
by hand to get workers on a small dataset."""
|
|
_force_start_method(monkeypatch, dnp, "fork")
|
|
_force_stdlib_start_method(monkeypatch, dnp, "fork")
|
|
monkeypatch.setenv(dnp.NUM_PROC_ENV_VAR, "16")
|
|
|
|
small = type("t", (), {"train_dataset": _Split(100)})()
|
|
assert dnp.resolve_responses_only_num_proc(small, None) == 16
|