Skip to content

Add some materialization tooling for mixtures - #313

Draft
undfined wants to merge 46 commits into
mainfrom
undfined/3p5-materialize
Draft

Add some materialization tooling for mixtures#313
undfined wants to merge 46 commits into
mainfrom
undfined/3p5-materialize

Conversation

@undfined

@undfined undfined commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

Summary

Tooling to materialize a token mixture: resolves a mix spec against .npy sources
in 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

Script Purpose
plan.py Resolve the mix, inventory S3, emit per-unit configs, manifests, and launchers
validate.py Check the plan is internally consistent before launching
preflight.py Re-check source drift and destinations immediately before job launch
materialize.py Provision workers, dispatch units, stream progress, verify output
verify.py Re-verify materialized output and write a report

output_report.py generates the HTML reports. setup_worker_storage.sh prepares
worker 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.py builds the pmr lifecycle commands and is
mixture-agnostic.

New: python/dolma/tokenizer/document_selection.py

Selects 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.py

  • Manifest sources: a config can name a CSV manifest of exact npy/csv.gz pairs with
    per-row repeat_count and optional partial selection, as an alternative to prefix
    globbing. Sizes and ETags are verified before and after download; a mismatch
    aborts the run.
  • merge_group copies per-document ranges via a selection index in addition to
    whole files, and emits progress callbacks.
  • reshard() errors if the destination contains objects. Token and metadata writes
    are no-clobber.

Breaking

  • ReshardingConfig is keyword-only. Positional construction raises.
  • weighted_bucket_sample removed. group_paths_by_max_num_files uses largest-first
    bin packing, output shard composition differs for the same input.
  • Populated destinations error instead of being reused.

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.yaml is the plan.py
default; inputs/dolma3p5-14t-optimal-mix.yaml is the superseded v1.
inputs/README.md documents both. Editing either file changes the build id and
invalidates existing plans.

READMEs: configs/dolma3p5-resharding/14t/README.md for the run procedure,
scripts/dolma3p5_resharding/README.md for code layout.

Tests

92 tests in tests/python/test_dolma3p5_resharding.py (new) and
test_resharding.py. Cover planning, allocation, worker lifecycle, manifest
handling, selection determinism, and overwrite refusal. No network or AWS
credentials required.

CI fixes (pre-existing failures on main)

  • pyproject.toml: pytest pythonpath = ["."], so tests can import from scripts/
    regardless of invocation directory.
  • Cargo.lock: cc 1.0.92 -> 1.4.4, cmake 0.1.50 -> 0.1.58. Older versions fail
    to detect Visual Studio 2026 on the windows-2025-vs2026 runner.
  • .github/workflows/CI.yml: removed the Linux container override, using
    maturin-action's per-target manylinux_2_28 defaults instead. The rust-cross
    x86_64 image ships an empty toolchain directory. Package install is now guarded on
    apt-get since 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)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant