Add some materialization tooling for mixtures - #313
Draft
undfined wants to merge 46 commits into
Draft
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Tooling to materialize a token mixture: resolves a mix spec against
.npysourcesin S3, plans an exact set of resharding jobs, runs them on pmr workers, verifies
the output. Includes config for the Dolma 3.5 14T mixture.
Pipeline
plan.pyvalidate.pypreflight.pymaterialize.pyverify.pyoutput_report.pygenerates the HTML reports.setup_worker_storage.shpreparesworker NVMe, single device or RAID0 depending on instance type.
Build state goes in one directory named by a hash of the mix and catalog. Tools
abort if either changed after the plan was built.
scripts/resharding/dispatch.pybuilds the pmr lifecycle commands and ismixture-agnostic.
New:
python/dolma/tokenizer/document_selection.pySelects a subset of whole documents from a token memmap to reach a target token
count. Key is blake2b over
(seed, start, end, document_id, source, source_index),so identical source metadata yields same selections. Two streaming passes;
metadata is not held in memory.
Changed:
python/dolma/tokenizer/reshard.pyper-row
repeat_countand optional partial selection, as an alternative to prefixglobbing. Sizes and ETags are verified before and after download; a mismatch
aborts the run.
merge_groupcopies per-document ranges via a selection index in addition towhole files, and emits progress callbacks.
reshard()errors if the destination contains objects. Token and metadata writesare no-clobber.
Breaking
ReshardingConfigis keyword-only. Positional construction raises.weighted_bucket_sampleremoved.group_paths_by_max_num_filesuses largest-firstbin packing, output shard composition differs for the same input.
New:
configs/dolma3p5-resharding/14t/Mix, source path catalog, and
settings.yaml(token target, worker instance grid,output shard sizing, residual bounds).
inputs/mix-v2.yamlis theplan.pydefault;
inputs/dolma3p5-14t-optimal-mix.yamlis the superseded v1.inputs/README.mddocuments both. Editing either file changes the build id andinvalidates existing plans.
READMEs:
configs/dolma3p5-resharding/14t/README.mdfor the run procedure,scripts/dolma3p5_resharding/README.mdfor code layout.Tests
92 tests in
tests/python/test_dolma3p5_resharding.py(new) andtest_resharding.py. Cover planning, allocation, worker lifecycle, manifesthandling, selection determinism, and overwrite refusal. No network or AWS
credentials required.
CI fixes (pre-existing failures on
main)pyproject.toml: pytestpythonpath = ["."], so tests can import fromscripts/regardless of invocation directory.
Cargo.lock:cc1.0.92 -> 1.4.4,cmake0.1.50 -> 0.1.58. Older versions failto detect Visual Studio 2026 on the
windows-2025-vs2026runner..github/workflows/CI.yml: removed the Linux container override, usingmaturin-action's per-target
manylinux_2_28defaults instead. The rust-crossx86_64 image ships an empty toolchain directory. Package install is now guarded on
apt-getsince the manylinux image is RPM-based.python/dolma/taggers/quality.py: black formatting only.Note: of the ~36k lines, 20,729 are data — the two mix weight tables (8,425 and 8,386 lines) and the source path catalog (3,918)