TensorFlow Large Model Support：TensorFlow中用于大模型支持（lms）的图形编辑库

参考：https://github.com/IBM/tensorflow-large-model-support

该库提供了一种训练大模型的方法，这些模型往往是GPU显存不足，它采用用户自定义的图，并自动添加swap-in和swap-out节点，用于将tensors从GPUs传输到主机设备，反之亦然。计算图是静态修改的。因此，它需要在回话开始前就完成。

IBM PowerAI 1.6包含比github更新的TensorFlow大模型支持的实现。新版本更容易去使用并且能够达到更高水平的tensor交换，从而可以实现更高分辨率，更深的模型，或者更大的批次。详细信息参见https://developer.ibm.com/linuxonpower/2019/05/17/whats-new-in-powerai-1-6-tensorflow-large-model-support/

install

git clone https://github.com/IBM/tensorflow-large-model-support.git
pip install ./tensorflow-large-model-support

如何使用

TensorFlow Large Model Support 需要知道关于用户定义模型的一些信息。用户定义的模型有一个要求：他必须具有optimizers/solvers。

嵌入LMS到模型中依赖于用户如何编写训练方式，以下指南有三种训练方法：

• 基于session的训练
• 基于estimator的训练
• 基于tf.keras的训练

基于session的训练

1、定义optimizer/solver范围

with tf.name_scope('adam_optimizer'):
	train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy)

2、定义LMS对象并运行它

from tensorflow_large_model_support import LMS
lms_obj = LMS({'adam_optimizer'})
lms_obj.run(graph=tf.get_default_graph())

基于Estimator训练

1、定义optimizer/solver范围

with tf.name_scope('adam_optimizer'):
      optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
      train_op = optimizer.minimize(
        loss=loss,
	global_step=tf.train.get_global_step())

2、定义一个LMSSessionRunHook

# Hook for Large Model Support
from tensorflow_large_model_support import LMSSessionRunHook
# LMSSessionRunHook and LMS share the same set of parameters. Here we just
# use the default keyword arguments.
lms_hook = LMSSessionRunHook({'adam_optimizer'})

3、增加LMSSessionRunHook到Estimator的hook 列表

mnist_classifier.train(
      input_fn=train_input_fn,
      steps=20000
      hooks=[logging_hook, lms_hook])

基于tf.keras的训练

使用LMSKerasCallback需要TensorFlow1.12或以上版本。

1、定义LMSKerasCallback

from tensorflow_large_model_support import LMSKerasCallback
# LMSKerasCallback and LMS share a set of keyword arguments. Here we just
# use the default options.
lms_callback = LMSKerasCallback()

2、将回调函数传递给Keras fit或fit_generator函数

model.fit_generator(generator=training_gen, callbacks=[lms_callback])

TensorFlow Grappler和TensorFlow Large Model Support

从TensorFlow1.14开始，TensorFlow的Grappler中依赖的优化器删除了TensorFlow Large Model Support的交换节点，为了避免这种情况，优化器必须禁用，以下是使用Keras模型执行此操作的例子：

import tensorflow as tf
from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.python.keras import backend as K
config = tf.ConfigProto()
config.graph_options.rewrite_options.dependency_optimization = rewriter_config_pb2.RewriterConfig.OFF
K.set_session(tf.Session(config=config))

TensorFlow有内存优化机制，虽然该机制能够与该模块共存，但依旧推荐将其模式切换为SCHEDULING_HEURISTICS以允许尽可能大的训练模型，这可以通过以下代码段完成：

config = tf.ConfigProto()
config.graph_options.rewrite_options.memory_optimization = \
	rewriter_config_pb2.RewriterConfig.SCHEDULING_HEURISTICS

Scaling tips

按照TensorFlow文档描述，通过以multi-tower方式构建模型并使用每个GPU的tower来实现扩展到多个GPU，还必须设置和调整TF_CUDA_HOST_MEM_LIMIT_IN_MB环境变量（在TensorFlow1.14中TF_CUDA_HOST_MEM_LIMIT_IN_MB）。

TensorFlow在CUDA主机（CPU）端分配内存的设置限制。该限制通常不足以充当多个GPU的张量交换空间，如果该设置不高，将导致此类内存不足错误：Allocator(cuda_host_bfc) ran out of memeory trying to allocate.注意提到的cuda_host_bfc allocator分配器而不是GPU alllocator分配器。注意在TensorFlow1.14中主机端allocator分类器名称已更改为gpu_host_bfc。

一个好的经验法，则是从GPU内存容量的4倍乘以将要使用的GPUs的数量的值开始。例如，如果你与16GB GPUs在系统中，将在训练中使用所有的四个GPUs。

TF_CUDA_HOST_MEM_LIMIT_IN_MB (TF_GPU_HOST_MEM_LIMIT_IN_MB in TensorFlow 1.14)应该被设置成262144并根据需要调整 (4 x 16384 (16GB as MB) x 4 GPUs) = 262144 MB.

对于LMS/LMSSessionRunHook/LMSKerasCallback参数化的参数

必要的参数

graph::将要为LMS修改的图，这应该是用户自定义的神经网络。（LMSSessionRunHook 或LMSKerasCallback不需要）

optimizer_scopes::optimizers/solvers作用域（LMSKerasCallback不需要）

可选参数

starting_scope:: 将此范围内的tensor操作转换成LMS，如果想要修改整个图的话，请在第一层就设置该作用域，默认是None。

starting_op_names::将此范围内的tensor名字转换成LMS，默认是None

excl_scopes::一些列操作范围，这些tensor不会被交换到主机，默认是empty

incl_scopes::一些列操作范围，这些tensor会被交换到主机，默认是empty

excl_types::一些列操作类型，这些tensor不会被交换到主机，默认是empty

incl_types::一些列操作类型，这些tensor会被交换到主机，默认是empty

n_tensors::从starting_scope开始计算的，LMS的张量数。要关闭LMS，设置n_tensors=0，默认是-1（所有可到达的tensors将被转换成LMS）

lb::用于LMS的Lowerbound下界值，在向后阶段，张量将在图中的至少lb个节点之间交换，默认是1

ub::用于LMS的Upperbound上界值，默认是10000

fuse_swapins::将close交换惭怍到一个操作中。这个也许会提升性能，默认是False。

ctrld_strategy::两种用于找到swapin操作的控制以来操作的两种策略:chain_rule和direct_order。chain_rule策略从正向操作开始，前进并找到相应的后向操作作为控制依赖操纵。direct_oder策略直接获取拓扑顺序中的向后才做作为控制依赖操作。两种策略都依赖于lb和ub来选择控制依赖型操作。虽然direct_oder相对于lb和ub比chain_rule更精确，但实验上通常会导致最大批量大小小于chain_rule，默认是chain_rule。

swap_branches::如果是true，LMS将在前置阶段交换分支中的张量。默认False。

branch_threshold::如果swap_branches启用了，并且在consuming操作和张量generating操作之间拓扑顺序距离大于branch_threshold，则交换张量，默认是0

debug::LMS的Debug模式，默认是false

debug_level::LMS的Debug等级（1或2），默认是1

性能调优LMS

一旦在代码中启用LMS图形修改，您将需要找到调整参数的组合，以便为您的模型提供更快的培训时间和最精确的效果。性能调整的目标是交换足够多的张量以允许您的训练运行而不会出现内存错误，同时不会交换太多，以至于额外的交换通信开销降低性能。

您应该关注的两个调整参数是n_tensors和lb。由于n_tensors控制将要交换的张量数量，因此设置的越高，GPU内存使用峰值就越低。lb控制张量在使用前重新换回的时间。较低的lb值可以使GPU上的训练暂停，并等待交换完成。这将会降低性能。较高的lb值允许张量交换在需要之前完成，并允许训练无间断的运行。过载交换的缺点是，在任何时间点GPU内存中都会有更多的张量，从而导致更高的GPU内存使用率。

因此，调整n_tensors和lb之间找到正确的平衡，从而为给定模型提供最佳性能。要开始性能调整，建议将n_tensors设置成-1，这将交换所有能够到达的张量。lb应该设置成默认值1，这是最低可能交换的值。如果tf.logging verbosity设置成tf.logging.INFO，则LMS将输出一个日志语句，其中包含已交换的张量数量的计算。在运行时设置n_tensors=-1，首先运行找到最大值然后向下调整。如果你的模型具有向某些UNet模型那样的分支，您可能希望设置swap_branches=Ture并调整分支阈值。在调整参数时，将训练脚本中的LMS参数表示微命令换藏书或者配置文件属性通常很有用。

默认情况下，LMS将分析您的图形以查找用于用于搜索张量交换候选的起始操作。您可以通过将明明范围内的起始操作并在starting_scope参数上提供范围，或在starting_op_names参数上提供启动操作的名称来绕过此分析。这可以酵素调整期间LMS的重复运行。进一步您可以启用debug=Ture和debug_level=1，LMS将打印出他找到的启动操作名称和类型。。这些名称可以在后续运行时传递给starting_op_names参数。