NVIDIA RAPIDS环境下CUDA内存不足问题求助

尝试将3.7GB的独热编码训练数据集导入NVIDIA RAPIDS的Jupyter Notebook，无论是用cudf.read_csv直接加载，还是用dask_cudf分块加载后训练dask-xgboost模型，都反复出现CUDA内存不足错误（std::bad_alloc: out_of_memory）。单GPU环境，GPU显存8GB。

硬件与软件配置 #

• CPU：i7 9700F @4.00GHz
• GPU：RTX 3070 8GB GDDR6
• RAM：16GB @3600MHz
• 系统：Windows 11（宿主机）；运行RAPIDS的Ubuntu 18.04.05（容器/虚拟机）
• RAPIDS版本：22.12
• CUDA版本：12.0
• NVIDIA-SMI版本：528.02

尝试过的操作与代码 #

直接加载数据集（失败） #

# 因内存问题无法运行
X_train = cudf.read_csv('one-hot-train.csv', index_col = 0)

分块加载数据集（成功加载，但训练失败） #

# 可正常导入数据
X_train = dask_cudf.read_csv('one-hot-train.csv', chunksize = "4GB")
X_train = X_train.drop(columns = ['Unnamed: 0'])

# 因y数据集过小，导入方式不影响
y_train = dask_cudf.read_csv('y-train.csv')
y_train = y_train.drop(columns = ['Unnamed: 0'])

dask-xgboost训练代码（失败） #

xgb_params = {
    'learning_rate': 0.3,
    'objective': 'binary:logistic',
    'tree_method': 'gpu_hist',
    'max_depth': 6,
    'seed': 555,
    'predictor': 'gpu_predictor',
    'eval_metric': 'aucpr',
    'n_estimators': 5000,
}

# 因内存问题无法运行
xgb_model = dask_xgb.XGBClassifier(**xgb_params)
xgb_model.fit(X_train, y_train)

数据集大小 #

!du -sh one-hot-train.csv
> 3.7G  one-hot-train.csv

!du -sh y-train.csv
> 10M   y-train.csv

错误信息 #

MemoryError: std::bad_alloc: out_of_memory: CUDA error at: ~/miniconda3/envs/rapids-22.12/include/rmm/mr/device/cuda_memory_resource.hpp

解决方案建议 #

1. 优化数据集加载与内存占用 #

• 缩小chunksize：当前chunksize="4GB"接近GPU显存上限，建议缩小到1GB或更小，让dask每次只将部分数据加载到GPU显存处理。
• 过滤冗余特征：检查数据集中方差极低、重复或无意义的列，提前删除可大幅降低内存占用。
• 压缩数据类型：将int64转为int32/int8（根据数值范围）、float64转为float32，CUDA对32位数据支持更高效，内存占用减半：

  # 批量转换浮点列
  float_cols = X_train.select_dtypes(include=['float64']).columns
  X_train = X_train.astype({col: 'float32' for col in float_cols})
  # 批量转换整数列
  int_cols = X_train.select_dtypes(include=['int64']).columns
  X_train = X_train.astype({col: 'int32' for col in int_cols})
  

2. 调整XGBoost参数降低显存占用 #

• 削减n_estimators：n_estimators=5000会生成大量树，显存占用随树数量线性增长，先降到1000以内测试，后续再根据验证集表现调整。
• 启用随机采样：通过采样样本或特征减少单棵树的数据用量：

  xgb_params.update({
      'subsample': 0.8,  # 随机采样80%样本训练单棵树
      'colsample_bytree': 0.8  # 随机采样80%特征训练单棵树
  })
  

• 限制树复杂度：将max_depth降到4-5，或改用lossguide生长策略并限制叶子节点数：

  xgb_params.update({
      'grow_policy': 'lossguide',
      'max_leaves': 32  # 限制每棵树的叶子节点数量
  })
  

3. 优化RAPIDS内存配置 #

• 限制RMM显存池：手动设置RAPIDS内存管理器的显存池上限，给XGBoost留出空间：

  import rmm
  rmm.reinitialize(
      pool_allocator=True,
      initial_pool_size=6*1024**3  # 分配6GB显存池，留2GB给系统和XGBoost
  )
  

• 关闭冗余库：代码中导入的tensorflow和torch会占用显存，训练前注释掉或不导入：

  # import tensorflow as tf
  # import torch
  

4. 优化系统内存 #

16GB内存对于3.7GB数据集不算充裕，建议增大Ubuntu的swap空间，避免dask因内存不足频繁读写磁盘，间接影响GPU内存使用。

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