使用 ray + light gbm + limited memory

Question

So, I would like to train a lightGBM on a remote, large ray cluster and a large dataset. Before that, I would like to write the code such that I can run the training also in a memory-constrained setting, eg my local laptop, where the dataset does not fit in-mem. That will require some way of lazy loading the data.

The way I imagine it, I should be possible with ray to load batches of random samples of the large dataset from disk (multiple.pq files) and feed them to the lightgbm training function. The memory should thereby act as a fast buffer, which contains random, loaded batches that are fed to the training function and then removed from memory. Multiple workers take care of training + IO ops for loading new samples from disk into memory. The maximum amount of memory can be constrained to not exceed my local resources, such that my pc doesn't crash. Is this possible?

I did not understand yet whether the LGBM needs the full dataset at once, or can be fed batches iteratively, such as with neural.networks, for instance. So far, I have tried using the lightgbm_ray lib for this:

from lightgbm_ray import RayDMatrix, RayParams, train, RayFileType

# some stuff before 
... 

# make dataset
data_train = RayDMatrix(
    data=filenames,
    label=TARGET,
    feature_names=features,
    filetype=RayFileType.PARQUET,
    num_actors=2,
    lazy=True,
)

# feed to training function
evals_result = {}
bst = train(
    params_model,
    data_train,
    evals_result=evals_result,
    valid_sets=[data_train],
    valid_names=["train"],
    verbose_eval=False,
    ray_params=RayParams(num_actors=2, cpus_per_actor=2)
)

I thought the lazy=True keyword might take care of it, however, when executing this, I see the memory being maxed out and then my app crashes.

Thanks for any advice!

Answer 1

LightGBM requires loading the entire dataset for training, so in this case, you can test on your laptop with a subset of the data (ie only pass a subset of the parquet filenames in).

The lazy=True flag delays the data loading to be split across the actors, rather than loading into memory first, then splitting+sending to actors. However, this would still load the entire dataset into memory, since all actors are on the same (local) node.

Additionally, when you do move to running on the remote cluster, these tips might be helpful to optimize memory usage: https://docs.ray.io/en/latest/train/gbdt.html?highlight=xgboost%20memro#how-to-optimize-xgboost-memory-usage .