Skip to content

Reduce peak GPU memory during histogram tree growth - #31

Merged
dev-rinchin merged 1 commit into
sb-ai-lab:masterfrom
SergeevVladislav:fix/reduce-histogram-memory
Oct 27, 2026
Merged

dev-rinchin merged 1 commit into
sb-ai-lab:masterfrom
SergeevVladislav:fix/reduce-histogram-memory

Conversation

@SergeevVladislav

Copy link
Copy Markdown
Contributor

Deep or wide trees can exhaust GPU memory while constructing cumulative histograms because the raw histogram and its cumulative copy coexist. Previous-level histograms and split-score tensors also remain live after their last use.

Accumulate each histogram in place, release the parent histogram after subtraction, and discard split scores and the full current histogram when they are no longer needed. This preserves the tree-growth algorithm and array dtypes.

Validation:

  • Added 12 GPU tests covering exact histogram values, sampled rows/columns, parent subtraction, non-default streams, constrained memory pools, full training, and the no-split/early-stopping path. Against upstream master, three fail with CUDA OOM; all 12 pass with this change. Without a CUDA device, all 12 skip.
  • A synthetic regression fit with 4,096 rows, 64 features, depth 8, and two trees fails on master under a 32 MiB training-pool limit and succeeds after this change. In an unrestricted run with validation, peak live pool allocations decrease from 37.57 to 21.64 MiB. This is live pool usage, not total device memory or a universal reduction.
  • Separate before/after comparisons cover regression, binary and multiclass losses, multiple outputs, random projection, constant Hessians, one-vs-all trees, sampling, and multiple gradient steps. Predictions and validation histories agree within floating-point tolerances.
  • Independently based on master; cherry-picks of Fix zero output batch size in GPU histogram construction #29 and Add NaN-aware binary loss and multilabel metrics #30 apply cleanly. The combined suite passes all 91 tests.

Tested on RTX 5090 with CuPy 14.2.0 and CUDA runtime 12.9. This is not a V100 hardware benchmark. Large histograms and split-score arrays can still exhaust VRAM at high depth.

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.

2 participants