CUDA Illegal Memory Access Error when using SketchBoost with FilterSketch

[Lopa10ko]

Issue Description

When fitting SketchBoost with FilterSketch method, the training process fails with a CUDA illegal memory access error. The error occurs during model training and appears to be related to GPU memory operations.

Traceback (most recent call last):
  File "cupy_backends/cuda/api/driver.pyx", line 234, in cupy_backends.cuda.api.driver.moduleUnload
  File "cupy_backends/cuda/api/driver.pyx", line 63, in cupy_backends.cuda.api.driver.check_status
cupy_backends.cuda.api.driver.CUDADriverError: CUDA_ERROR_ILLEGAL_ADDRESS: an illegal memory access was encountered

Dataset Specifications

Training Data

• Features: (529, 1185) - 529 samples, 1185 features
• Target: (529, 27) - 529 samples, 27 output dimensions

Test Data

• Features: (133, 1185) - 133 samples, 1185 features
• Target: (133, 27) - 133 samples, 27 output dimensions

Environment Information

• CUDA Version: 12.2
• Py-Boost Version: 82b9676
• Python Version: 3.10.5

Reproduction Steps

from py_boost import GradientBoosting, SketchBoost

params = {
    # Basic training parameters
    'loss': 'multilabel',
    'metric': None,
    'ntrees': 10000,
    'lr': 0.1,
    'es': 15,
    'seed': 42,
    'verbose': 10,
    
    # Tree structure parameters
    'max_depth': 6,
    'min_gain_to_split': 0,
    'min_data_in_leaf': 10,
    'lambda_l2': 1,
    
    # Data sampling parameters
    'subsample': 0.05,
    'colsample': 1.0,
    
    # Quantization parameters
    'quantization': 'Quantile',
    'max_bin': 256,
    'min_data_in_bin': 3,
    'quant_sample': 2000000,
    
    # Gradient boosting parameters
    'gd_steps': 1,
    'use_hess': False,
    
    # Sketching parameters
    'sketch_method': 'filter',  # this is important!
    'sketch_outputs': 1,
    'sketch_params': None,
    
    # Callbacks
    'callbacks': None
}

with np.load('train_compressed.npz') as data:
    X_train = data['X_train']
    y_train = data['y_train']
with np.load('test_compressed.npz') as data:
    X_test = data['X_test']
    y_test = data['y_test']

model = SketchBoost(**params)
model.fit(X_train, y_train, eval_sets=[{'X': X_train, 'y': y_train}, {'X': X_test, 'y': y_test}])

Data archive

train_test_data.zip