无论模型有多大或数据存储在哪里,IPU访问存储的独特方式、一流的处理器内存储(In-Processor Memory)设计以及Exchange Memory软件功能相结合,使得用户可以快速执行机器学习模型。
IPU存储访问设计
最初创建IPU时,我们花费了大量时间来考虑存储。这是IPU设计方式中最令人兴奋且具有前瞻性的方面之一。它采用了全新的方法来组织存储。首先,我们着重于拥有大量的处理器存储,即每个MK2 Colossus GC200 IPU都具有900MB的存储,这是前所未有的,足以在许多情况下容纳整个模型(尤其是在不浪费存储的情况下)。其次,我们确保IPU可以通过其独特的Exchange Memory通信技术访问其他存储源,以启用更大的模型和程序。该设计对于现代机器学习工作负载至关重要,即访问存储的方式与获取数据后执行计算的方式同样重要。
Exchange Memory访问设计围绕两个原则进行:
- 在硬件设计中尽可能少做存储访问行为的假设。
- 允许对存储访问进行完全、显式的软件控制。
这两个设计原则,既是针对IPU上高速处理器内存储的,也是针对访问芯片外部的流存储的。
流存储
像许多现代处理器一样,IPU可以处理分层的存储层次结构。由于IPU在多芯片计算中协同工作,因此最好的方法是IPU系统的存储规格,例如IPU-Machine: M2000 (IPU-M2000):

每个IPU-Machine具有可通过4个IPU进行寻址的450GB存储,分为每个IPU上900MB的处理器内存储和每个IPU上16GB的流存储。流存储被置于在IPU-Machine 1U服务器(与IPU位于同一服务器)上的DDR4 DIMM(双列直插式存储)模块中,确切的存储量将取决于这些DIMM的规模。
使用处理器内存储和流存储的软件支持
如前所述,IPU设计的指导原则之一是存储使用由软件高度指导。这使我们能够为具有最大灵活性的应用程序提供一个平台,以实现高效的存储访问。那么,什么软件功能能够实现这一目标呢?
Exchange Memory
Exchange Memory是我们用于Poplar SDK功能的总称,该功能对处理器内存储和流存储的使用进行管理。我们在Poplar SDK 1.2中提供这些功能,并将在2020年及以后继续扩展它们。这些功能多种多样,需要经过深思熟虑的设计。像任何存储层次结构一样,这些工具必须在快速本地存储(IPU具有比大多数处理器更多的本地存储)和流存储(提供大量空间但访问时间较慢)之间高效实现平衡。
Poplar®Graph编程框架中的
显式存储管理
在较低级的Poplar计算图编程框架中,所有存储空间都是显式的。可以在流存储中创建远程缓冲区,并且控制程序可以精确地控制何时将数据从这些缓冲区复制到处理器存储中的变量,反之亦然。
// Add a remote buffer in streaming memory
auto r = graph.addRemoteBuffer("MyRemoteBuffer", FLOAT, 128, 10);
// Add a 2D graph variable in in-processor memory
auto v = graph.addVariable(FLOAT, {8, 16});
...
...
// Copy the variable into the remote buffer
prog.add(Copy(v, r));
...
Poplar代码使用显式存储控制
机器学习框架中的智能变量放置
在机器学习框架(例如TensorFlow)级别,智能变量放置(intelligent variable placement)选择计算图中的变量何时处于流存储中以及何时处于处理器存储中。目前,关于哪些变量驻留在哪个存储中的决定是根据用户注释来实现的。展望未来,我们计划添加更多灵活的注释和自动化方法以简化开发。
在PopART™(Poplar Advanced RunTime)里训练神经网络模型时,程序可以将计算图的各个部分标记为“阶段”。在每个阶段之前,所有必需的模型参数(以及为神经网络训练的后向传递而存储的激活都将从流存储移入处理器内存储。在每个阶段之后,任何需要保存的参数或激活都被移入流存储。这意味着,在任意某一时间,仅一部分模型需要驻留在处理器内存储中。
# Add a matmul on a particular IPU (virtualGraph) for a particular
# ping-pong phase
with builder.virtualGraph(vgid), builder.pingPongPhase(i):
x = builder.aiOnnx.matmul([x, qkv])
PopArt代码展现了模型部分的注释
在上面的代码中,pingPongPhase注释标记了计算图部分,从而将其参数和激活存储在流存储中,并且将该阶段的执行方案设置为“乒乓”。在乒乓方案中,一个IPU将进行计算,而另一个IPU将从流存储中加载/保存数据。当该阶段结束时,这两个IPU将交换角色,于是现在可以对上一阶段中加载的数据进行计算,并且先前进行计算的IPU可以加载/保存更多数据。实际上,使用者有两个IPU来实施一个存储提取管道。乒乓只是使用流存储的框架中可用的执行方案之一。
该软件可以利用模型的分阶段执行以及重新计算来进一步减少存储访问时间。这只有在非常大的处理器存储空间中才有可能。在梯度下降训练算法的后向传递中,可以从阶段开始就重新计算所有激活,并将其存储在该大存储中。这意味着整个阶段里,只有一组激活需要存储并从流存储中重新加载。
当您将权重移入流存储时,可以在PopART™中启用的另一个优化是复制权重切分(replicated weight sharding)。此优化对要参与数据并行训练的多个IPU并行获取的权重进行分区。这个微软的项目证明了这项技术。通过在芯片之间使用快速的本机IPU-Links™,然后在数据并行副本之间共享权重,可以在IPU上进一步增强此功能。使用并行数据加载和IPU-Links™的组合将权重的加载/存储时间减少到总运行时间的很小一部分。
在TensorFlow模型中,缺省状态下,所有变量都驻留在处理器内存储中,而优化器状态变量(动量项等)缺省到流存储。此缺省可以更改。未来,我们将不断为TensorFlow添加控制变量放置的新选项。
return
pipelining_ops.pipeline(computational_stages=...,
pipeline_depth=...,
...
offload_weight_update_variables=True,
...
TensorFlow代码配置优化器状态变量以驻留在流存储中
优化的数据通信
流存储可用来优化将数据移入和移出本地IPU(到CPU或其他服务器中的其他IPU)所需的时间。流存储可以充当通信的中间缓冲区。这些优化将在未来被引入Poplar SDK。这些优化的一大好处是它们可以与您的应用程序自动配合工作,无需更改代码。
总结:用于智能计算的智能存储
良好的存储和数据管理是高性能机器智能工作负载的关键。Graphcore公司将之视为其技术最重要、最令人兴奋的功能之一。我们的方法的核心就是围绕软件/硬件协同设计,以全新的眼光来思考存储。这使我们能够在软件和新硬件平台上进行探索和创新,以最大程度地利用存储。