Validation ensemble predictions use float64 and consume twice the necessary GPU memory #34

Description

Validation ensemble predictions are allocated as float64, while training ensemble predictions use float32. For the same shape, the validation buffer therefore consumes twice as much GPU memory.

This is avoidable memory overhead in multioutput/multiclass training and can contribute to out-of-memory failures.

Root cause

GradientBoosting._create_build_info() explicitly selects float32 for the training ensemble but omits dtype for validation ensembles:

ens = cp.empty(
    (y.shape[0], self.base_score.shape[0]), order="C", dtype=cp.float32,
)
val_ens = [
    cp.empty((x.shape[0], self.base_score.shape[0]), order="C")
    for x in y_val
]

CuPy's default dtype for empty() is float64.

Source in master.

Minimal reproduction

Requires a working CUDA/CuPy installation. This subclass only prints the allocated buffers and otherwise uses the existing training implementation.

import numpy as np
from py_boost import GradientBoosting

class InspectBoost(GradientBoosting):
    def _create_build_info(self, *args, **kwargs):
        builder, info = super()._create_build_info(*args, **kwargs)
        train = info["data"]["train"]["ensemble"]
        valid = info["data"]["valid"]["ensemble"][0]
        print("train", train.dtype, train.shape, train.nbytes)
        print("valid", valid.dtype, valid.shape, valid.nbytes)
        return builder, info

rng = np.random.default_rng(42)
X = rng.normal(size=(64, 4)).astype(np.float32)
y = (2 * X[:, 0] - X[:, 1]).astype(np.float32)

model = InspectBoost(
    loss="mse", ntrees=2, max_depth=2, es=0, verbose=1000,
)
model.fit(X, y, eval_sets=[dict(X=X.copy(), y=y.copy())])

Observed output:

train float32 (64, 1) 256
valid float64 (64, 1) 512

Impact

For each validation set with N rows and K outputs, this buffer uses 8 * N * K bytes instead of 4 * N * K bytes. At 100,000 rows and 1,000 outputs, that is 800 MB instead of 400 MB, excluding other arrays and allocator overhead.

Related to the memory discussion in #8, but fixing this allocation alone will not make arbitrary large-class datasets fit on a GPU. Histogram memory improvements in #31 address a different allocation.

Expected behavior and suggested fix

Allocate validation ensemble predictions with an explicit dtype=cp.float32, consistently with the training ensemble. Verify any effects on validation metrics and early stopping, since changing accumulation precision may change results slightly.

Verification

The output above was reproduced on Linux with an NVIDIA A100, Python 3.10.22, NumPy 2.2.6, and CuPy 13.6.0. Training with the inspection subclass completed successfully. The allocation remains unchanged in master at the linked commit.

Suggested priority: P2.

Acceptance criteria

• Training and every validation ensemble use float32.
• Validation buffers use 4 * N * K bytes rather than 8 * N * K bytes.
• Coverage includes multiple validation sets and a multioutput target.
• Validation metrics and early stopping remain correct within an appropriate floating-point tolerance.