XGBoost 多线程训练时内存耗尽，通常是因为 CPU 版直方图构建和线程间数据拷贝放大了内存峰值；切换到 Dask 并行后，任务被分散到多个 worker，但显存（GPU 显存）监控反而容易变成盲区，因为 Dask 的 worker 进程和 XGBoost 的 GPU 上下文并不总是以一对一方式暴露给常用监控命令。你需要监控的对象不是 Jupyter 或调度器所在进程，而是每个 worker 进程实际绑定的 GPU 卡号、显存占用和 GPU 利用率，同时也要留意显存不足导致回退到 CPU 的静默行为。
切换 Dask 并行后，显存监控的核心是定位 worker 进程对应的 GPU 卡，而不是只看系统整体空闲显存。建议先用 nvidia-smi 锁定进程 PID 和 GPU 卡号，再用 Python 脚本周期采样每张卡的 used 和 per-process used memory。若出现显存不足，应先降低每 worker 的并发线程数或改用 device='cuda' 的显式分配策略，而不是直接调大 Dask 集群规模。
先确认 Dask 与 XGBoost 的实际并行边界
Dask 的并行单位是 worker 进程，每个 worker 通常包含一个线程池。XGBoost 的 Dask 接口（dask_xgboost 或 xgboost.dask）会为每个 worker 创建一个 XGBoost 训练器对象，但 GPU 资源分配由 XGBoost 的 n_gpus 参数和 CUDA 环境决定，并不自动等于 Dask worker 数量。很多情况下，多个 Dask worker 会共用同一块 GPU，或者一个 worker 内的多个线程同时向 GPU 提交计算，导致显存峰值显著高于串行训练。
因此第一步不是写监控脚本，而是先配置 Dask 集群，让 worker 数量与 GPU 卡数对应。通常做法是：先获取当前机器可见的 GPU 卡数量，再在创建 LocalCluster 或 SSHCluster 时指定 n_workers=GPU_N，并给每个 worker 设置环境变量 CUDA_VISIBLE_DEVICES。如果你使用 Kubernetes 或 Slurm，则应在部署文件中按 GPU 卡号隔离。否则，显存监控会因为没有稳定的进程-设备映射而失去意义。
用 nvidia-smi 建立进程与显存映射基线
训练开始时先运行一次 nvidia-smi，记录内存占用较高的进程。第二步执行：
nvidia-smi `--query-gpu`=index,memory.used,memory.total,utilization.gpu `--format`=csv -l 2这个命令每 2 秒输出一行，能显示整体显存趋势，但看不到具体是哪个 Dask worker 占用。若单个 GPU 显存持续升高，需要进一步查看进程级信息：
nvidia-smi `--query-compute-apps`=gpu_uuid,pid,used_memory `--format`=csv -l 5把输出的 PID 与 Dask worker 进程 PID 对照：在 Python 中可用 multiprocessing.current_process().pid 打印，或在命令行用 ps aux | grep dask-worker 查看。确认哪个 worker 的 PID 对应哪张 GPU，之后再进入周期采样监控。
写一个轻量监控脚本，按 worker 采样显存
下面的脚本适合在训练脚本外部单独运行，目的是记录每张 GPU 上所有 XGBoost 相关进程的显存变化，而不是等 OOM 后看日志。脚本会在当前目录生成 tsv 格式日志，训练结束后你可以在 CSV 里核对峰值时间点。
import csv, subprocess, time, os, sys
def get_pids_by_name(name):
out = subprocess.check_output(['pgrep', '-f', name]).decode().strip()
return set(out.split()) if out else set()
def sample():
pids = get_pids_by_name('dask-worker') | get_pids_by_name('xgboost')
rows = []
for pid in pids:
try:
out = subprocess.check_output([
'nvidia-smi', '`--query-compute-apps`=pid,used_memory', '`--format`=csv,nounits'],
stderr=subprocess.DEVNULL).decode().strip().split('\n')
except subprocess.CalledProcessError:
continue
for line in out[1:]:
cols = [c.strip() for c in line.split(',')]
if cols[0] == pid:
rows.append((pid, cols[1]))
return rows
if __name__ == '__main__':
interval = float(sys.argv[1]) if len(sys.argv) > 1 else 2.0
logname = 'gpu_mem_log.tsv'
with open(logname, 'w', newline='') as f:
w = csv.writer(f, delimiter='\t')
w.writerow(['time', 'pid', 'used_mib'])
while True:
t = time.strftime('%Y-%m-%d %H:%M:%S')
for pid, mem in sample():
w.writerow([t, pid, mem])
f.flush()
time.sleep(interval)运行：python gpu_watch.py 2。如果训练进程名不是默认的 dask-worker 或 xgboost，需要修改 get_pids_by_name 参数。该脚本只负责记录，不会干预训练；若显存持续逼近上限，应在训练脚本内设置回调提前终止，避免 OOM 后整个集群重启。
从 Dask 侧限制显存使用
监控发现显存不足时，优先调整 XGBoost 的 nthread 和 max_bin，而不是盲目增加 Dask worker。每个 worker 内线程数过高会同时申请多个 GPU 流，显存碎片化更严重。通常建议先设置：
params = {
'tree_method': 'gpu_hist',
'n_gpus': 1, # 每个 worker 只使用 1 块 GPU
'nthread': 4, # 控制线程并发，避免显存峰值
'max_bin': 256,
'gpu_id': 0 # 需要配合 CUDA_VISIBLE_DEVICES 使用
}在 Dask 接口下，gpu_id 不一定直接生效，更可靠的做法是启动 Dask worker 前设置环境变量：
CUDA_VISIBLE_DEVICES=0,1 dask-worker ... `--nworkers` 2 `--nthreads` 4这样每个 worker 只能看到分配的 GPU，避免擅自占用其余卡。若显存仍然不足，可以调低 max_bin 或改用 hist 的 CPU 直方图，但两者训练速度都会下降，需要结合具体数据量实测才能确定可接受范围。
常见问题
Dask 调度的 CPU 内存占用也会影响 GPU 训练吗？
会。Dask 调度器负责收集 worker 的梯度直方图，当树分裂时，worker 需要将 GPU 上的直方图数据复制到 CPU 内存，再发给调度器合并。若 worker 的 CPU 内存不足，训练会提前崩溃，且不会显示为显存不足。因此监控显存的同时，也建议用 free -m 记录系统内存变化。
为什么 nvidia-smi 显示空闲，但 XGBoost 报显存不够？
常见原因是 CUDA context 默认预留了一部分显存，且 XGBoost 的缓存池不会立即释放。如果多个 Dask worker 进程共享一张 GPU，每个进程都会创建独立的 CUDA context，累计预留量会明显上升。这时应减少共享同一 GPU 的 worker 数量，并检查 nvidia-smi 的 per-process used memory，而不是只看整体空闲显存。
监控脚本会影响训练性能吗？
每 2 秒调用一次 nvidia-smi 产生的开销通常很小，不过在极端高并发下建议延长间隔到 5 秒。更保守的做法是只用 watch -n 2 nvidia-smi 手工观察，确认瓶颈后再启动持久化脚本，避免训练日志和监控日志同时写入同一块磁盘。