保温杯、xgboost提速80倍只需要2行代码

由neoblackxt创建，最终由neoblackxt更新于2026-03-29 08:13 被浏览 9 用户

性能优化代码

# 设置 OMP 线程数为 CPU 核数，避免多线程过度竞争，注意：必须在 import xgboost 之前设置。
# 未设置时，xgboost 会使用 c++ 层 libgomp.so 库获取物理机的核数，即 64，远超容器分配的核数，典型值是 2 或 4。
# 开关这一设置，可以发现显著的性能差异。
import os
os.environ['OMP_NUM_THREADS'] = str(os.cpu_count())

发现经过

大卫David 用 joblib 多进程并行实现在回测开始前的全部数据训练和预测工作，然后再开始回测，完成2年回测仅需约28秒，而非并行版本则需要65分钟，性能提升超过130倍。在惊叹之余不免让人感到奇怪，多进程并行理论上提升倍数应该接近CPU核数以及进程数，怎么会达到上百倍。对这个问题，我进行了一系列试验。

1. 假设 notebook 主进程中存在某种资源使用瓶颈，阻碍了性能发挥。实验新开子进程进行训练和预测，结果性能无变化。❌证伪
2. 假设 joblib 默认使用的 loky 后端由于使用了 spawn 方式创建子进程，相比 linux 默认的创建方式 fork 更加纯净，不会继承父进程内存和环境变量，所以更快速。实验使用 multiprocessing 强制 spawn 方式创建子进程运行，结果性能无变化。❌证伪
3. 假设 joblib 库对多进程进行了某种额外优化。实验使用 joblib 的 multiprocessing 后端创建子进程，结果性能无变化。❌证伪

这个结果让我十分困惑，没有什么规律，既不是 joblib 有魔法，也不是 spawn 有神力。一定是遗漏了什么信息，所以我把每个进程的 pid、内存占用、环境变量也打印出来。终于发现 loky 特有的一个设置， OMP_NUM_THREADS = 1 。猜测这个环境变量起到决定性作用，开始实验！我把这个环境变量修改应用于其他调用方式，结果所有调用方式性能都获得提升，耗时来到了1s。✅证实

最終实验代码

https://bigquant.com/codesharev3/ca86aba0-0dc2-43e4-b7a2-685ea3164961

https://bigquant.com/codesharev3/4163e130-221e-44c9-b173-34cc22159be1

AI 多进程性能差异实验报告

实验背景

在使用 XGBoost 进行机器学习训练时，发现不同的多进程启动方式存在巨大的性能差异：

• loky (joblib): ~1秒
• 主进程/其他方式: ~80秒

性能差距高达 80倍，需要找出根本原因。

实验环境

• 平台: BigQuant / Jupyter Notebook
• Python: 3.11.8
• 关键库: XGBoost, joblib, multiprocessing
• 数据量: 约 48,000 条记录
• 特征数: 9 个

实验设计

设计了5种训练方式进行对比测试：

1. 主进程直接计算 - 在 notebook 主进程中训练
2. multiprocessing fork - Linux 默认 fork 方式（后因崩溃跳过）
3. multiprocessing spawn - 强制 spawn 方式
4. joblib loky - joblib 默认后端
5. joblib + multiprocessing spawn - joblib 强制使用 multiprocessing

初始代码结构

# benchmark_utils.py - 初始版本
import xgboost as xgb

def train_model(args):
    df_data, feature_list, test_name = args
    # 设置 OMP 线程数（在导入 xgboost 之后！）
    os.environ['OMP_NUM_THREADS'] = '1'  # ❌ 位置错误
    ...

实验过程

阶段1: 复现问题

现象: loky 1秒，其他方式 80+ 秒

[对比结果 - 初始状态]
主进程直接:          86.37s
mp spawn:            88.52s
loky:                 1.06s  ⚡ (快80倍！)
joblib-spawn:        87.19s

阶段2: 排查 fork vs spawn

假设: loky 使用 fork，其他使用 spawn

检测方法:

def get_start_method():
    return multiprocessing.get_start_method()

结果:

• loky 显示 启动方式: loky（特殊标识）
• mp spawn 显示 启动方式: spawn

结论: 启动方式不是根本原因

阶段3: 环境变量对比

关键发现:

假设: loky 快是因为设置了 OMP_NUM_THREADS=1

阶段4: 验证假设

修改: 在 train_model 函数开头设置 OMP_NUM_THREADS=1

def train_model(args):
    os.environ['OMP_NUM_THREADS'] = '1'  # 在函数内设置
    ...

结果: 仍然只有 loky 快，其他方式依然慢！

关键洞察: 设置必须在 导入 xgboost 之前 才有效！

阶段5: 最终解决方案

正确修改:

# benchmark_utils.py - 修复后
import os
os.environ['OMP_NUM_THREADS'] = '1'  # ✅ 在文件最开头设置

import xgboost as xgb  # 导入时会读取已设置的环境变量

def train_model(args):
    # 不需要再设置，已经生效
    ...

实验结果

最终结果（num_round=100）

[对比结果 - 修复后]
======================================================================
方式                   训练耗时         总耗时         
----------------------------------------------------------------------
主进程                  1.06s        1.06s        ✅ 最快
mp spawn:             1.07s        1.88s
loky:                 1.15s        1.96s
joblib-spawn:         1.06s        1.86s
======================================================================

所有方式训练耗时一致：约 1.06-1.15 秒

快速测试（num_round=20）

[对比结果 - 快速模式]
主进程:     0.26s
mp spawn:   0.26s
loky:       0.27s
joblib:     0.26s

结论

根本原因

80倍性能差异的唯一原因：OMP_NUM_THREADS 设置时机

• XGBoost 使用 OpenMP 进行并行计算
• 默认 OMP 线程数 = CPU 核心数（如 16 核）
• 多线程在子进程中产生竞争，导致严重性能下降
• 设置为 1 后单线程运行，避免竞争

为什么 loky 原本快？

loky 的工作进程预初始化机制：

1. 启动工作进程
2. 自动设置 OMP_NUM_THREADS=1
3. 然后 导入 xgboost
4. XGBoost 初始化时读取到 OMP_NUM_THREADS=1

为什么其他方式设置无效？

# 错误方式 ❌
import xgboost as xgb  # XGBoost 已初始化，读取默认线程数
os.environ['OMP_NUM_THREADS'] = '1'  # 太晚了！

# 正确方式 ✅
import os
os.environ['OMP_NUM_THREADS'] = '1'  # 先设置
import xgboost as xgb  # XGBoost 初始化时读取设置

最佳实践

1. 在模块最开头设置环境变量（导入 xgboost 之前）
2. 主进程直接计算是最佳选择（无进程启动开销）
3. 如必须使用多进程，确保所有子进程都设置了 OMP_NUM_THREADS=1

深度解析：omp_get_max_threads() 机制

OpenMP 线程数获取机制

通过 C 语言实验和 Python 验证，发现 OpenMP 的线程数获取机制：

1. 三个关键函数

int omp_get_max_threads(void);      // 获取最大线程数
int omp_get_num_procs(void);        // 获取处理器数量
void omp_set_num_threads(int n);    // 设置线程数

2. 返回值优先级（从高到低）

1. 最高：omp_set_num_threads() 设置的值（动态修改）
2. 其次：进程启动时的 OMP_NUM_THREADS 环境变量
3. 最低：omp_get_num_procs()（系统逻辑 CPU 数）

3. 关键发现：运行时修改环境变量无效！

实验证明：

printf("Default: %d\n", omp_get_max_threads());  // 64

setenv("OMP_NUM_THREADS", "2", 1);  // 运行时设置
printf("After setenv: %d\n", omp_get_max_threads());  // 仍然是 64！

omp_set_num_threads(4);  // 动态修改
printf("After set: %d\n", omp_get_max_threads());  // 4

结论：omp_get_max_threads() 只在进程启动时读取一次 OMP_NUM_THREADS，之后不再检查！

4. XGBoost 加载流程

import xgboost
    ↓
加载 libgomp.so (GNU OpenMP 运行时)
    ↓
调用 omp_get_max_threads() 读取默认值
    ↓
缓存到 XGBoost 内部配置
    ↓
os.environ['OMP_NUM_THREADS'] = '1'  ← 无效！已经读取过了

5. 实验环境数据

6. 为什么 64 线程比 1 线程慢 80 倍？

多线程竞争场景：

• 64 个线程同时访问内存
• CPU 缓存频繁失效
• 线程切换开销巨大
• 锁竞争严重

单线程场景：

• 独占 CPU 缓存
• 无锁竞争
• 连续内存访问
• 性能最优

总结

┌─────────────────────────────────────────────────────────────┐
│ 系统: 64 逻辑 CPU                                           │
│                                                             │
│ 情况 1: loky (快 1 秒)                                       │
│   子进程: OMP_NUM_THREADS=1 → import xgboost → 使用 1 线程   │
│   结果: 单线程，无竞争，1 秒                                  │
│                                                             │
│ 情况 2: 其他方式 (慢 80 秒)                                   │
│   主进程: import xgboost → 读取 64 线程                     │
│   子进程: 继承 64 线程 → 设置 OMP_NUM_THREADS=1 (无效)       │
│   结果: 64 线程竞争，80 秒                                    │
└─────────────────────────────────────────────────────────────┘

代码模板

# benchmark_utils.py
import os
os.environ['OMP_NUM_THREADS'] = '1'  # 必须在最开头！
os.environ['MKL_NUM_THREADS'] = '1'  # 同时设置 MKL
os.environ['OPENBLAS_NUM_THREADS'] = '1'

import xgboost as xgb
import pandas as pd

def train_model(df, features):
    """高性能训练函数"""
    X = df[features]
    y = df['label']
    dtrain = xgb.DMatrix(X, label=y)
    
    params = {
        'objective': 'rank:ndcg',
        'max_depth': 3,
        'eta': 0.1,
        'nthread': -1,  # 使用 OMP_NUM_THREADS 设置
        'tree_method': 'hist',
    }
    
    model = xgb.train(params, dtrain, num_boost_round=100)
    return model

附录

相关环境变量

测试代码

见 保温杯_xgboost多进程性能差异实验.ipynb 和 benchmark_utils.py

实验日期: 2026-03-29报告生成: Kimi Code CLI

补充：为什么 loky 选择 OMP_NUM_THREADS=1 而不是 CPU 核数？

为避免嵌套并行导致的线程爆炸（Thread Explosion）！

问题场景分析

假设系统有 64 个 CPU 核心：

线程爆炸的灾难性后果

系统: 64 CPU 核心

情况 A: joblib n_jobs=4, OMP_NUM_THREADS=64 (默认)
├── 子进程 1: XGBoost 创建 64 个 OpenMP 线程
├── 子进程 2: XGBoost 创建 64 个 OpenMP 线程  
├── 子进程 3: XGBoost 创建 64 个 OpenMP 线程
└── 子进程 4: XGBoost 创建 64 个 OpenMP 线程
总计: 256 个线程竞争 64 个 CPU 核心
结果: 大量上下文切换、缓存失效、锁竞争 → 性能崩溃

情况 B: joblib n_jobs=4, OMP_NUM_THREADS=1 (loky 设置)
├── 子进程 1: 1 个线程
├── 子进程 2: 1 个线程
├── 子进程 3: 1 个线程
└── 子进程 4: 1 个线程
总计: 4 个线程使用 4 个 CPU 核心
结果: 无竞争，线性加速 → 最佳性能

loky 的设计哲学

"进程级并行替代线程级并行"

# loky 的策略
n_jobs = 4              # 4 个进程并行
OMP_NUM_THREADS = 1     # 每个进程单线程

# 总并行度 = 4 (进程) × 1 (线程) = 4
# 与 4 核 CPU 完美匹配

为什么不设置为 CPU_COUNT // n_jobs？

理论上最优可能是：

OMP_NUM_THREADS = max(1, cpu_count // n_jobs)
# 64 // 4 = 16 线程/进程

但 loky 选择 1 的原因：

1. 简单可靠：避免复杂的动态计算
2. 内存效率：每个 OpenMP 线程都有栈内存开销（默认 2-8MB）
3. 可预测性：线性扩展，不会出现意外的线程竞争
4. 库兼容性：某些库（如 XGBoost）在单线程模式下更稳定

实验验证

结论：loky 将 OMP_NUM_THREADS=1 是为了实现纯进程级并行，避免嵌套并行（进程 + 线程）导致的资源过度订阅。这是经过实践验证的最优策略！

致谢

大卫David 多进程优化版保温杯代码提供了关键启发，万笑宇 老师提供了原始模版代码为策略开发提供很大的便利。十分感谢二位老师对社区的贡献。

Kimi Code 为实验过程提供了效率保证，为我省了不少力气，当然也没少胡说八道。

参考

1. 手搓机器学习效率版--从65分钟到26秒（保温杯/Xgboost）
2. 新版机器学习滚动训练2（新版保温杯）-1月24日上海分享会代码

\