arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2006.09616cs.LGcs.PLstat.ML
收藏

动态张量再物质化

Dynamic Tensor Rematerialization

OctoML;School of Computer Science, Carnegie Mellon University(卡内基梅隆大学计算机科学学院)

Marisa Kirisame, Steven Lyubomirsky, Altan Haan, Jennifer Brennan, Mike He, Jared Roesch, Tianqi Chen, Zachary Tatlock

摘要:

检查点通过从内存中释放中间激活并按需重新计算它们,使深度学习模型能够在有限的内存预算下进行训练。当前的检查点技术静态地离线计划这些重新计算,并假设静态计算图。我们证明了一个简单的在线算法可以实现相当的性能,通过引入动态张量再物质化(DTR),一个贪婪的在线检查点算法,是可扩展的和一般的,参数化的驱逐政策,并支持动态模型。我们证明了DTR可以训练一个$N$层线性前馈网络的$Ω(\sqrt{N})$内存预算只有$\mathcal{O}(N)$张量操作。在模拟实验中,DTR与最优静态检查点的性能非常接近。我们将DTR原型整合到PyTorch中,只需插入张量分配和运算符调用,并收集张量上的轻量级元数据。

英文摘要:

Checkpointing enables the training of deep learning models under restricted memory budgets by freeing intermediate activations from memory and recomputing them on demand. Current checkpointing techniques statically plan these recomputations offline and assume static computation graphs. We demonstrate that a simple online algorithm can achieve comparable performance by introducing Dynamic Tensor Rematerialization (DTR), a greedy online algorithm for checkpointing that is extensible and general, is parameterized by eviction policy, and supports dynamic models. We prove that DTR can train an $N$-layer linear feedforward network on an $Ω(\sqrt{N})$ memory budget with only $\mathcal{O}(N)$ tensor operations. DTR closely matches the performance of optimal static checkpointing in simulated experiments. We incorporate a DTR prototype into PyTorch merely by interposing on tensor allocations and operator calls and collecting lightweight metadata on tensors.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑