流式构建LightGBM Dataset后分布式训练停滞问题求助

你操作中存在的核心问题 #

• 调用非公开内部方法：_push_rows是LightGBM Dataset类的内部实现细节，未对外提供兼容性承诺，其流式添加行的逻辑没有考虑分布式环境下的特征分桶同步，是导致后续问题的根源。
• 特征分桶（Total Bins）不一致：投票并行模式的分布式训练要求所有worker的特征分桶边界完全对齐，一旦各节点分桶数量或边界不同，迭代更新时worker间无法达成树结构共识，会卡在同步步骤，出现CPU/网络空闲但进程停滞的情况。
• 数据集构建流程不规范：手动补充零值、标记数据集完成的操作未在分布式环境下同步执行，导致不同worker的数据集状态存在差异，进一步加剧分桶不一致问题。

正确的流式/分布式数据集构建方式 #

方法1：使用官方流式Dataset API（LightGBM 3.0+） #

LightGBM 3.0及以上版本提供了官方支持的流式数据集构建能力，无需依赖内部方法：

import lightgbm as lgb

# 所有worker必须使用完全相同的参考数据集配置（特征数量、顺序、类别特征等）
ref_dataset = lgb.Dataset(...)

# 初始化流式数据集，绑定参考数据集以确保分桶一致
streaming_dataset = lgb.StreamingDataset(
    ref_dataset=ref_dataset
)

# 循环加载数据块并添加
for data_chunk in load_data_chunks():
    # 确保每个数据块的特征结构与参考数据集完全匹配
    streaming_dataset.push(data_chunk)

# 所有worker同步设置标签与权重
streaming_dataset.set_label(labels)
streaming_dataset.set_weight(weights)

# 配置分布式投票并行训练参数
params = {
    'boosting_type': 'gbdt',
    'objective': 'binary',
    'metric': 'auc',
    'num_leaves': 31,
    'verbose': 1,
    'parallel': 'voting'
}

# 启动训练（需提前完成分布式通信环境初始化，如MPI）
booster = lgb.train(
    params,
    train_set=streaming_dataset,
    num_boost_round=100
)

方法2：预统一特征分桶（适用于LightGBM 3.0以下版本） #

如果使用旧版本LightGBM，需先在所有worker上同步特征分桶信息：

1. 提前在单节点生成参考数据集的分桶信息，序列化后分发到所有worker；
2. 每个worker基于统一的分桶信息构建数据集：

import lightgbm as lgb

# 所有worker加载相同的参考数据集（或其序列化的分桶信息）
ref_dataset = lgb.Dataset.load('ref_dataset_binning.bin')

# 逐步合并数据块，每次合并都绑定参考数据集确保分桶一致
temp_dataset = None
for data_chunk in load_data_chunks():
    chunk_dataset = lgb.Dataset(data_chunk, reference=ref_dataset)
    temp_dataset = chunk_dataset if temp_dataset is None else temp_dataset + chunk_dataset

# 同步设置标签与权重
temp_dataset.set_label(labels)
temp_dataset.set_weight(weights)

# 启动分布式训练
booster = lgb.train(params, temp_dataset, num_boost_round=100)

关键注意事项 #

• 禁止使用非公开内部方法：内部实现无分布式场景适配，版本迭代中可能随时变更。
• 强制分桶一致性：所有worker必须使用完全相同的特征分桶规则，这是投票并行训练的核心前提。
• 同步数据集操作：分布式环境下，所有worker的数据集构建步骤（添加数据、设置标签等）必须完全同步执行，避免状态差异。

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