Hardware Co-Design Scaling Laws via Roofline Modelling for On-Device LLMs

• AI Lab, The Yangtze River Delta Institution of Automation, Chinese Academy of Sciences
• University of Chinese Academy of Sciences
• Li Auto
• University College London
• The University of Edinburgh

• 提出一个硬件协同设计缩放定律，用于在资源受限的设备端大型语言模型（LLM） 中联合优化模型精度与推理延迟。
• 研究目标是通过结合损失缩放定律和基于屋顶线模型的延迟建模，明确刻画在固定硬件约束下精度与延迟的帕累托前沿，从而系统化地指导模型架构设计。
• 解决的问题是：在设备端部署场景（如自动驾驶汽车、移动机器人）中，如何选择最优的LLM架构，使其在严格的推理延迟和硬件效率约束下，达到最佳的性能-精度权衡。

• 提出首个实用的、可操作的硬件协同设计缩放定律框架，用于设备端LLM部署，首次通过屋顶线模型联合捕捉模型精度与推理延迟。
• 进行了大规模架构评估和训练实验：在NVIDIA Jetson Orin上评测了1942个候选架构，并训练了170个选定模型（每个10B token），以拟合将架构与训练损失关联的缩放定律。
• 将经验框架扩展为一个原则性的理论表述，将架构搜索表述为在精度和性能上的联合优化问题，推导出工业硬件和应用预算下的可行设计区域。
• 显著缩短架构选择周期：将传统需要数月的架构选择过程缩短至数天，并计划开源代码、检查点和评估协议以支持社区发展。

• 核心方法是硬件协同设计定律：使用架构超参数的函数显式建模训练损失（精度），并通过屋顶线分析建模推理延迟（性能）。
• 损失建模：通过拟合170个模型的训练结果，得到一个经验性缩放定律（如公式2），将验证损失近似为深度、宽度、MoE激活率、FFN扩展比、KV-cache维度等参数的函数。
• 延迟建模：基于屋顶线模型，为Transformer的每个算子（如注意力、FFN）进行分析，将端到端延迟表达为计算量、内存访问量和硬件特性（峰值算力、内存带宽）的函数（如公式3,4）。
• 架构搜索：将上述损失和延迟模型结合，通过帕累托前沿分析，在给定的硬件约束（如目标延迟）下识别最优的架构配置。

• 训练数据：使用了包含通用语料、数学推理和代码数据的混合数据集，总训练量为10B tokens。
• 模型规模：搜索的架构覆盖了从稀疏（MoE）到密集的多种Transformer配置，参数量级在0.5B附近（如对比的Qwen2.5-0.5B）。
• 训练资源：未在提供的HTML片段中明确说明，但附录提及了训练细节（Appendix B Pre-training Details），通常此类实验涉及大量GPU/TPU资源。

• 评估环境与基准：主要在NVIDIA Jetson Orin硬件平台上进行评估，并与其他设备端模型如Qwen2.5-0.5B进行对比。
• 主要评估指标：

- 推理延迟：单位毫秒（ms）。

- 模型精度：在WikiText-2数据集上使用困惑度（Perplexity） 衡量。

• 关键实验结果：

- 在目标硬件上，与Qwen2.5-0.5B达到相同延迟时，论文协同设计的架构在WikiText-2上取得了19.42% 更低的困惑度。

- 通过构建损失-延迟帕累托前沿，有效识别了在给定延迟约束下能提供最佳精度的架构配置。