Feat/quality based upsampling - #461
Open
rrutmann wants to merge 34 commits into
Open
Conversation
Builds a training blend by filtering documents on quality signals and choosing how heavily each dataset is sampled, with a fast preview of the resulting token budget. Selection addresses two kinds of signal with one syntax: metrics a corpus already carries in its records, and external per-document annotations that are joined on by key. Four key kinds cover the corpora we have, including those that store no identifier at all and one whose identifier points into a separate source corpus. The source data is never copied or modified. Selection emits a filtered .idx, and PackedDataGenerator already tokenizes exactly the documents its index lists, so pack_encoded_data consumes the output unchanged and no packing code was touched. Previewing is instant because the per-document table is aggregated into a cube once; a threshold on a bin edge is answered exactly, one inside a bin is reported as interpolated rather than silently guessed. WeightedCombinedDataset applies the ratio at training time, so a dataset can contribute 2.5 or 0.3 epochs without being duplicated on disk. This retires the _epoch_1/_epoch_2 directory convention. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Profiling the pipeline against the real 43 TB blend rather than a fixture found three stages far slower than they needed to be, two of which could not be parallelised at all. Native metrics declared as a plain field path are now read by direct dictionary lookup instead of jq. jq re-serialises the whole document on every call, which on a 21 KB record cost more than twenty times the rest of building a sidecar row: 29 MB/s with jq against 374 MB/s without, 13x end to end. Complex patterns still use jq, and the builder warns when one does, because that pattern then dominates the stage. build_cube groups with Arrow's kernels instead of a Python loop over documents, batching row groups so cardinality saturates before each pass: 246k to 2.39M rows/s. That stage was not shardable, so its 8.5 hours were a hard floor; it is now under an hour. build-sidecar and the new bucket-annotations stage both take --shard_id/--num_shards. Sidecar work is divided per file across every dataset, so one array covers the whole blend rather than leaving a floor of the slowest dataset. Bucketing moved out of join-annotations, which now refuses to run against an incomplete bucketing run rather than silently dropping the annotations a missing task was carrying. One-time setup goes from ~88 h to ~4 h on two nodes. Previewing a selection is unchanged at ~10 s, since it only reads the cubes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Adds the four array scripts and a runbook for running the quality pipeline against /data/annealing, following the numbered-sbatch convention already used for the long-context preprocessing. Two environment findings are recorded because they cost time to discover. /home/richard.rutmann/data is a symlink to /data, so the registry's existing paths already point at the right tree. And nemo_25_11.sif cannot run this branch -- it ships torch 2.9 where modalities needs >=2.10 -- so the runbook uses a plain CPU venv instead, these stages having no use for a GPU. Also adds the packing template the pipeline needs for calibration and for generating per-file packing configs. Its tokenizer is a decision, not a default: every token figure downstream depends on it and a wrong choice fails silently, so the file says so. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Ran the pipeline against /data/annealing/italian/Finewiki end to end. It works, and it corrected something I had documented too narrowly. Duplicate annotation keys are not specific to the Nemotron-CC split as the registry implied: the finewiki split reported 868,586 duplicates in 43.1 M rows. The join already keeps the first occurrence and reports the count, but the docs now say to expect this generally and to check join_report.json rather than assume a one-to-one join. Also records measured stage timings, and that preview pays ~6 s of import overhead before doing any work, so a full-blend preview is nearer 15-20 s than the 10 s the cube arithmetic alone implies. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
quality calibrate read 20,000 lines from every file of a dataset to build a 2,000-document sample, so its cost tracked the file count rather than the sample size. Over the real blend that was 30.3 TB and ~30 hours; dolmino alone, at 40,003 files, was 26 TB of it. Found because a real run was still going after an hour with 9 of 19 datasets done. The sampler now probes at most 32 files, spaced evenly across the dataset, reading only enough lines from each to fill the sample. Full calibration of all 19 datasets takes 4 minutes and reads about 1 GB. Spread is preserved and the seeded trim keeps it reproducible. calibration.yaml is also written after each dataset rather than once at the end, so interrupting the stage no longer discards what it measured. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
run_all_timed.sh runs the stages in order, prints the elapsed time for each and a summary at the end. Array jobs go through sbatch --wait so the timing is the job's real duration rather than the submission's, which also gives the ordering steps 4 and 6 onwards need. Takes step ids to run a subset, stops at the first failure while still printing the summary, and accepts REGISTRY/SELECTION/TOKENIZER_CONFIG overrides so a blend variant needs no edits to the shared configs. The sbatch scripts honour the REGISTRY override too. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The README said ~3 h on 64 tasks. The real run projects ~7 h, and the reason is an arithmetic error rather than anything being slow: the storage floor was derived from the blend's 43 TB read once, but the first build-sidecar reads every file twice -- once for IndexGenerator to write the .idx, then once to build the sidecar. The floor is 2 x 43 TB / 3.8 GB/s, about 6.3 h. Measured mid-run: 51.3 % of the blend's bytes in 3.77 h, which is 3.28 GB/s of actual reads against a measured ceiling of 3.81 GB/s. The stage is storage-bound at 86 % of what the filesystem delivers, so more tasks would not help. Also records that the index pass is not wasted: pack_encoded_data needs those .idx files and later runs reuse them, so a second build-sidecar over the same data is roughly twice as fast. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The first real build-sidecar run took ~15 h against a ~7 h projection. The cause was placement, not throughput: SLURM packed 12 of the 64 tasks onto one node and 8 onto another while four nodes sat idle, because at 2 CPUs per task 16 fit on a node. Each of those tasks got ~95 MB/s where a single stream gets 282 MB/s, so the job spent nine hours in a tail running at a third of its initial speed. Both array scripts now request 8 CPUs per task. A task is single-threaded and needs one; the request is there to cap tasks per node at 4 so an array of 64 spreads across the cluster. The underlying mistake was benchmarking storage on the login node and assuming the 3.8 GB/s aggregate transferred to compute nodes. It does not: per-node bandwidth binds first, which is why raising the array size would not have helped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Companion to f8728e3. The estimate assumed the 3.8 GB/s aggregate measured on the login node transferred to compute nodes; it does not. Per-node bandwidth binds first, so packing tasks onto few nodes caps the job well below the cluster's capacity, and raising the array size does not help. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
bucket-annotations OOM-killed all 64 tasks of a real run at 24 GB each. The writer flushed a bucket once it held 100,000 rows, which bounds nothing: spread 50 M rows over 1024 buckets and each holds ~49 k, so no bucket ever reaches the threshold and the whole input accumulates as Python dicts until close. The cap is now on the total buffered across all buckets, and reaching it flushes everything. On the exact shard that failed -- 50 M rows, 1024 buckets -- peak RSS is 2.36 GB instead of unbounded, at an unchanged 121k rows/s. bucket_annotations now also refuses to write into a directory holding output from a run with a different num_shards, since a sharded run cannot clear the directory and the leftovers would be silently mixed into a bucket. The completeness guard from the earlier sharding work did its job: the join refused the partial buckets rather than producing sidecars with no labels. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The mixed-array-size guard from 6d9e45d read every _meta.*.json with a bare json.loads, while the writer used Path.write_text, which truncates before writing. Sibling tasks of the same array read those files, so 12 of 64 tasks of a real run died on an empty file with JSONDecodeError. Metadata now goes to a per-task .tmp and is renamed onto the final name, so a reader sees either the old file or the complete new one. Both readers skip a file they cannot parse and ignore *.tmp; for read_bucket_metadata that is the safe direction, since an unread file leaves its shard id unseen and the run reports incomplete rather than joining missing annotations. Verified on the real finewiki split: four sequential shards into one directory, 43.1 M rows, metadata reading back complete with a matching total and no temporary file left behind. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The runbook told the reader to export six variables by hand in every shell. A
fresh shell without them turns "$QDIR/slurm/x.sbatch" into "/slurm/x.sbatch",
and "rm -rf $WORK/buckets" into "rm -rf /buckets". Both happened; the second
was harmless only because /buckets does not exist.
env.sh sets the variables, fills in only what is unset, and refuses to
continue if any is empty or points at a missing path, so the failure is a
clear message rather than a confusing sbatch error or an rm against the wrong
directory.
reset_buckets.sh replaces the bare rm in the instructions. It requires WORK
via "${WORK:?}", rejects paths like /buckets outright, reports what it is
about to delete, and leaves the sidecars alone.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
join_annotations looped sidecar parts outside and annotation buckets inside, so the entire bucketed split was read once per part -- and since a part's documents hash across every bucket, those were full re-reads. Measured on the real blend: 454 TB, 451 hours. nemotron-cc alone was 5,319 parts against a 23 GB split. Parts are now processed in batches and each bucket is read once per batch, taking the amplification from the part count to the batch count: 5.9 TB and 5.8 h serial. Buckets are filtered with is_in before anything reaches Python, so memory follows the batch rather than the bucket, and bucket file lists are globbed once instead of per batch. Duplicate keys are counted once among the keys the join wants, rather than afresh per part: finewiki-it reported 868,586 where the real figure is 36,747. Adds 3a_join_annotations.sbatch, one array task per annotated dataset with the dataset resolved from the registry, so wall time is the slowest dataset (~2 h) rather than the sum, plus 3b_build_cubes.sbatch to follow it. Verified on the real finewiki-it sidecar: 1,799,759/1,799,759 annotated in 48 s, against 107 s for a quarter as many documents before. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
A nemotron-cc join reached 5319/5319 parts and was killed by a 12-hour limit during its final write-back, leaving 10 of 5,319 parts unlabelled. build-cube then died on that with KeyError: Field "educational_value" does not exist, after building 9 cubes and before attempting 6 healthy datasets. join-annotations --resume skips parts that already carry the label columns. Finishing the interrupted run took 9 minutes instead of 12 hours, skipping 5,309 parts. Off by default, because resuming after re-bucketing would keep the old labels. build_cube now reads every part's schema instead of trusting the first, and raises CubeError naming the dataset, the missing columns and the --resume command that fixes it. build_cubes builds every healthy dataset, reports the failures together and re-raises, rather than losing six buildable cubes to one bad dataset. Join reports are written per dataset as well: sixteen parallel --only tasks had been overwriting one shared join_report.json. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
evaluate_blend caught the SelectionError a cube raises for an ungrouped field and quietly scanned the per-document sidecar instead. That is exact but reads every document, so a preview advertised as seconds ran past ten minutes on the real blend: nemotron-cc thresholded commercial_bias, which the join attaches but the cube does not group, and dolmino thresholded dclm_plus2, which exists only under one subdirectory and was null in all of 400 sampled parts. The fallback is now opt-in via --allow_fallback. Without it, every unanswerable predicate is reported together with the field, the dimensions the cube does carry, and the document count a scan would read. The real selection fails in 16 s naming both problems. build-cube --label_dimension chooses the grouped columns, so a field a selection needs can be added; it replaces the default seven rather than extending them. The example selection no longer thresholds fields the default cubes lack, and says at each site why and how to re-enable. First full preview of the real blend: 13.7 s over 19 cubes, 3.04 T effective tokens against a 400 B target. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
A sidecar row locates its document by (file_id, byte_offset, byte_len), and file_id was a position in a file list re-derived from the filesystem at every stage. That makes every recorded offset depend on the source tree never changing, with nothing recording what it looked like. An ongoing transfer then re-sharded four corpora after their sidecars were built. Nemotron-CC went from 606 MB files to 137 MB files while keeping nearly the same file count, so the count comparison that existed passed and the offsets pointed past end of file. Eleven of nineteen datasets were unusable and only one -- whose file count fell to zero -- failed loudly. build-sidecar now records its file list in sidecar/<dataset>/_files.json, written atomically since a sharded build has every task describing the same list. apply resolves ids through that manifest instead of re-globbing, so an added file cannot renumber anything, and refuses to run if a recorded file changed size or vanished. Paths are relative, so moving or snapshotting a tree stays valid. New quality verify-sidecar seeks to recorded offsets and compares the document found against the recorded text length, catching a file rewritten at the same size. Rows at offset zero are skipped: the first document of any JSONL file parses, so they succeed against a completely different file, which is how the broken sidecars looked healthy. build_sidecars no longer defaults its index root to the source tree. SidecarBuilder writes a .idx beside each JSONL when given no index root, which would modify a shared read-only corpus. Adds a snapshot-based smoke test covering all four join-key kinds plus the native-metrics-only path, so the full pipeline can be exercised in minutes rather than 15 hours. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
--resume counted only the parts it re-joined, so a run that skipped everything wrote a report saying 0 documents and 0.0 coverage. On the smoke run that overwrote three datasets' genuine coverage with zeros, which reads as a failed join rather than a skipped one. Coverage describes the sidecar, not the run, so skipped parts now contribute their existing labels to the totals and n_parts_resumed records how many were not redone. Updates the resume test, whose assertion encoded the old behaviour. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The smoke run compared preview estimates against a real packing run for the first time. Document counts matched exactly; finewiki-de tokens were 10.7% out. Two independent defects. The sample was a prefix: _sample_documents took the first N documents of each probe file. Deterministic, so it gave identical ratios across every seed, which reads as stability rather than as a sample that never moves. On FineWiki the first 2,000 documents give 3.531 bytes/token where the whole file gives 4.214. I introduced this when bounding an earlier 30 TB read. And one global ratio cannot describe a corpus: FineWiki runs from 3.571 for documents under a kilobyte to 34.648 for the nine above 256 KB, which hold 5.7% of all bytes. The global estimator's error swung from -19.4% to +4.2% depending on whether the sample caught them. Now: offsets spread across each file, taking the document containing each offset so selection is length-proportional (top stratum went from 2 samples per 2,000 to 44); ratios measured per log-spaced size stratum and applied from the byte length the sidecar already records; inverse-probability weights for the corpus-wide fallback, since a plain sum ratio under this sampling came out 62% low; and each document measured on a <=64 KB slice from a random position, because tokenizing the multi-megabyte documents in full took calibration past 10 minutes. Worst error over 5 seeds against ground truth: 0.8%, down from 19.4%. Pre-existing calibrations have no strata and fall back to the global ratio, so they load but should be re-measured. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Re-ran the smoke pipeline after the estimator fix. Total estimated vs packed tokens now -0.02%, worst dataset -1.60%; finewiki-de was -10.72% before. Document counts match exactly for all five datasets, which is the stronger check since the filtered index names exactly the selected documents. The blend loads through WeightedCombinedDataset with fractional repeat factors, exercising the partial-pass permutation on real packed files for the first time, and nothing was written under the source root. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The join built a dict of key to (part, row) over every document, a per-row dict of labels, and a list comprehension per label column -- a handful of Python objects per document, at 1.7 bn documents for Nemotron-CC. Each part is now resolved with one index_in and one take per label column, and the key column stays Arrow in the outer loop, where materialising it had been 1.7 bn Python strings before any joining began. Profiling what remained showed the next cost was not per-row work: 22 s of reads and 14 s of a thousand separate is_in calls out of 47 s, because a 554 MB split is spread over 1024 files of ~540 KB. The routing those buckets exist for also decided nothing -- a batch's keys hash across every bucket, so all 1024 files were read anyway, after a blake2b call per key to choose them. Routing is replaced by one pyarrow.dataset scan with the key filter pushed in, which keeps memory bounded by matching rows rather than split size. finewiki-en, 6.6 M documents: 88.7 s -> 42.1 s, 13.41 us -> 6.36 us per document. Equivalence verified twice: 79,375,860 label values identical on the benchmark and 1,441,584 across all four key kinds on the smoke blend, with duplicate counts matching exactly. This does not address the larger cost. The split is scanned once per batch, and nemotron-cc has 85 batches over a 22 GB split -- about 1.9 TB, which is where its twelve hours mostly went. The benchmark has one batch, so it measures the overhead fixed here and none of the re-scanning. See CHANGELOG for the numbers and the route to fixing it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Matches the visual language of the existing training/pruning canvas: rounded boxes, orange stage labels, owner badges, grouped containers, '=' underlined section header. Laid out as two converging lanes -- documents along the top, annotations below, meeting at the join. A single row needs arrows that skip past boxes, and those cross other boxes wherever they are routed; the first attempt had three such crossings. Includes the generator, so the layout can be extended as the pipeline changes, and a check that no arrow segment passes through a box it does not terminate on. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> for also decided nothing -- a batch's keys hash across every bucket, so all 1024 files were read anyway, after a blake2b call per key to choose them. Routing is replaced by one pyarrow.dataset scan with the key filter pushed in, which keeps memory bounded by matching rows rather than split size. finewiki-en, 6.6 M documents: 88.7 s -> 42.1 s, 13.41 us -> 6.36 us per document. Equivalence verified twice: 79,375,860 label values identical on the benchmark and 1,441,584 across all four key kinds on the smoke blend, with duplicate counts matching exactly. This does not address the larger cost. The split is scanned once per batch, and nemotron-cc has 85 batches over a 22 GB split -- about 1.9 TB, which is where its twelve hours mostly went. The benchmark has one batch, so it measures the overhead fixed here and none of the re-scanning. See CHANGELOG for the numbers and the route to fixing it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Measured on nemotron-cc itself, after two proxy datasets gave per-document costs 6x apart and a projection that exceeded the time actually observed with slower code. 60 parts of the real nemotron-cc sidecar, 18,748,230 documents with real UUID keys against the real 747 M-row split in 16,384 files: 772.6 s -> 505.0 s, 1.53x, with all 224,978,760 label values identical. That also settles whether replacing bucket routing with a filtered dataset scan regressed this shape -- it does not, this is where the gain is. But peak RSS doubled, 24.9 -> 51.9 GiB, against a 64 G request: 19% headroom on a stage already OOM-killed once. Raised to 160 G. Decomposition: one scan of the split costs 252.5 s, non-scan work 13.47 us/document, so the full dataset is roughly 6 h of scanning plus 6 h of per-document work at today's 85 batches. Fewer batches is worth about 1.8x, not the 10x my original estimate implied, and 6 h is the floor. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
A ratio gives every surviving document the same repeat factor, so a dataset filtered to "content quality at least adequate" repeats its barely-adequate documents as often as its excellent ones. `upsampling` replaces the scalar with a curve whose factor rises with quality. Method and functional form from Dolma 3 / Olmo 3 (arXiv:2512.13961), which measured it against flat quality filtering on 1B models and found it better at every matched repetition factor. Quality sits on a [0,1] axis where a bucket's width is its share of the dataset's tokens; the factor is C*(x-a)**p above the discard threshold, integral pinned to the token target, no bucket above max_factor, monotone. Reproducing their published example returns exactly their figure: bottom eight vigintiles dropped, top at 7.00x. Two deliberate departures. Their family carries an extra exp(lam*(x-a)); fixing lam=0 makes the solution unique rather than a feasible region and every integral analytic. And the exponent is capped at 8, because steepest-admissible is right only when the cap binds: when the target is a small fraction of the pool -- any blend drawing far fewer tokens than it holds -- nothing violates the cap and the budget collapses onto the top bucket, which is the hard top-k filtering the curve exists to beat. The first attempt did exactly that, exponent 512 with one surviving bucket. apply writes one index tree per bucket, since the packer emits one file per source file and differing factors only mean anything in separate indexes. Each becomes a manifest row named <dataset>__<level> carrying source_dataset so the registry lookup still resolves, and its own repeat factor in WeightedCombinedDataset. The curve is re-solved from the exact counts apply observes rather than the cube's interpolated ones. Ordinal fields only for now; a numeric axis needs the cube's quantile edges in materialisation. Not ablated -- the cluster was busy, so this is verified against the paper's example and by unit tests, not a training run. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The transfer renamed HPLTv4 to hpltv4 in all four languages, so the registry resolved nothing for those datasets and apply failed on them. Not my change -- it came from the working tree while I was on the upsampling curves, and is separated out so the feature commit stays about the feature.
Two ideas from a colleague's TokenFabric design. Most of that architecture applies policy at training time and does not transfer to an offline pipeline, but two of its principles do. Repetition is per pass, and nothing was checking it. target_tokens was documented as "only used to report the gap", and that is all it did. A ratio is relative to one pass over the blend, so a run consuming more tokens than the blend yields wraps and multiplies every factor. The curves added in the previous change can assign 7x; two passes make that 14x, silently. target_tokens now means what the run consumes, and max_total_exposure caps how often anything may be seen -- a curved dataset uses its own max_factor instead. preview reports it, apply refuses, --allow_overexposure overrides. Caps apply to the requested factor even with no target declared, since a ratio of 9 against a cap of 2 is a violation at any number of passes. On the smoke blend, a declared 300 M run against its 151 M yield: climbmix-en / high asked for 2.53x and would be seen 5.01x against its own 4x cap. preview --explain attributes retention to individual predicates: what each matches alone, its marginal effect with the others left in place, and a pairwise overlap matrix. Milliseconds, since the per-cell weights are the whole computation. It found two redundant predicates on the smoke blend immediately -- content_integrity on finepdfs-es (21 k marginal of 65 M) and educational_value at_least basic on klettermix-de (matches 100.0%). Two bugs found while building it: the overlap diagonal multiplied a predicate's factor by itself, which squares a fractional survival rate and undercounts (8.19 M read where 9.75 M matched); and _cube_weights tested the cumulative exactness flag, so once any predicate was interpolated every later one was reported as interpolated too. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
target_tokens was already set but did nothing; it now means what the run consumes. max_total_exposure is set to 4.0, above the largest ratio in the selection (3.0), so it fires on an accident rather than on the intended design -- a blend trimmed too far, or a ratio raised without re-checking the yield. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Pointer resolution destroyed the pointers it could not resolve. Nemotron-ClimbMix was moved to /data/annealing_unused, so resolution returned nothing, wrote 225 M nulls over the pointers, logged it as INFO and carried on -- surfacing two hours later as 0% coverage with the pointers gone. It now refuses to write a part that resolved nothing. The join scan was not bounded by matches. The comment claimed memory stays bounded by matching rows; to_table buffers the whole scan. On HPLT, 65,536 bucket files and 3.76 bn rows, that hit 167 GB and was OOM-killed against a 160 G request even though only ~10 M rows can match. Now streamed via Scanner.to_batches with bounded readahead. The join report merge raced: non-atomic per-task writes let a sibling read a half-written file, so finewiki-es died on JSONDecodeError after its own join had already succeeded. Same fix as the bucket metadata, os.replace plus a tolerant reader. And the delivered Nemotron-CC changed shape: warc_record_id renamed to id, and re-chunked from 1.70 bn documents of ~1 KB to 504 M of ~12 KB in the same 1.8 TB. Registry updated; 400/400 ids match in both subdirectories. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The streamed scan fixed the wrong layer. HPLT still exceeded 160 GiB, and the measurements show peak memory tracks bucket-file count rather than data volume: nemotron-cc at 16,384 files and 504 M documents peaked at 81.6 GB, while HPLT at 65,536 files and only 10.7 M documents was OOM-killed. Four times the files, one forty-seventh of the data, and it is the one that dies -- a dataset holds a fragment per file with that file's parquet metadata, and batching record batches never touched that. The split is now scanned in groups of SCAN_FILE_GROUP files. All three HPLT splits completed in 42-45 min, and hplt-it, which had succeeded before, returns identical coverage. All 16 annotated datasets now join at 100%. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
pack_encoded_data rebuilds its components, tokenizer included, from the config on every call -- 24.7 s measured on a compute node, against ~3 s of real work for a Dolmino file. The blend renders 54,738 configs, 40,003 of them Dolmino, so the CLI-per-file shape spends ~375 core-hours on startup for ~48 core-hours of tokenising. pack_many.py loads the tokenizer once per process (1.2 s, paid once) and constructs a PackedDataGenerator per config, which is what pack_encoded_data does internally. Slices are strided, since the config list is grouped by dataset and contiguous slices would give one task all of Dolmino; a failing config logs and continues; --skip_existing makes reruns resumable. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The smoke checker assumes a handful of packed files and loads every document index at once; at 54,738 files and 6.0 TB that does not hold. scan_pbins.py reads every packed file's header and reports any with an empty or unreadable data section. On the first full run it found exactly one bad file in 54,738: finewiki-it/000_00000.pbin, 151 MB on disk reporting data_len=0, a leftover that pack_many.py --skip_existing had skipped because the file existed. Existence taken for health -- the same proxy-instead-of-check mistake this pipeline has now made in several places. --skip_existing now reads the header. Had verification counted files instead of reading them, that dataset would have trained 567 M tokens short. verify_blend.py compares packed tokens against manifest estimates, packed documents against documents selected, and file counts, reading one index at a time. Real blend after repacking: 18 datasets, all document counts exact, 1,637,755,654,948 packed tokens against 1,646,919,435,199 estimated, -0.56% total, worst dataset -4.39%, nothing written into the source tree. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Committed them straight from the scratch copies without reading the linter output first. Multi-statement lines, an unused import and an over-long line. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
4_pack.sbatch now drives pack_many.py rather than one pack_encoded_data call per config, which is what made packing 1 h 40 m instead of an estimated ten hours of tokenizer startup. 5_verify.sbatch runs the three checks that each caught something real: header scan, token and document counts against the manifest, and loading the blend the way training will. scan_pbins.py, verify_blend.py and load_blend.py take --work_dir instead of hardcoded paths, so they work for any blend. The README runbook now carries measured timings from the full run over /data/annealing rather than estimates, and says which stage each check protects. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Each cost hours on the 20 TB blend and none had a test. They share a shape: a stage fails, reports it somewhere quiet, and destroys or hides the evidence, so the symptom shows up much later and further downstream. - pointer resolution wrote nulls over the pointers it could not resolve, so a retry needed a full sidecar rebuild. The test asserts the pointers survive a failed resolution. - the annotation scan held a fragment per bucket file, so a 65,536-file split with 10.7 M documents exceeded 160 GiB where a 16,384-file split with 504 M did not. The test varies SCAN_FILE_GROUP from 1 to 1024 and asserts the join result is identical. - skip-if-exists treated a truncated .pbin as finished; one file in 54,738 was 151 MB on disk reporting data_len=0. The test covers absent, truncated, garbage and healthy. Self-contained fixtures rather than moving the shared ones out of test_quality_pipeline.py, to avoid disturbing the 55 tests that use them. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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
1 participant
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.
What does this PR do?
This PR adds quality-based document selection and up/downsampling to the data preprocessing pipeline, letting an ablation choose which documents enter a training blend by thresholding quality annotations and native corpus metrics, and how heavily each surviving subset is repeated. It was run end to end on the 20.21 TB annealing blend: 19 datasets, 16 annotated at 100 % join coverage, 1.64 T tokens packed into 6.0 TB, with packed token counts landing within −0.56 % of the pipeline's own estimates and document counts matching exactly.
The design keeps selection virtual for as long as possible. A per-document sidecar records where each document is and what it is worth; a contingency cube aggregates that so a threshold can be costed in seconds; and only apply materialises anything, writing filtered .idx files that name the surviving byte ranges. The source corpora are never written to. Trying a different threshold costs 14 seconds rather than a repack.
General Changes
Breaking Changes
None. Everything is additive: a new quality command group, a new WeightedCombinedDataset component, and new config files. No existing entry point, config schema or dataset class changes behaviour.
Checklist before submitting final PR
python tests/tests.py)CHANGELOG_DEV.md)