Databricks中SparkXGBClassifier训练小数据集内存不足问题排查

在Databricks环境中，通过xgboost.spark.SparkXGBClassifier结合PySpark训练仅约20k行的小数据集时，任务因内存不足失败，报错信息如下：

Job aborted due to stage failure: Could not recover from a failed barrier ResultStage. Most recent failure reason: Stage failed because barrier task ResultTask(66,0) finished unsuccessfully.

经排查（结合训练阶段Ganglia等实时指标，内存完全耗尽），确定为内存问题。该数据集的特征列是由输入text列（每行含400-600个字符长度不一的单词）生成的IDF矩阵。

即便仅使用默认组件构建浅层管道模型（未加入自定义转换器），内存占用仍急剧飙升。已尝试以下优化方案：

• 使用更多worker节点（最多12个，每个含4核CPU+14GB内存）
• 对DataFrame重新分区
• 调整集群配置（如executor核数、内存大小）
• 释放冗余变量

但仅当将数据集缩减至3k行，或把XGBoost树最大深度设为2时，内存占用才有效降低，其他措施均无明显效果。此外，无论任务成功还是失败，worker节点CPU使用率均极低（平均约15%）。

相关代码片段：

xgboost = SparkXGBClassifier(
    features_col="vectorizedFeatures", 
    label_col="label", 
    num_workers=2 # set to 2 because of 2 worker nodes
)

train_data, validation_data = df.randomSplit([0.8, 0.2], seed=42)

# Tokenization into separate words
tokenizer = Tokenizer(inputCol="Text",outputCol="tokens")
vectorizer = CountVectorizer(inputCol="tokens",outputCol="raw_features")
labelEncoder = StringIndexer(inputCol="category",outputCol="label", handleInvalid="keep")
idf = IDF(inputCol="raw_features",outputCol="vectorizedFeatures")

# Pipeline defining the order of the dataflow through the model
pipeline = Pipeline().setStages([tokenizer, vectorizer, labelEncoder, idf, xgboost])

# Definition of an evaluator
evaluator = MulticlassClassificationEvaluator(metricName="accuracy")

# Fitting the model
model = pipeline.fit(train_data)

内存占用过高的核心原因 #

1. IDF特征矩阵维度爆炸：每行文本含400-600个单词，20k行数据会生成极高维度的稀疏特征矩阵。虽然Spark用稀疏向量存储，但XGBoost Spark版在训练时，会将稀疏向量转换为密集格式（尤其是barrier模式下的分布式训练），导致内存瞬间膨胀。树深度较大时，XGBoost需要存储更多分裂节点统计信息，进一步加剧内存消耗——这也是调低树深度到2时内存占用明显下降的原因。

2. XGBoost barrier模式特性：SparkXGBClassifier默认使用barrier模式执行任务，要求所有worker节点同步数据。如果特征维度极高，每个worker需要加载和处理的特征数据量远超预期，即使增加worker数量，也无法有效分摊内存压力，因为每个节点都需要持有部分全局特征数据用于树分裂计算。

3. Pipeline隐式数据持久化：Pipeline执行时会自动缓存中间结果（比如IDF生成的vectorizedFeatures列），若未显式清理或调整缓存级别，这部分数据会持续占用executor内存，叠加XGBoost训练的内存需求，最终触发OOM。

CPU使用率极低的原因 #

1. 内存瓶颈导致CPU等待：executor内存耗尽时，系统会频繁触发GC（垃圾回收），CPU大部分时间用于内存回收而非训练计算，因此使用率极低。

2. 并行度不匹配：虽然设置了num_workers=2，但高维度特征导致XGBoost每个任务处理的数据量过大，无法充分利用CPU核心；同时barrier模式下任务需等待所有节点就绪才能执行，进一步降低CPU利用率。

3. 小数据集分布式训练开销：20k行属于小数据集，分布式训练的通信开销远超计算开销，CPU大部分时间在等待节点间的数据同步，而非执行计算逻辑。

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