Repository navigation
perf(data)!: bin once and share: cached dual-layout index, row views, bin-routed margins, binned matrices - #168
Closed
brndnmtthws wants to merge 7 commits into
Closed
brndnmtthws wants to merge 7 commits into
brndnmtthws wants to merge 7 commits into
Conversation
GHistIndex keeps its two layouts, both written by one parallel pass over 1,024-row blocks instead of binning rows and then transposing them: - rows: each row's present entries as global bins (u16 up to 65,536 bins, else u32), CSR-shaped, which row-wise histogram gathers read; - columns (dense indexes and sparse ones at least half full): each feature's bins local to the feature, feature-major, in u8 up to 256 bins per feature (the default max_bin), else u16/u32, plus one bit per entry marking missing values, which partitions, vector-leaf sweeps, budget/quantized column paths and the GPU uploads read. The columns were global u16 with a u16::MAX sentinel; local bins with a separate bitmap fit 256-bin features in a byte. HistCuts streams dense rows through 65,536-row transposed blocks (every sketch still takes its feature's values in row order) instead of transposing the whole matrix, and CSR input goes through its CSC view a feature at a time as before. Every histogram keeps its per-bin addition order (sum_order is keyed on the dense flag as before), so models are bit for bit unchanged.
A DMatrix and its metadata copies (with_labels, with_weights, Python's copy + set_info) share their feature values; they now also share the histogram indexes binned from them. DMatrix::hist_index(max_bin) keys a BinCache by max_bin, sample weights, feature types and missing value, and hist training, EBM pair ranking and budget training read their index from it, so repeated and concurrent trainings of one matrix bin once. The index is freed with the values; new values (dense_values_mut, map_values, select_rows) start a new cache. The first request for a key builds it and concurrent requests wait. A parallel build runs on a pool of its own and publishes from there: a builder waiting on its own rayon joins may steal a waiter (two trainings of one matrix started from one pool), which would otherwise wait on its own stack forever. Serial builds (small matrices, one thread) never enter rayon and run on the caller; a builder that cannot start threads builds unshared and its waiters build their own. A cache hit is the index a miss builds, so training is bit for bit unchanged.
`DMatrix::select_rows` now returns a view: the selected rows' metadata is gathered, the feature values stay shared with the source and are read through a composing row map. Raw-value readers go through the row-mapped `DenseRows`/`CsrRows` accessors, and writing a view's values (`dense_values_mut`) copies only its rows. Cross-validation folds, target-statistics folds and every other internal selection are views. `select_rows_with_cuts(rows, ViewCuts)` picks the view's histogram cuts: `ViewCuts::Own` (the default) cuts the view's rows, so training on a view is training on a copy of its rows, bit for bit; `ViewCuts::Source` gathers the view's rows of the source's cached index (`GHistIndex::select_rows`), so a source and its views share one binning pass and split on the same thresholds. Python: `DMatrix.slice(rows, cuts="own" | "source")`. The routing guard of the `hist` builder's row partitions is now `DMatrix::sketched_rows`: a source-cut view of rows its source weighs zero holds values its cuts never saw, so those rows' training margins route by value (and linear leaves and device rounds keep off the partitions), which keeps training equal to continuing it. BREAKING CHANGE: `select_rows` keeps the source's missing value instead of converting the selection to NaN-missing.
A sampled round's margin update and exact leaves' row routing walked every training row through the new tree by raw value, reading scattered row-major `f32`s. When the index's cuts were sketched from every row (no zero sample weight, `SketchedRows::All`), a row's bins route it through any tree split on those cuts as its values do, so both now walk the index's narrow feature columns instead (`tree::binned::BinnedTree`, `margins::RowRouting`); every row still receives one addition of the same leaf value, so models are unchanged bit for bit. Trees whose thresholds fall inside a bin (`HistCuts::bins_left_of`), linear leaves, refreshed trees, sparse indexes without columns, and eval sets keep routing by value. `HistCuts::categories_left` replaces the builder's private `category_left`, now shared by partitions, device growth, and the walker. Prediction no longer loads every row block when summing an empty tree range, which every training's initial margins did.
`data::BinnedMatrix` (Python `hessboost.BinnedMatrix`) keeps a matrix's `hist` index and metadata and drops the values, as XGBoost's `QuantileDMatrix` does: `BinnedMatrix::new(&dmatrix, max_bin)` takes the index the matrix's own training builds (shared with it, so binning after a training costs nothing), and training on it is training on the matrix, bit for bit. On the 4M x 173 panel the index takes 1.4 GiB against 2.6 GiB of values, and 20 rounds of quantile training peak at 1.56 GiB once the `DMatrix` is dropped. `Trainer::new` and `train` take `impl Into<TrainingData>`, a `DMatrix` or a `BinnedMatrix`. Training refuses what would read the values it dropped: another `max_bin`, continuation, `device = cuda` (whose recovery routes rows by value), and every option outside an allow-list of what `hist` growth on bins serves (`validate_binned_request` over `refuse_changes_from`, so a new option is refused until it is vetted). Binning refuses zero-weight rows, whose values the cuts never saw. The weights and feature types that shaped the cuts are fixed; labels, groups, margins, bounds, and feature weights can be set. The values-free storage (`Storage::Bins`) panics if a value is ever read, so a missed refusal fails loudly. Training rows route through finished trees by their bins on every index layout now: `BinnedTree` also walks a sparse index's row bins. Python: `train` is generic over the matrix, and a custom objective receives it (`Objective[M]`); prediction, eval sets, and `cv` refuse a `BinnedMatrix` with a `TypeError`. `DMatrix` and `BinnedMatrix` share their metadata readers (`_Metadata`).
Review fixes for the bin sharing: - A binned matrix's rows are all sketched (`SketchedRows::All`): binning proved it, and its weights are fixed. Re-deriving it from the weights failed for a source-cut view that weighs rows its unweighted source sketched anyway, sending its training to read values it no longer has. - A forked child bypasses the bin caches it inherits: the fork may have caught one mid-build (its builder thread did not fork) or with a lock held, either of which would hang the child's first training. - Budget mode refuses a source-cut view that weighs rows its source never sketched: its margins and loss decrements follow the builder's partitions, which route those rows by bin, not value. - A source-cut view keeps its source's feature types, which its cuts were computed for (`with_feature_types` refuses others). - SGLB leaf re-estimation reads the builder's partitions under the same guard as the margins, routing rows by value for such views (scalar and vector leaves).
… perf/bin-once-share
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: de49264310
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Comment on lines
+138
to
+140
| let mut columns: Vec<ColumnCuts> = (0..n_features) | ||
| .map(|f| ColumnCuts::new(is_categorical[f], present[f], max_bin, ingest)) | ||
| .collect(); |
There was a problem hiding this comment.
For wide dense matrices, constructing all ColumnCuts up front retains every feature's quantile queue, summaries, and sort scratch until all row blocks have been processed. With fewer than CUT_BLOCK_ROWS rows, the transposed block already occupies the same rows × features memory as the former full transpose, so datasets such as thousands of rows by tens of thousands of features can now consume hundreds of megabytes or more in additional per-feature sketch state and fail from memory pressure. Process bounded feature batches, or otherwise avoid keeping one populated sketch per column concurrently.
Useful? React with 👍 / 👎.
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.
What
Bin each matrix's feature values once and share the result: across trainings of a matrix and its metadata copies, across cross-validation folds and other row selections, and across concurrent trainings. Plus a matrix type that trains from the bins alone (XGBoost's
QuantileDMatrix).Stacked on #164.
perf(data)): one parallel pass over 1,024-row blocks writes CSR global row bins (u16up to 65,536 total bins) and, for dense indexes and sparse ones at least half full, feature-major local columns (u8up to 256 bins per feature, elseu16/u32) with a missing-entry bitmap. Cuts stream dense rows in 65,536-row transposed blocks. Every builder keeps its exact per-bin addition order, andsum_orderstill keys on the oldis_denserule, so no sum changes.perf(data)):DMatrix::hist_index(max_bin)caches the index with the values, keyed bymax_bin, weights, feature types and missing value. The first caller builds it and concurrent callers wait (a parallel build runs on its own pool, so waiting rayon workers cannot deadlock). Clones andwith_*copies share it, and it is freed with the values.feat(data)!):select_rowsreturns a view (a row map over the shared storage, with only the metadata gathered).ViewCuts::Own(the default) cuts the view's own rows, bit-identical to a copy of them.ViewCuts::Sourceinstead gathers the view's rows of the source's index, so a source and its views share one binning pass. Python:DMatrix.slice(rows, cuts="own" | "source"). CV folds are views.perf(training)): sampled rounds and exact leaves route every training row through the finished tree by its bins (BinnedTree) instead of reading scattered rawf32rows. This happens only when the cuts were sketched from every row (SketchedRows::All); otherwise rows are routed by value, as before. Prediction also skips loading row blocks when the tree range is empty, which every training's initial margins did.feat(data)):data::BinnedMatrix/hessboost.BinnedMatrixkeeps the index and metadata and drops the values.Trainer::newandtraintakeimpl Into<TrainingData>(aDMatrixor aBinnedMatrix), and training on the bins is training on the matrix, bit for bit. Everything that would read values is refused through an allow-list overrefuse_changes_from, so options added later are refused until vetted: anothermax_bin, continuation,device = cuda,tree_methodapprox/exact, non-gbtree boosters, linear leaves, refresh, SGLB, and reuse penalties. Binning refuses rows with zero weight, whose values the cuts never saw.feat(python)):DMatrix(x, copy=False)reads a read-only, C-contiguous, alignedfloat32array in place, so a memory-mapped.npy(np.load(..., mmap_mode="r")) stays file-backed; the matrix and every copy or slice of it keep the array alive (aBinnedMatrixkeeps only the bins). Refused: writeable arrays, non-contiguous ones, arrays whose base chain (followed through memory views and the buffer protocol) can still write the buffer, and buffers another Rust extension holds a mutable borrow of. The core API isDMatrix::from_dense_shared(Arc<dyn AsRef<[f32]> + Send + Sync>, ..); the extension's oneunsafeblock (python/src/borrowed.rs,unsafe_codelowered fromforbidtodenywith that block allowed) turns the kept-alive numpy buffer into the slice. The contract, documented oncopy: the array, and every array sharing its memory, stays unchanged while a matrix borrows it — numpy offers no way to enforce that (anyWRITEABLEflag can be set again), so it is the caller's promise, as in every zero-copy numpy bridge.Numbers
4M x 173 panel (41.6% NaN), Neoverse-V3, cores pinned; before = #164's head (
53862434, main merged in), after = this branch. Another workload shared the cores during these runs: memory is unaffected, and the per-round rows are the fastest of 12 runs interleaved with the baseline.copy+set_infoBinnedMatrixof the panel, then drop theDMatrix, then 20 roundscopy=False)When copying, the 4- and 8-head peaks are one value copy (2.64 GiB) + one index (1.42 GiB) + about 0.12 GiB per head; borrowing drops the copy, since the memory-mapped pages are file-backed and not counted. A profile of 20 sampled rounds puts the bin-routed margin update at 4.5% of samples, against about 10.5% for the raw walk it replaces (
DMatrix::get+predict_row+goes_left). The skipped row loads for an empty tree range were another 7.4%.Bit identity
Model and margin hashes for 7 configurations (vector quantiles, weighted and dense; 2-target squared error by node; vector lossguide; scalar quantile; scalar squared), each at 1, 4, and 16 threads, are identical to #164 after every commit.
tests/binned_matrix.rstrains every option binned training accepts on both theDMatrixand itsBinnedMatrixand compares model bytes and eval histories: sampled rows and columns, lossguide, forests, GOSS, quantized gradients with leaf renewal, extra trees, path smoothing, monotone and interaction constraints, exact leaves, vector leaves, multiclass, ranking, and a sparse index without columns, with symmetric trees.Design notes
DMatrix::sketched_rows(All/Weighted/Partial) replaces the weights-only check. It sees through source-cut views: a view that gives weight to rows its source weighs zero getsPartial, so its training margins route by value. That keeps "train k + j rounds" equal to "train k, then continue for j" (tests/row_views.rs).sweep_pays) is carried onto the new column layout in the merge of feat(training): opt-in exact quantile leaves (leaf_estimation_method = exact) #164's head (de492643).Review
An independent review found five defects, all around source-cut provenance;
0e74a6a8fixes them and adds a regression test for each:with_feature_typeson a source-cut view contradicted its cuts; it is now refused.The re-review of the fixes found nothing further.
Breaking changes
DMatrix::select_rowsreturns a view that keeps the source's missing value; it no longer converts the selection to NaN-missing.Trainer::newandtraintakeimpl Into<TrainingData<'_>>.&DMatrixcallers compile unchanged; callers passing&&DMatrixneed one deref.Objectiveis generic over the training matrix type (Objective[M]).The 0.4.0 bump in #165 covers these.