基于Ray+LightGBM的内存受限场景大数据集训练问题咨询

1. 精准控制RayDMatrix的分片大小 #

lazy=True仅实现延迟加载，但如果每个Ray Actor分配到的文件分片过大，依然会触发内存溢出。改用shard_size参数按内存容量限制每个Actor加载的数据量，而非依赖num_actors均分文件：

data_train = RayDMatrix(
    data=filenames,
    label=TARGET,
    feature_names=features,
    filetype=RayFileType.PARQUET,
    shard_size="100MB",  # 按内存大小分片，避免大文件导致单Actor过载
    lazy=True,
)

2. 分离训练/验证集的内存占用 #

当前代码将训练集直接作为验证集，会导致两份全量数据的内存开销。单独创建验证集的RayDMatrix，还可通过采样进一步缩小验证集规模：

# 单独创建验证集，设置更小分片或采样比例
data_valid = RayDMatrix(
    data=valid_filenames,
    label=TARGET,
    feature_names=features,
    filetype=RayFileType.PARQUET,
    shard_size="50MB",
    lazy=True,
    sampling=0.1  # 仅用10%数据做验证
)

# 训练时传入独立验证集
bst = train(
    params_model,
    data_train,
    evals_result=evals_result,
    valid_sets=[data_valid],
    valid_names=["valid"],
    verbose_eval=False,
    ray_params=RayParams(num_actors=2, cpus_per_actor=2)
)

3. 配置LightGBM核心内存控制参数 #

在params_model中添加算法层面的内存优化参数，降低单轮训练的内存消耗：

params_model = {
    # 原有业务参数
    "objective": "binary",
    "metric": "auc",
    # 内存优化参数
    "max_bin": 128,  # 降低直方图分箱数（默认256）
    "bagging_fraction": 0.8,  # 每次迭代采样80%数据
    "feature_fraction": 0.8,  # 每次迭代采样80%特征
    "min_data_in_leaf": 50,  # 限制叶子节点最小数据量，减少内存占用
    "verbose": -1
}

4. 强制Ray集群的内存资源管控 #

通过Ray初始化和训练参数，限制每个Actor及整体集群的内存使用上限：

import ray
# 本地初始化Ray时限制全局内存
ray.init(
    memory=4 * 1024**3,  # 4GB进程内存
    object_store_memory=2 * 1024**3  # 2GB对象存储内存
)

# 训练时指定每个Actor的内存上限
ray_params = RayParams(
    num_actors=2,
    cpus_per_actor=2,
    memory_per_actor=2 * 1024**3  # 每个Actor最多占用2GB内存
)

5. 手动实现分批增量训练（兜底方案） #

如果上述方法仍无法解决，可基于Ray Dataset手动实现全可控的分批迭代训练：

import ray
from lightgbm import train as lgb_train

# 加载Parquet数据集为可迭代对象
ds = ray.data.read_parquet(filenames).repeat()  # 重复迭代数据集

# 初始化模型
bst = None

# 分批训练
for batch in ds.iter_batches(batch_size=100_000, batch_format="pandas"):
    # 转换为LightGBM原生Dataset
    lgb_data = lgb.Dataset(batch[features], label=batch[TARGET])
    # 增量训练
    bst = lgb_train(
        params_model,
        lgb_data,
        num_boost_round=10,
        init_model=bst,
        verbose_eval=False
    )
    # 可按需添加早停、验证集评估逻辑

内容的提问来源于stack exchange，提问作者user19976975