CubicML: Automated ML for Large ML Systems Co-design with ML Prediction of Performance

📄 arXiv: 2409.04585v2 📥 PDF

作者: Wei Wen, Quanyu Zhu, Weiwei Chu, Wen-Yen Chen, Jiyan Yang

分类: cs.LG, cs.AI, cs.DC

发布日期: 2024-09-06 (更新: 2024-09-21)

💡 一句话要点

CubicML：利用机器学习自动优化大规模机器学习系统训练性能

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 机器学习 超参数优化 分布式训练 性能预测 大规模系统

📋 核心要点

  1. 大规模机器学习系统协同设计面临超参数空间爆炸的挑战，难以手动优化。
  2. CubicML利用机器学习模型预测训练性能，作为代理模型指导超参数搜索，提升效率。
  3. 实验表明，CubicML有效优化了Meta内部广告推荐模型和大型语言模型的训练速度。

📝 摘要（中文）

扩展深度学习模型已被证明可以有效提高机器学习模型的智能，特别是对于工业推荐模型和大型语言模型。大规模分布式机器学习系统和算法的协同设计（以最大化训练性能）对于其成功至关重要。随着规模的扩大，协同设计超参数的数量迅速增长，这给可行地找到系统性能最大化的最佳设置带来了挑战。在本文中，我们提出了CubicML，它使用机器学习来自动优化大型分布式机器学习系统的训练性能。在CubicML中，我们使用机器学习模型作为代理来预测训练性能，以提高搜索效率和性能建模的灵活性。我们证明了CubicML可以有效地优化Meta内部具有730亿参数的广告推荐模型和高达4050亿参数的大型语言模型的训练速度。

🔬 方法详解

问题定义：论文旨在解决大规模分布式机器学习系统训练性能优化问题。现有方法，如人工调参或网格搜索，在高维超参数空间中效率低下，难以找到最优配置，成为制约模型训练速度的关键瓶颈。

核心思路：论文的核心思路是利用机器学习模型来预测不同超参数配置下的训练性能，将性能预测问题转化为回归问题。该模型作为代理模型，代替实际训练过程，从而加速超参数搜索过程，降低计算成本。

技术框架：CubicML的整体框架包含以下几个主要阶段：1) 离线数据收集：通过随机或启发式方法采样不同的超参数配置，并在实际系统上运行训练，收集训练性能数据（如训练时间、吞吐量等）。2) 代理模型训练：使用收集到的训练性能数据训练机器学习模型，该模型用于预测给定超参数配置下的训练性能。3) 超参数搜索：使用优化算法（如贝叶斯优化、进化算法等）在超参数空间中搜索，并使用代理模型预测性能，选择最有希望的配置进行实际训练。4) 模型更新：定期使用新的训练数据更新代理模型，以提高预测精度。

关键创新：CubicML的关键创新在于使用机器学习模型作为代理来预测大规模分布式机器学习系统的训练性能。与传统的基于规则或经验的调参方法相比，CubicML能够自动学习超参数与性能之间的复杂关系，从而更有效地优化训练过程。

关键设计：代理模型的选择至关重要，需要考虑模型的预测精度和计算效率。论文中可能采用了梯度提升树、神经网络等模型。超参数搜索算法的选择也影响搜索效率，贝叶斯优化等算法能够更好地平衡探索和利用。此外，代理模型的更新频率和数据量也需要仔细调整，以保证模型的预测精度和泛化能力。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

CubicML在Meta内部的广告推荐模型（730亿参数）和大型语言模型（高达4050亿参数）上进行了验证，结果表明CubicML能够有效优化训练速度。具体的性能提升数据（例如加速比、训练时间缩短百分比等）需要在论文中查找。

🎯 应用场景

CubicML可广泛应用于大规模机器学习系统的训练性能优化，例如推荐系统、自然语言处理、计算机视觉等领域。通过自动优化训练配置，可以显著缩短模型训练时间，降低计算成本，加速模型迭代和部署，从而提升业务效率和竞争力。该方法尤其适用于拥有大量超参数和复杂依赖关系的大规模模型。

📄 摘要（原文）

Scaling up deep learning models has been proven effective to improve intelligence of machine learning (ML) models, especially for industry recommendation models and large language models. The co-design of large distributed ML systems and algorithms (to maximize training performance) plays a pivotal role for its success. As it scales, the number of co-design hyper-parameters grows rapidly which brings challenges to feasibly find the optimal setup for system performance maximization. In this paper, we propose CubicML which uses ML to automatically optimize training performance of large distributed ML systems. In CubicML, we use an ML model as a proxy to predict the training performance for search efficiency and performance modeling flexibility. We proved that CubicML can effectively optimize training speed of in-house ads recommendation models with 73 billion parameters and large language models up to 405 billion parameters at Meta.