TF中Placement的最后一道防线——Placer

二环宇少

发布于 2020-08-25 11:14:24

1.2K0

发布于 2020-08-25 11:14:24

问题引入

在使用TensorFlow构建模型时，为了能够使用GPU的Device，你可能会用到下面的这样的写法。

with tf.device('/gpu:0'):
  a = tf.get_variable(.....)
  b = .......
  c = .......

那么，上面代码中的a、b和c就真的一定会放在GPU:0上吗？如果c不存在GPU上的实现会怎么样？进一步地，有没有其他约束会让用户的设置失效？

事实上，当你打开session config的log_device_placement选项后，仔细逐个检查每个Op被放置的位置，你会发现某些Op并没有如你所愿被你控制，而是被“悄悄地”放到别的Device上了。

这并不是Bug，而是Placer模块发挥了保护作用。Placer是TensorFlow中Placement设置的最后一道防线。它工作在TF底层，在尽可能满足用户诉求的前提下，暗中纠正部分不合理的Placement。

且听我从设计初衷与源码上，为你娓娓道来~

Placement设计初衷

受限于单个Device的计算能力和存储大小，模型分片是重要的需求点之一。它的本质是将模型和相关的计算切分到不同的Device，如此不但可以解决单个Device放不下大模型的问题，还有可能带来计算加速的收益。

在深度学习框架方面，显然在TensorFlow上做模型分片比Caffe更加容易，这主要得益于TensorFlow的Placement机制。Placement是TensorFlow引入的特有概念，它指定某个Op与具体Device的绑定关系，因此模型分片问题实际上就是该模型上每个Op的Placement问题。

在Python层面，一共存在两个API与Placement相关的接口，它们不但广泛存在与框架代码中，还可以被用户拿来直接使用。

但是用户指定Placement信息存在一定的不可靠性，它与Op的实际情况往往存在一定的矛盾，这就需要TensorFlow中的Placer模块来解决。

Placer功能描述

Python构完图之后，请你把GraphDef打印出来，我们要关注每一个Node的NodeDef结构（如下图），这里有两个地方和Placement相关。

device属性：它显示指定了这个Node应该被放在何种Device上，它由用户通过with tf.device指定。
字符串标记loc:@xxxx：这是placement的约束条件，隐式指明该Node的Placement应该和哪些Node保持一致。xxxx代表某个Group的名字，该Node应该和Group名为xxxx内的所有Node的Placement保持一致。

可以想象，以上两个信息可能会出现矛盾的情形。

Placer不但要处理二者的矛盾，还要通过一些规则尽可能避免因Placement不当带来的性能问题。每个Node在经过Placer处理后都会得到最终的Placement信息，它将重新覆盖NodeDef中的device属性内容。

所以，通俗地讲，Placer的功能就是推断并填入所有NodeDef的device属性。

一些前驱内容

梳理逻辑时难免会碰到一些为解决这个问题专门设立的名词和经典的算法，所以建议在阅读Placer模块相关内容之前先确认已经弄清楚下面的东西，避免走一些弯路。

显式Placement：指用户通过with tf.device直接指定的Placement信息，它将写入上一小节中NodeDef中的device属性。
隐式Placement：指间接指定的Placement信息，这个信息与上一小节中NodeDef中的loc:@xxxx对应。上一节说到，xxxx是一个Group的名字，该Group内所有的Node都要求具有相同的Placement信息，这个Group被叫做Colocation Group，属于一种约束（Constraint）条件。
Find-Union算法：并查集算法，Placer内最重要的算法。TensorFlow通过Find-Union算法高效地处理了Node的Colocation问题。简单而言，逻辑上，多个具有相同Colocation Group的Node应该被“并”到同一个组中，从而“查”某个Node的Placement信息时，可以更快速地获取整组的信息。在实现时，如何设计更好的数据结构，并高效地实施“并”和“查”两个过程，是并查集算法的核心。

Placer决策基本原则

Placer会根据会对Graph进行一定程度的分析，并结合用户的要求对每个Node的Placement进行微调，微调的原则可以概括为下面四点

尽可能满足用户要求（User Requirement First）：每个Node的placement会尽量满足用户的要求
尽可能使用计算更快的设备（High Performance Device）：若某个Node的Placement没有被用户指定，则优先分配计算更快的设备
保证程序可运行（Runable）：若某个Node不存在用户要求的Placement相关实现版本，会退而求其次选择其它实现版本，保障程序可以用
尽可能考虑近邻特性（Near When Possible）：在做Placement的微调时考虑节点的近邻特性，尽可能减少无意义的拷贝

原则原理详细展开

1. 尽可能满足用户要求（User Requirement First）

用户要求分为两种，一种是显示指定，表现为在Node中设置的device信息；另一种是隐式指定，表现为loc:@xxxx属性，即Colocation Group。

Placer会根据用户这两方面的要求并结合实际情况做Placement信息补全和微调。

文章开头的截图展示了某个Node的NodeDef信息，它表明类型为MatMul的Op被用户显示指定放到'/device:GPU:0'上，同时希望放入名为global_step的Colocation Group中。

NodeDef中的device属性和loc:@xxxx属性分别由下面两个python级别的API引入，它们都由用户来控制，有些被用在高层API内部封装中。

# device attributes
@tf_export("device")
def device(device_name_or_function):

# colocation attributes
@tf_export("colocate_with")
def colocate_with(op, ignore_existing=False):

2. 尽可能使用更快的计算设备（High Performance Device）

如果某个Node的device属性中不含device_type（即GPU或CPU），那么Placer必须决定使用何种Device。每种Device注册到TensorFlow中时都带有优先级，通常高优先级的Device具有更好的计算性能。

当某个Op具有多种Device实现时，Placer将选取优先级最高的Device实现版本，通过设置device_type为所有实现版本中最高优先级的Device来实现这种选取。

3. 保证程序可运行（Runable）

这是通过Soft Placement机制保证的（在session config里可以设置）。

如果某个Node被显示指定精确放在某Device上，但系统中却没有该Device上的实现版本，那么为了保证程序可用，Soft Placement将发挥作用，它将忽略device type，在系统中按照Device优先级选取另一个可用的实现版本重新改写Placement。

举例而言，假设某Node的op是SparseToDense，device_type被指定为GPU，但目前SparseToDense在TensorFlow中只有CPU的实现，那么Soft Placement将改写该Node的device_type为CPU。

4. 尽可能考虑近邻特性（Near When Possible）

这块就比较复杂了，但我们要抓住重点，你就不会乱：关注三类特殊的Op类型，他们的特殊性，决定了其近邻是需要特殊处理的，分别是：

Generator类Op：入度为0，出度为1的Op
MetaData类Op：直接在Tensor的元数据MetaData上操作，不改变Tensor本身的内容，比如Reshape）
Ref类或Resource类：例如Variable这种可能发生赋值的Op（或者叫左值）

在Placer中使用以下三种启发式规则来分别应对上面三种特殊的Op。

若某个Node是GeneratorNode，将其与Consumer与其放在同一个Device上可以防止无意义的跨Device拷贝。这一步在算法中被称之为启发式规则A;
若某个Node是MetaDataNode，将其与Producer放在相同的Device上也可以防止无意义的跨Device拷贝。这一步在算法中被称为启发式规则B;

若某个Node的输入是Reference type或者是Reource type，那么尽量将其与输入放在同一个Colocation Group中（比如Variable，对其assign等操作肯定直接在Variable所在之地执行即可，如果Variable在A处，对其的assign在B处，显然是不合理的）。算法中没有为这个步骤起名字，为了方便我们称之为启发式规则C。

Placer决策总体流程

总体流程分为四个步骤，下图展示了宏观层面的流程图。其中最后两个步骤相对较为复杂，下一节中将会细化其流程图。

Placer分布详解与关键代码

注意！本节看源码的时候，要注重结构，而不是每个细节都去纠缠。

第一步——根据外部指定Colocation聚合Group

一般情况下，没有被用户指定Colocation Group信息的Node会被单独放入一个Group中作为唯一的成员，并以该Node的Name作为Group的名字，所以Graph中每个Node都会有自己的Colocation Group。

从逻辑上来说，合并多个Group是非常简单的问题，但是这个场景中的Group不仅是Node的集合，还包含若干属性，比如某个Group的possible device表示这个Group可用的所有Device集合。

因此我们需要一种数据结构和算法，帮助我们在合并两个Group时很方便地生成新Group及相关属性（方便Union），并且能够根据某个Node快速查看所属Group的所有属性（快速Find），这就是Find-Union的优势所在。

Find-Union算法原理将不在这里描述，这里只给出代码中Find-Union用到的基本数据结构——Member，它用来描述Group的基本信息。在阅读下段代码注释前，需要对Find-Union中的树形结构含义有基本的理解。

// Represents a node in the disjoint node set forest, and the
  // accumulated constraints on the device used by that node.
  struct Member {
    Member() = default;
    // The id of the node that is the parent of this one, or its own
    // id if it is a root. parent <= 0 indicates that this member is invalid.
    int parent = -1;

    // A proxy for the depth of the tree that is used to prefer
    // connecting smaller trees to larger trees when merging disjoint
    // sets.
    int rank = 0;

    // The intersection of all device types supported by this node,
    // and those of all of its children, in priority order
    // of the preferred device.
    DeviceTypeVector supported_device_types;

    // The merged form of the device requested for this node, with
    // those of all of its children.
    DeviceNameUtils::ParsedName device_name;

    // If this node is a root, stores a list of Devices to which this node
    // and all of its children have been assigned, or nullptr if this
    // has not yet been computed.
    std::vector<Device*> possible_devices;
  };

下面的代码是处理这一步骤的核心代码。首先创建ColocationGraph对象，这是一个处理Colocation Group的工具类，里面使用了Find-Union算法对Group进行聚合。

在调用InitiailizeMembers对Find-Union算法的基本数据结构进行初始化之后，就直接调用ColocationAllNodes根据用户指定的所有colocation信息进行聚合。

ColocationGraph colocation_graph(
      graph_, devices_,
      options_ == nullptr || options_->config.allow_soft_placement(),
      default_device_);

  TF_RETURN_IF_ERROR(colocation_graph.InitializeMembers());

  // 1. First add all of the nodes. Note that steps (1) and (2)
  // requires two passes over the nodes because the graph (and hence
  // the constraints) may not be acyclic.
  TF_RETURN_IF_ERROR(colocation_graph.ColocateAllNodes());

第二步——应用启发式规则C（处理Ref类Op Placement）

这一步将对Colocation Group进行调整。在遍历Graph的每个Node时，需要根据Node input来决定是否将该Node所在的Group与Source Node所在的Group合并。

如果Node的input是Reference type或者DT_RESOURCE（关于DT_RESOURCE一般会在使用ResourceVariable时才会碰到。ResourceVariable与Variable相比具有很多新特性，这些特性是TF2.0中主推的内容。关于它的优势我们不在这里展开，只对其Op的类型做一个说明。

Variable在C++层面的Op类型是VariableV2，而ResourceVariable在C++层面的Op类型为VarHandleOp。后者产生的Tensor就是一种DT_RESOURCE），那么就尝试做合并。在合并之前需要做必要的可行性检查，适当地主动报错。比如在合并时除了要考虑这一对节点的连接以外，还需要考虑这个Node的其他输入是否属于Reference type或者DT_RESOURCE。这一部分的代码比较长，但逻辑比较简单，这里不再展示。

第三步——应用启发式规则B（处理MetaData类的Op Placement）

从这一步开始，Placer才开始真正的为每个Node分配Device，下面的流程图中展示了这一步骤。

如果当前的Node的device属性中已经有值，那么Placer将不再对其做重复的assign操作，直接跳过这个Node;
如果当前Node是GeneratorNode，先将其放入一个名为second_pass的vector中;
如果不是以上两种情况，那么该Node正是这一步骤需要处理的对象。先从该Node所在的Colocation Group中获取可用的Devices（获取会受到Soft Placement的影响）作为候选。如果该node是MetaData node，那么会尝试应用启发式规则B，否则，将分配候选集中优先级最高的Device。

    int assigned_device = -1;

    // Heuristic B: If the node only operates on metadata, not data,
    // then it is desirable to place that metadata node with its
    // input.
    if (IsMetadata(node)) {
      // Make sure that the input device type is in the list of supported
      // device types for this node.
      const Node* input = (*node->in_edges().begin())->src();
      // TODO(vrv): if the input is empty, consider postponing this
      // node's assignment to the second pass, so that we handle the
      // case where a metadata node's input comes from a backedge
      // of a loop.
      if (CanAssignToDevice(input->assigned_device_name(), *devices)) {
        assigned_device = input->assigned_device_name_index();
      }
    }

    // Provide the default, if necessary.
    if (assigned_device == -1) {
      assigned_device = graph_->InternDeviceName((*devices)[0]->name());
    }

    AssignAndLog(assigned_device, node);

第四步——应用启发式规则A（处理Generator类的Op Placement）

这一步将对second_pass数组中的所有的Node分配Device，下面的流程图中展示了这一步骤。

放在second_pass中的代码全部是GeneratorNode，所以只需要应用启发式规则A即可，和步骤3一样，启发式规则A的应用也是尝试性的，如果实在不能满足，会直接分配候选Device中优先级最高的Device，下面是启发式规则A的应用部分代码。

    int assigned_device = -1;

    // Heuristic A application.
    if (IsGeneratorNode(node)) {
      const Node* output = (*node->out_edges().begin())->dst();
      int output_device_name = output->assigned_device_name_index();

      const bool consumers_on_same_device = std::all_of(
          node->out_edges().begin(), node->out_edges().end(),
          [output_device_name](const Edge* e) {
            return e->dst()->assigned_device_name_index() == output_device_name;
          });

      if (consumers_on_same_device &&
          CanAssignToDevice(output->assigned_device_name(), *devices)) {
        assigned_device = output_device_name;
      }
    }

    // Provide the default, if necessary.
    if (assigned_device == -1) {
      assigned_device = graph_->InternDeviceName((*devices)[0]->name());
    }

    AssignAndLog(assigned_device, node);

至此，所有Node的Placement信息都已经分配并微调完毕。

总结

经过Placer处理的GraphDef解决了显式和隐式Placement信息的所有冲突，可谓是最后一道防线。

在Placer之后，GraphDef将被送入GraphPartitioner模块中根据每个Node的device做子图切分，并插入Send，Recv以及必要的ControlFlow节点。因此，此步必不可少。

我们也可以看出，Placer模块的核心是对Placement进行微调，由于启发式规则相对简单，性能问题并未完全解决。甚至，我们马上可以想到，在分布式模式下，粗糙的Placement方案会让作业性能变得非常差，因为它会引入计算之外的通信开销。

TensorFlow为了高度灵活性，将Placement策略的负担丢给了用户，这也是为什么有些用户写出的TensorFlow分布式程序性能非常差的原因之一。站在TensorFlow框架的功能角度，它应该能够解放用户的编写程序负担，让用户能够完全专注在模型算法层面的研究中。

但是自动搜索Placement最佳策略的难度非常大，因为它要考虑集群通信的带宽，以及每个Op的计算量，是一个与硬件和环境高度联系的复杂问题。不仅如此，通常深度学习模型含有成千上万个Node，这使得方案的搜索空间巨大无比。

对于这个问题的解决办法，目前是百家争鸣。如果你对策略感兴趣，我这里给你推荐一篇Google发表的论文，它利用强化学习搜索更好的分片策略。有兴趣的同学可以参考这篇ICML的论文：Device Placement Optimization with Reinforcement Learning。

本文参与腾讯云自媒体同步曝光计划，分享自微信公众号。

原始发表：2020-08-24，如有侵权请联系 cloudcommunity@tencent.com 删除

node.js