TensorFlow的圖切割模塊——Graph Partitioner

背景

[做者： DeepLearningStack，阿里巴巴算法工程師，開源TensorFlow Contributor]

在通過TensorFlow的 Placer策略模塊調整以後，下一步就是根據Placement信息對Graph作切割，而後分發到不一樣的Device上去執行的過程了。在對Graph作切割時，爲了保證跨Device執行的邏輯與切割前一致並保證原圖中Node節點之間的依賴關係不受到破壞，不但須要插入Send、Recv通訊節點對，還須要維護相對複雜的Control Edge。這些功能被設計在了TensorFlow的Graph Partitioner模塊中。從該模塊的代碼量和原理上看，其內容很是好理解，但在涉及到對含有while_loop、loop_contition、exit、enter、merge、switch等Control Flow Op的圖作切割時，其處理就顯得相對複雜。本篇主要介紹Graph Partitioner的總體過程和相關源碼，但考慮到Control Flow Op相關的處理還須要一些前置知識，而這些前置知識在TensorFlow源碼閱讀與架構梳理系列中還沒有完成書寫，所以本篇暫時過濾掉對Control Flow Op相關邏輯的處理。

功能描述

顧名思義，Graph Partitioner是用來根據每一個節點的Placement信息對原圖作切割的，它主要包括三個核心步驟：

1. 對原圖的Placement信息作劃分，產生多個子圖Sub graph；

2. 爲具備跨Device依賴的節點對插入Send類和Recv類節點對；

3. 插入必要的Control Edge

一個完成了圖切割的Graph能夠在多個機器的分佈式集羣環境中執行，可是應當注意到在單機單卡時這一步驟也是必要的，由於TensorFlow是個異構框架，單機單卡也涉及到GPU和CPU之間的圖切割。圖切割的依據是Placement信息，若是想深刻了解Placement模塊相關內容，請參考本系列的這篇文章——《 TensorFlow中的Placement啓發式算法模塊——Placer》。

Graph Partitioner模塊十分通用，在單機單卡運行過程當中，DirectSession會讓Graph Partitioner根據不一樣的Device進行切割。而在分佈式運行過程當中，Graph Partitioner會被執行兩次，一次是SplitByWorker，另外一次是SplitByDevice。

Graph Partition切割流程

爲了描述方便，特地將圖切割過程分爲如下幾個子過程，總體流程以下圖所示，圖右邊的文字是對每一個過程的簡短描述，本篇咱們重點闡述標記爲深色的子過程。

第一步——分析構建Control Flow相關信息

這個過程在代碼中是經過AddControlFlow函數實現的，因爲改代碼深度依賴於Control Flow Op的相關模塊，且對於不含有Control Flow Op的Graph幾乎什麼都沒有作，所以咱們先忽略這個過程，等到對Control Flow模塊作詳細解讀時再回過頭來研究其在Graph Partitioner中的意義。

 1 GraphInfo g_info;
 2 if (!opts.control_flow_added) {
 3   // Add the "code" for distributed execution of control flow. Code is
 4   // added only for the frames that are placed on multiple devices. The
 5   // new graph is an equivalent transformation of the original graph and
 6   // has the property that it can be subsequently partitioned arbitrarily
 7   // (down to the level of individual device) for distributed execution.
 8   status = AddControlFlow(opts, g, &g_info);
 9   if (!status.ok()) return status;
10 }

第二步——構建Op的Input和Output Memory類型信息

在介紹這個過程以前，首先須要明確兩種概念，他們是DeviceMemory和HostMemory。前者指的是計算設備的Memory類型，後者指的是CPU的Memory類型，它們在TensorFlow中被定義爲Enum類型，代碼以下所示。

1 // MemoryType is used to describe whether input or output Tensors of
2 // an OpKernel should reside in "Host memory" (e.g., CPU memory) or
3 // "Device" Memory (CPU memory for CPU devices, GPU memory for GPU
4 // devices).
5 enum MemoryType {
6   DEVICE_MEMORY = 0,
7   HOST_MEMORY = 1,
8 };

對Op的Input和Output Memory信息進行檢索並構建緩存的函數是BuildMemoryDeviceInfo，該過程構建的信息對後面真正作圖切割很是重要。由於TensorFlow的Op在註冊時須要不但須要指定其在各個Device上的實現版本（好比CPU版本的Op和GPU版本的Op都是分別註冊到系統中的），還須要指出其Input和Output Tensor的類型以及所使用的Memory類型，即便某個Op存在GPU上的實現，它的GPU版本也有可能須要在CPU上讀入數據或輸出結果。例如，GPU版本的Reshape Op註冊代碼以下。

 1 #define REGISTER_GPU_KERNEL(type)                               \
 2   REGISTER_KERNEL_BUILDER(Name("Reshape")                       \
 3                               .Device(DEVICE_GPU)               \
 4                               .HostMemory("shape")              \
 5                               .TypeConstraint<type>("T")        \
 6                               .TypeConstraint<int32>("Tshape"), \
 7                           ReshapeOp);                           \
 8   REGISTER_KERNEL_BUILDER(Name("Reshape")                       \
 9                               .Device(DEVICE_GPU)               \
10                               .HostMemory("shape")              \
11                               .TypeConstraint<type>("T")        \
12                               .TypeConstraint<int64>("Tshape"), \
13                           ReshapeOp);

上面的宏顯示，雖然Reshape Op確實在GPU上有註冊的實現版本，可是它依然要使用HostMemory。另外，某些Tensor的類型也決定了其是否能夠被放置到Device Memory上，通常狀況下float類型的數據對於計算設備是很是友好的，而String類型就不是這樣，因此在types.cc文件中規定了一些強制被放在HostMemory的數據類型，以下代碼所示。

 1 bool DataTypeAlwaysOnHost(DataType dt) {
 2   // Includes DT_STRING and DT_RESOURCE.
 3   switch (dt) {
 4     case DT_STRING:
 5     case DT_STRING_REF:
 6     case DT_RESOURCE:
 7       return true;
 8     default:
 9       return false;
10   }
11 }

TensorFlow的設計哲學認爲，參與計算的Tensor應該被放在DeviceMemory上，而參與控制的Tensor應該放在HostMemory上。這樣的設計思路雖然有必定道理，但也確實對一些case產生了負面的性能影響。在後面的過程當中咱們能夠看到，Partition過程會根據每一個Op的Input和Output Memory類型決定是否插入Send類和Recv類節點對，所以會常常遇處處於同一個Device上的兩個節點也須要插入Send類和Recv類節點對的狀況，顯然這有可能帶來性能降低。

第三步——對原圖進行分析，併產出切割後的多個子圖

在面兩個步驟的準備工做完成以後，就能夠進行圖切割和Send類、Recv類節點對的插入，以及Control Edge的插入了，這個過程以下圖所示。由於流程圖繪製的比較簡潔，咱們將在下面對該圖進行詳細說明。

 

 

1.將原圖中取出一個節點dst，根據其Device將其分配到對應的Sub Graph中，而後以dst節點爲終點節點，沿着其接收Tensor的方向向輸入節點src進行分析；

2.Node之間的鏈接依靠的是Edge，所以對於dst來講須要根據其Input的Edge來分析src節點的位置，因此這裏要得到dst的全部Input Edge；

3.在逐個遍歷分析Input Edge時，第一個要處理的就是src和dst處於同一個Device，但依然須要插入Send類和Recv類節點對的狀況。根據第二步BuildMemoryDeviceInfo提供的信息，某些Op的註冊和特殊之處確實會獲得這種狀況；

4.若是決定須要插入Send類和Recv類節點對，那麼優先考慮是否能夠重用Recv節點，若是根據信息拼出的Key可以在緩存中搜索到該Recv Node，那麼則取出重用。這種Recv Fusion是一種性能優化手段，能避免屢次沒必要要的通訊，真正作到達到一次通訊屢次使用的目的，下面的代碼展現了這一個過程；

 1       // Check whether there is already a send/recv pair transferring
 2       // the same tensor/control from the src to dst partition.
 3       const bool on_host = IsDstInputOnHost(edge, g_info);
 4       DupRecvKey key{src->id(), edge->src_output(), dst_graph, on_host};
 5       auto iter = dup_recv.find(key);
 6       if (iter != dup_recv.end()) {
 7         // We found one. Reuse the data/control transferred already.
 8         const string& recv_node_name = iter->second.recv->name();
 9         if (edge->IsControlEdge()) {
10           AddInput(dst_def, recv_node_name, Graph::kControlSlot);
11         } else {
12           AddInput(dst_def, recv_node_name, 0);
13         }
14         ref_control_inputs.push_back(recv_node_name);
15 
16         // We want the start_time for the recv to be the smallest of the start
17         // times of it's consumers. So we update this whenever we use a recv,
18         // and write it out to the attribute at the end of the subroutine
19         if (iter->second.start_time > recv_start_time) {
20           iter->second.start_time = recv_start_time;
21         }
22         continue;
23       }

5.若是緩存中沒有找到可重用的節點，那麼只能建立新的Send類和Recv類節點對了。插入通訊節點對時須要考慮多種狀況，有時插入Send和Recv節點就能完成任務，有時還須要插入Control Edge以保證依賴順序，有時甚至還要插入一些其餘的輔助節點。事實上，分紅這三種邏輯處理已經覆蓋任何狀況了，後面一章將詳細闡述這三種處理邏輯。

第四步——必要的後處理

這是一些收尾的工做，過程很是簡單，好比完善Send和Recv節點的Incarnation信息，補全各個子圖的version信息等，代碼以下所示。

 1   const FunctionLibraryDefinition* flib_def = opts.flib_def;
 2   if (flib_def == nullptr) {
 3     flib_def = &g->flib_def();
 4   }
 5 
 6   // Set versions, function library and send/recv incarnation.
 7   for (auto& it : *partitions) {
 8     GraphDef* gdef = &it.second;
 9     *gdef->mutable_versions() = g->versions();
10     // Prune unreachable functions from `flib_def` before adding them to `gdef`.
11     *gdef->mutable_library() = flib_def->ReachableDefinitions(*gdef).ToProto();
12 
13     // Traverse the graph to fill every send/recv op's incarnation
14     // information.
15     SetIncarnation(opts, gdef);
16   }

Send和Recv節點對插入的三種狀況

在代碼中，聲明插入Send和Recv節點的代碼段很是簡單，以下所示。

 1       // Need to split edge by placing matching send/recv nodes on
 2       // the src/dst sides of the edge.
 3       NodeDef* send = AddSend(opts, g_info, src_graph, edge, send_from,
 4                               send_start_time, &status);
 5       if (!status.ok()) return status;
 6 
 7       NodeDef* real_recv = nullptr;
 8       NodeDef* recv =
 9           AddRecv(opts, g_info, dst_graph, edge, &real_recv, &status);
10       if (!status.ok()) return status;

可是對於不一樣的狀況卻有着豐富的處理邏輯，因此下面在展現示意圖的同時，會將相關的代碼段摘出來作展現。

在同一個Device上插入Send和Recv節點對

由於同一個Device上的Send和Recv節點在執行過程當中實際上Memory Copy，而Recv的kernel又是異步的，因此須要有一種機制保證保證Recv必定要在Send以後執行，所以須要在Send和Recv之間插入一個Control Edge，從圖的依賴上保證它們的執行順序。

這個過程的關鍵是在插入Send和Recv節點以後，須要插入額外的Control Edge，代碼以下。

// Fix up the control flow edge.
// NOTE(yuanbyu): 'real_recv' must be the real recv node.
if (src_graph == dst_graph) {
  // For same device send/recv, add a control edge from send to recv.
  // This prevents the asynchronous recv kernel from being scheduled
  // before the data is available.
  AddInput(real_recv, send->name(), Graph::kControlSlot);
}

跨Device根據DataFlow插入Send和Recv節點對

這是最容易理解的一種狀況，Send節點須要插入到和src節點相同的Device上，Recv須要插入到和dst節點相同的Device上。而且爲了減小沒必要要的通訊開銷，儘量的重用Recv節點。

該過程的關鍵在於複用Recv節點，前面在獲取緩存時已經闡述過，這裏不重複展現。

跨Device根據ControlFlow插入Send和Recv節點對 

當存在跨Device的Control Flow依賴時，問題變得相對複雜。由於Control Edge只是用做控制，它並不傳輸真正的Tensor，但在跨Device的狀況下，必需要向dst所在的Device發送消息，讓其知曉存在依賴控制。TensorFlow選擇發送DummyConst的方式通知dst節點，具體而言，須要在src的Device上插入shape爲0的DummyConst節點，而後將其做爲Send的惟一輸入，並將src節點做爲它的Control Dependncy。另外一方面，在dst的Device上插入Recv節點以後，還須要插入一個identity節點負責讀取發送來的DummyConst，而後將Indentity做爲dst的Control Dependency。如此一來，這種跨Device的依賴關係就能夠被徹底等價的表示出來。

這個過程的關鍵在於src端的DummyConst插入和dst端的Identity插入，這兩部分的邏輯處理寫在了兩個地方。DummyConst和相關控制依賴的代碼以下。

 1       NodeDefBuilder::NodeOut send_from;
 2       if (edge->IsControlEdge()) {
 3         // Insert a dummy const node that will generate a tiny
 4         // data element to be sent from send to recv.
 5         VLOG(1) << "Send/Recv control: " << src->assigned_device_name() << "["
 6                 << src->name() << "] -> " << dst->assigned_device_name() << "["
 7                 << dst->name() << "]";
 8         NodeDef* dummy = AddDummyConst(opts, src_graph, edge, &status);
 9         if (!status.ok()) return status;
10         // Set the start time for this dummy node.
11         if (opts.scheduling_for_recvs) {
12           AddNodeAttr("_start_time", send_start_time, dummy);
13         }
14         AddInput(dummy, src->name(), Graph::kControlSlot);
15         send_from.Reset(dummy->name(), 0, DT_FLOAT);
16       } else {
17         send_from.Reset(src->name(), edge->src_output(), EdgeType(edge));
18       }

Indentity即相關依賴的插入邏輯被寫在了AddRecv中，下面展現了這個片斷。

 1   // Add the cast node (from cast_dtype to dtype) or an Identity node.
 2   if (dtype != cast_dtype) {
 3     const string cast_op = (host_memory) ? "_HostCast" : "Cast";
 4     NodeDefBuilder cast_builder(opts.new_name(src->name()), cast_op);
 5     cast_builder.Attr("DstT", dtype);
 6     cast_builder.Device(dst->assigned_device_name())
 7         .Input(recv->name(), 0, cast_dtype);
 8     NodeDef* cast = gdef->add_node();
 9     *status = cast_builder.Finalize(cast);
10     if (!status->ok()) return nullptr;
11     return cast;
12   } else if (edge->IsControlEdge()) {
13     // An Identity is only needed for control edges.
14     NodeDefBuilder id_builder(opts.new_name(src->name()), "Identity");
15     id_builder.Device(dst->assigned_device_name())
16         .Input(recv->name(), 0, cast_dtype);
17     NodeDef* id = gdef->add_node();
18     *status = id_builder.Finalize(id);
19     if (!status->ok()) return nullptr;
20     return id;
21   } else {
22     return recv;
23   }

關於使用bfloat16壓縮通訊

TensorFlow支持經過使用bfloat16減小通訊量，雖然bfloat16理論上是有損精度的，可是大量的實踐證實這個精度損失是基本感知不到的。bfloat16的通訊功能能夠經過如下配置項打開，只要在建立Session時傳入打開該功能的config便可。

graph_options = tf.GraphOptions(enable_bfloat16_sendrecv=True)
session_config = tf.ConfigProto(gpu_options=gpu_options)　　

而TensorFlow在底層插入bfloat的轉換節點就是在Graph Partitioner的AddSend函數和AddRecv函數中插入的，可是這個轉換隻會在跨Device的Send和Recv先後插入，這也很是符合邏輯，由於處於同一個Device的Send和Recv本質上是本地的Memory Copy，其帶寬很是高，因此通訊並非瓶頸，而插入兩個轉換節點只能帶來額外的轉換開銷。

總結

本文介紹了TensorFlow中的圖切割模塊——Graph Partitioner。考慮到Graph Partitioner在處理含有Control Flow Op的Graph時具備更加複雜的邏輯，而本系列還沒有完成Control Flow模塊的編寫，所以在梳理源碼時只對通常狀況做了詳細闡述。事實上，僅僅是這些內容也已經可讓讀者對TensorFlow的圖切割過程有了較好的理解。不管是SplitByDevice仍是SplitByWorker，Graph Partitioner做爲TensorFlow的圖切割模塊都具備良好的模塊化通用化特色，它的關鍵點在於如何保證切割後的多個子圖和原圖具備徹底的邏輯等價性。Graph Partitioner可以正常工做的前提是Graph中的每一個Node都具備了Device Placement信息，所以在一次Run過程當中，Graph Partitioner是在Placer模塊完成以後才進行的。從此咱們在梳理單機多卡和分佈式執行引擎時，咱們還會看到Placer和Graph Partitioner的身影，這也是本系列中屢次強調其重要性的緣由。