TM-Training：一种用于神经处理单元（NPU）中深度学习训练的高效分层内存系统

《ACM Transactions on Storage》：TM-Training: An Energy-Efficient Tiered Memory System for Deep Learning Training in NPUs

DNN加速系统中DRAM成本占比高，需通过分层内存系统实现能效与成本平衡。TM-Training采用HBM与闪存双层架构，通过静态预测数据分布并动态调整分片比例，减少 tiers间数据迁移达64%，在12.5%容量配置下能效较基线提升55%，同时降低87%硬件成本。研究重点包括工作负载重复性带来的静态预测可行性、基于时序监控器的数据分片算法、以及对比验证各策略的能效与吞吐表现。分层内存设计需考虑硬件兼容性，TM-Training通过优化数据分片和预取策略，在保持高吞吐的同时显著降低能耗，其核心创新在于结合工作负载周期性特征实现静态数据放置，而非依赖在线替换策略。

　　在当今的深度学习加速系统中，动态随机存取存储器（DRAM）占据着系统总拥有成本（TCO）的较大比例。随着深度学习模型的规模不断增大，对内存容量的需求也随之上升，这促使了对分层内存系统的探索。分层内存系统通过使用不同性能和密度的存储技术来实现内存容量和性能的平衡，其中高层通常配备高性能但成本较高的技术，而低层则采用高密度但性能较低的技术。然而，无论分层内存系统的设计如何，数据在各层之间的迁移仍然消耗大量能源。特别是在模型规模和内存需求持续增长的情况下，这种数据迁移变得更为频繁，从而对能源效率构成了挑战。

为了解决这一问题，研究者们提出了一些创新性的方案，其中一种思路是通过预测性地将数据放置到合适的层级，并在需要时及时从该层级获取所需数据，从而减少不必要的数据迁移。这一策略的核心在于识别和利用深度学习训练任务中数据访问模式的重复性。例如，变压器模型（Transformer）在执行过程中会反复使用相同维度的张量，这种特性使得静态预测成为可能。基于此，我们提出了一个名为TM-Training的新型深度学习加速系统，它利用闪存（Flash Memory）构建分层内存结构，旨在通过优化数据放置策略，显著提升能源效率。

TM-Training的设计灵感来源于对深度学习训练任务中数据访问模式的深入分析。该系统基于水平分层内存架构，将数据划分为多个部分，并分别存储在不同层级中。在系统运行时，TM-Training会根据任务的执行情况，动态调整数据在各层级中的分布。这种数据分割方法允许系统在不牺牲性能的前提下，尽可能减少数据迁移次数，从而降低整体能耗。同时，TM-Training引入了一种新的数据放置管理机制，能够在不显著增加硬件或软件开销的情况下，实现高效的数据迁移管理。

TM-Training的关键组件包括请求时间监控器（Request Timing Monitor）和数据放置管理器（Data Placement Manager）。请求时间监控器负责记录数据在系统中不同层级间的访问时间，包括数据从分层内存到片上SRAM缓冲区的迁移时间、计算单元（如MXU）对数据的使用时间以及操作完成的时间。这些时间信息对于数据放置管理器来说至关重要，因为它可以根据这些信息计算出最优的数据分割比例，并据此调整数据在不同层级中的存储策略。数据放置管理器利用简单的算法，结合内存带宽、数据量等信息，以实现高效的数据分配。

在实验评估中，TM-Training展示了显著的性能提升。与传统仅使用HBM的系统相比，TM-Training在训练任务中实现了每瓦特55%的吞吐量提升。此外，TM-Training还能够将分层内存系统中数据迁移的能耗降低64%。这一结果表明，通过合理利用数据访问模式的重复性，TM-Training能够有效减少数据迁移，从而提升系统的整体能效。TM-Training的这些优势使其在应对大规模模型训练任务时，能够提供更高的性价比。

为了进一步优化TM-Training的性能，研究者们还分析了不同层级的内存带宽对系统能耗的影响。例如，使用NAND闪存作为分层内存的第二层，由于其相对较低的带宽，可能导致更多的数据迁移和更高的能耗。因此，TM-Training通过数据分割技术，使得数据在分层内存中分布得更加均匀，从而减少对高带宽内存的需求。同时，TM-Training还考虑了内存访问模式的预测性，通过提前迁移和存储关键数据，以减少实际执行过程中的数据迁移需求。

此外，TM-Training的实现还考虑了硬件和软件层面的开销。在硬件方面，系统通过设计请求时间监控器和在专用IOMMU中添加额外的硬件组件，以支持数据迁移和存储的高效管理。而在软件层面，TM-Training的算法能够基于简单的逻辑进行数据分割，从而减少对复杂软件功能的依赖。这种设计使得TM-Training在保持高性能的同时，也能够显著降低能耗。

通过这些设计和优化，TM-Training不仅在能效方面表现出色，还能够在大规模模型训练任务中提供更高的吞吐量。例如，在使用不同层级的内存系统时，TM-Training能够有效减少数据迁移，从而提高系统的整体性能。这一成果对于当前深度学习加速系统的设计具有重要的参考价值，尤其是在应对日益增长的模型规模和内存需求时，TM-Training的分层内存策略能够提供更优的解决方案。

总的来说，TM-Training通过利用深度学习训练任务中数据访问模式的重复性，设计了一种高效的分层内存管理策略。这种策略不仅能够显著减少数据迁移带来的能耗，还能够通过优化数据放置，提升系统的整体性能。在实验评估中，TM-Training展示了其在能效和吞吐量方面的优越性，这为未来深度学习加速系统的优化提供了新的思路和方法。