NNTile:用于单节点训练超大规模 GPT 类语言模型的机器学习框架
机器学习
2025-04-21 v1 数学软件
摘要
本研究提出了 NNTile 框架,用于在异构集群上训练大型深度神经网络。NNTile 基于 StarPU 库实现,采用任务级并行机制并将所有提供的任务调度到所有可用处理单元(CPU 和 GPU)上。这意味着,训练大型神经网络所必需的特定操作可在 CPU 核心或 GPU 设备上执行,取决于自动调度决策。这种方法将决定计算位置和通信时机的负担从人类转移到自动决策者(无论是简单的贪心启发式方法还是复杂的基于 AI 的软件),从而实现高效的资源调度。通过大量数值实验,展示了该工具在训练大型语言模型方面的性能。
引用
@article{arxiv.2504.13236,
title = {NNTile: a machine learning framework capable of training extremely large GPT language models on a single node},
author = {Aleksandr Mikhalev and Aleksandr Katrutsa and Konstantin Sozykin and Ivan Oseledets},
journal= {arXiv preprint arXiv:2504.13236},
year = {2025}
}