MxNet系列——how_to——model_parallel_lstm

博客新址: http://blog.xuezhisd.top
邮箱：xuezhisd@126.com

---

使用模型并行的方式在多个GPUs上训练LSTM

由于复杂的数据依赖，LSTM评价很困难。LSTM的训练过程，在反向传播阶段有更严重的逆序的数据依赖，更加难以并行化。获取有关LSTM更通用信息，请查看优秀的 简介 （作者：Christopher）有关使用模型并行的方式在多个GPUs上训练LSTM的资料，请查看 example/model-parallelism-lstm/。

模型并行：将多个GPUs作为流水线

最近，在应用机器学习中，有很多关于模型并行的热烈讨论。它最初是为GoogleNet中的超大卷积层设计的。我们借鉴了一个网络层位于一个GPU上的想法。模型并行的基元是神经网络模型的网络层。它带来的好处是GPU没有必要将所有网络层的参数保存在内存中了。这降低了大规模任务对内存要求，例如，机器翻译。

在上图中，不同的LSTM模型位于不同的GPUs上。GPU 1完成第一层的（第一个句子的）计算之后，输出结果传递给GPU 2。同时，GPU 1 获取下一个句子并开始训练。这和数据并行明显不同，因为在每一次迭代的最后，没有更新共享模型的冲突。绝大多数的通信发生在不同GPUs之间的流水线式（先后顺序）的传递中间结果。

在这个实现中，网络层定义在 lstm_unroll()。

划分工作负载

为了在不同GPUs上划分网络，实现模型并行需要有关训练任务的丰富知识。尽管它需要详细的分析（已经超出了该节内容的范围），你可以应用一些通用的原则：

• 为了避免数据传输，将相邻层定义在同一个GPU上。
• 为了避免流水线的瓶颈，需要均衡GPUs之间的负载。
• 记住不同的网络层有不同的计算-内存属性。

浏览以下上图中的两个流水线。他们都有八个网络层，一个解码层和一个编码层。根据第一个原则，将所有相邻层放在不同GPUs上是不明智的。我们也想均衡不同GPUs上的负载。尽管LSTM层比编码/解码层消耗更少的内存，但因为展开LSTM的依赖性，它们会花费更多的计算时间。因此，左边的流水线的负载更加均衡，因此比右边的更快。

在这个实现中，网络层在 lstm.py中划分和在 lstm_ptb.py中使用 group2ctx 选项配置。

在模型并行中使用 Bucketing

为了在模型并行中使用 bucketing，你需要为每一个bucket把LSTM展开，以为每一？生成一个执行器。关于模型时如何约束的细节，请查看 lstm.py。

另一方面，因为模型并行划分模型/网络层，输入数据必须变换/转置成一致的形状。更多细节，请查看 bucket_io。