利用单批次模型化训练超参数配置以提升单GPU利用率的optuna-saturate方法

《SoftwareX》：optuna-saturate: training hyperparameter configurations as one batched model to raise the utilisation of a single GPU

　　

小型神经网络（参数规模约104–106，在适中的批次上训练）的超参数优化（Hyperparameter Optimization, HPO）会使现代GPU大部分时间处于空闲状态：试验（trial）逐个执行，而单次训练步骤仅启动少量计算，导致设备在内核（kernel）之间等待。研究人员提出optuna-saturate库，将Optuna搜索空间划分为两类需要不同执行策略的超参数。保形（Shape-preserving）超参数——学习率、动量、权重衰减、dropout概率、softmax温度——不改变任何张量形状，因此此类配置可通过torch.func的stack_module_state与vmap作为单一批量化模型进行训练。变形状（Shape-changing）超参数改变网络架构，无法批量化。在与顺序Optuna相同试验数量、数据在训练过程中加载的条件下，八成员分组将24次试验研究的墙钟时间（wall-clock time）在Fashion-MNIST CNN上缩短6.3倍、KMNIST多层感知机（Multilayer Perceptron, MLP）上缩短7.0倍、CIFAR-10 Transformer上缩短4.0倍，同时GPU板级能耗降低2.2至3.6倍，最佳准确率保持在顺序执行臂的种子间波动范围内。部分增益来自输入流水线：分组中每个批次仅被加载一次供所有成员共享。使用四个加载器工作进程（worker process）时，顺序执行臂自身提速两倍，批量化优势在短研究中降至1.6倍，在单次试验工作量增大22倍的任务上仍保持4.0倍。研究还测量了Optuna基于线程的n_jobs、独立CUDA流上的线程、带内存准入控制的K个工作进程以及Ray Tune在同一任务上的表现：基于线程的策略慢于顺序执行，Ray Tune达到1.8至2.3倍加速，普通工作进程达到2.8至3.8倍加速。若以顺序执行臂的墙钟时间预算而非试验数量为基准，批量化搜索可评估6.6倍数量的配置（尽管覆盖的架构种类更少），并达到更高的最佳准确率。一致性测试（parity test）表明，在确定性float32条件下（SGD带动量与权重衰减、Adam和AdamW，分别用于线性、卷积和Transformer网络），批量化训练与独立训练在atol=10-5、rtol=10-4容差内复现一致；另有针对bfloat16训练和随机dropout的独立测试。这些测试支持在所述条件下训练等价性成立，而分组改变了搜索流程，其效果由评估部分单独度量。

论文解读

一、研究背景与问题

超参数优化是应用机器学习中的标准环节，Optuna、Ray Tune、SMAC3和Syne Tune等成熟框架提供了采样器、剪枝器和存储后端，它们决定了评估哪些配置以及评估多久，但都未解决一个关键问题：单个配置在评估期间对加速器的占用效率有多高。对于小型神经网络而言，这一缺口尤为严重，而小型网络在工程应用中极为常见——状态监测、剩余寿命预测和电池健康估计等领域 routinely 使用紧凑的卷积和循环模型，需要反复训练以进行交叉验证、种子稳定性分析和超参数搜索。一个参数量为10⁴–10⁶、层宽较窄的网络在适中批次上训练时，会发出大量小型CUDA内核，其训练步骤的墙钟时间主要由内核启动开销和Python端工作主导，而非算术运算本身。由于试验逐个执行，设备在整项研究中大部分时间处于空闲状态。

研究者在主基准测试中发现，顺序执行期间每次NVML利用率轮询均返回非零值，而实际运行的CUDA内核累计时长仅占墙钟时间的5.6%，且设备时钟频率仅为最大值的三分之一。NVML利用率定义为采样周期内存在内核执行的百分比，单个多处理器上的内核即计为完全占用，因此非零轮询无法区分"内核间空闲"与"真正忙碌"的设备。现有工具的并行方案同样无法弥合这一差距：Optuna的study.optimize(n_jobs=K)以线程运行试验，受Python全局解释器锁（GIL）和单设备争用限制；基于进程的工作者（Ray Tune和Syne Tune本地后端也采用此方式）虽避免锁竞争，但会倍增CUDA上下文、在设备间时间片轮转，且需要准入控制以避免内存耗尽。这些策略均未改变工作的形态——每个试验仍启动各自的小型内核。

将重复训练任务融合为单一任务、使小型内核变为大型内核的思想并非本研究首创：Horizontally Fused Training Arrays在算子级别融合重复作业的模型，在GPU和TPU上获得高达15倍的吞吐增益；Hydro在数据中心调优服务中应用试验间和试验内融合。但二者均构建为独立框架，以融合算子实现或集群调度器为核心。optuna-saturate的贡献更为聚焦，面向单GPU实践者：融合通过PyTorch自身的函数变换torch.func表达，无需维护融合算子；分组由未修改的Optuna study驱动并回传结果，用户已有的采样器、存储和仪表板继续可用；用户对可批量化超参数的声明在运行时被检查；批量化分组训练每个成员与独立训练等价这一声明，对库中提供的每种优化器都进行了数值测试。

二、研究内容与核心发现

本研究的核心观察是：超参数搜索空间并非同质，可分为两类需要不同执行策略的超参数。保形超参数——学习率、动量、权重衰减、dropout概率、softmax温度、损失权重——不改变任何张量形状，以可广播的值进入计算，仅在这些参数上不同的配置可被展平为一组带前导成员维度的参数，以单次传递训练。变形状超参数——隐藏宽度、嵌入大小、层数——改变参数形状，在这些参数上不同的成员无法共享批量化模型；不改变形状但改变计算结构的超参数（如固定嵌入大小下的注意力头数）同样无法批量化。该区分无法自动推断，因为超参数是否改变形状取决于目标函数如何使用它，且形状匹配是必要而非充分条件。因此optuna-saturate要求用户声明可批量化的参数名，组内固定所有其他超参数，并在运行时对照成员参数形状检查声明，不匹配即报错而非静默产生错误结果。

分组构造采用"先探后组"策略。由于分类需要完整的超参数名列表，而该列表仅在目标函数运行后才存在，研究者通过探针运行（probe run）打破循环：目标函数在首次请求模型时中止，此时所有声明已完成且尚未开始任何训练，因此发现搜索空间不消耗训练步骤。随后请求一个试验并读取其变形状值，再以固定这些值的方式入队K-1个试验，使采样器自由选择可向量化的超参数，最终得到K个共享架构、仅在可批量化参数上不同的真实试验。批量化训练中，torch.func的stack_module_state将K个模块实例展平为带前导成员维度的张量字典，functional_call在vmap下以单次传递评估所有成员。损失在批次维度上归约但不跨成员归约，返回K个值；由于成员间不共享参数，对成员i损失的梯度即为其自身损失的梯度，一次反向传播即可训练整个分组。库内置三个向量化优化器VectorizedSGD、VectorizedAdam和VectorizedAdamW，其更新规则逐操作对应torch.optim的SGD、Adam和AdamW，包括Nesterov动量、耦合与解耦权重衰减、Adam偏置校正以及参数级步数计数。dropout以基于逐成员概率的Bernoulli掩码实现，vmap以randomness="different"调用使各成员获得独立掩码。前向和损失可在torch.autocast下运行，参数与优化器状态保持float32，与普通混合精度循环一致，不应用损失缩放，因此bfloat16为经过测试的推荐数据类型。

三、关键技术方法

研究者采用的主要技术方法包括：利用PyTorch的torch.func函数变换（stack_module_state与vmap）实现多模型批量化训练；开发三个向量化优化器（VectorizedSGD、VectorizedAdam、VectorizedAdamW），其更新规则逐操作对应标准torch.optim实现；通过探针运行（probe run）在零训练成本下发现搜索空间；采用运行时声明检查机制，对照成员参数形状验证用户声明的可批量化超参数；使用NVML轮询和CUDA内核追踪构建双指标分析器（busy_fraction与kernel_time_fraction）以量化设备空闲程度；提供recommend_group_size()函数基于PyTorch分配统计拟合内存-组大小线性关系以推荐最优组大小。基准测试使用Fashion-MNIST CNN、KMNIST多层感知机和CIFAR-10 Transformer三个模型族，数据在训练过程中加载。

四、研究结果

3.1 顺序执行的设备利用率测量：通过NVML轮询和CUDA内核追踪发现，顺序执行期间每次轮询均返回非零利用率，但实际内核累计时长仅占墙钟时间的5.6%，设备时钟为最大值的三分之一，证明NVML利用率无法区分内核间空闲与真正忙碌的设备。

3.2 现有并行策略对比：在相同任务上测量Optuna线程式n_jobs、独立CUDA流上的线程、带内存准入控制的K个工作进程以及Ray Tune。基于线程的策略慢于顺序执行，Ray Tune达到1.8至2.3倍加速，普通工作进程达到2.8至3.8倍加速，均未改变工作形态。

3.3 批量化加速效果：在相同试验数量下，八成员分组将24次试验研究的墙钟时间在Fashion-MNIST CNN上缩短6.3倍、KMNIST MLP上缩短7.0倍、CIFAR-10 Transformer上缩短4.0倍，GPU板级能耗降低2.2至3.6倍，最佳准确率保持在顺序执行臂的种子间波动范围内。部分增益来自输入流水线——分组中每个批次仅加载一次供所有成员共享。

3.4 组大小与内存/加速权衡：通过recommend_group_size()拟合内存-组大小线性关系和加速-组大小曲线，测量了组大小增长时内存近似线性增长而加速增益递减的规律。

3.5 输入流水线影响：使用四个加载器工作进程时顺序执行臂自身提速两倍，批量化优势在短研究中降至1.6倍，在单次试验工作量增大22倍的任务上仍保持4.0倍。

3.6 固定时间预算下的搜索效率：以顺序执行臂的墙钟时间预算而非试验数量为基准，批量化搜索可评估6.6倍数量的配置（覆盖的架构种类更少），并达到更高的最佳准确率。

3.7 训练等价性测试：在确定性float32条件下，批量化训练与独立训练在atol=10^-5、rtol=10^-4容差内复现一致，覆盖SGD带动量与权重衰减、Adam和AdamW在三种网络上的组合；另有bfloat16训练和随机dropout的独立测试，支持在所述条件下训练等价性成立。

五、讨论与结论

讨论部分指出，分组改变了搜索流程本身——采样器在组内对可向量化超参数的自由选择受到架构固定约束，这使搜索在更少的架构种类上评估更多配置，其效果由评估部分单独度量。批量化训练与独立训练的等价性在所述条件下经数值测试验证，但分组对搜索流程的改变是独立于训练等价性的另一维度。结论部分表明，optuna-saturate通过将超参数搜索空间分类为保形与变形状两类，以PyTorch自身的函数变换实现多配置批量化训练，在单GPU上以相同试验数量获得4.0至7.0倍墙钟时间加速和2.2至3.6倍能耗降低，且最佳准确率保持在顺序执行的种子间波动范围内；在固定时间预算下可评估6.6倍数量的配置并达到更高最佳准确率。该库以MIT许可证开源（v0.2.0），代码仓库位于https://github.com/kurekj/optuna-saturate/tree/v0.2.0，支持Python 3.10+、PyTorch 2.2+和Optuna 4.0+。