中文

TensorBank:面向基础模型训练的张量湖仓

机器学习 2024-03-22 v3 人工智能 数据库 信息检索

摘要

随着超越自然语言的基础模型兴起,为基础模型训练存储与流式传输高维数据成为一项关键需求。在本文中,我们介绍TensorBank,一个 petabyte 规模的张量湖仓,能够基于复杂关系查询以线速将张量从云对象存储(COS)流式传输至GPU内存。我们使用分层统计索引(HSI)进行查询加速。我们的架构允许使用HTTP范围读取在块级别直接寻址张量。一旦进入GPU内存,数据可使用PyTorch变换进行转换。我们提供一种通用PyTorch数据集类型及相应的数据集工厂,将关系查询与所请求的变换实例化为一个对象。通过利用HSI,无关块可在不被读取的情况下跳过,因为这些索引在不同分层分辨率级别上包含其内容统计。这是一种由开放标准驱动并重度使用开源技术的定制化架构。尽管已使用地理时空数据强化至生产可用,该架构可推广至其他用例,如计算机视觉、计算神经科学、生物序列分析等。

关键词

引用

@article{arxiv.2309.02094,
  title  = {TensorBank: Tensor Lakehouse for Foundation Model Training},
  author = {Romeo Kienzler and Leonardo Pondian Tizzei and Benedikt Blumenstiel and Zoltan Arnold Nagy and S. Karthik Mukkavilli and Johannes Schmude and Marcus Freitag and Michael Behrendt and Daniel Salles Civitarese and Naomi Simumba and Daiki Kimura and Hendrik Hamann},
  journal= {arXiv preprint arXiv:2309.02094},
  year   = {2024}
}