Tensorflow XLA HLO — BufferLiveness

XLA Sevice 的 High Level Optimization可以分为设备无关优化和设备相关优化两部分, 设备无关优化是Buffer优化, 以便使用更少的主存或显存(以及其他硬件backend), 设备相关的优化主要是根据硬件特性优化执行流, HLO设备无关优化主要分为3个阶段:

  1. 支持分析指定Buffer的Liveness
  2. 构造逻辑Memory对象到物理Memory对象的映射
  3. 优化物理Memory, 交给特定Backend做specific的进一步优化

本文分析第一个阶段–BufferLiveness分析

所以,BufferLiveness, 就是获取Buffer生命周期关系, 以便决定Buffer复用策略. 源码整体调用栈如下, 该阶段进一步也分3个过程: LogicalBufferAnalysis->TuplePointsToAnalysis->BufferLiveness

-1- Memory优化入口
-3-16- 获取BufferLiveness, 对比-16-, 获取到的BufferLiveness会用于支撑Buffer优化决策,用OOP的方式优雅的实现pipeline
-8- 获取LogicalBufferAnalysis
-9- 根据HloInstruction的Shape, 构造一组相应的LogicalBuffer, Shape本身是一个树结构,表示一个HloInstruction的输出形状, 一个LogicalBuffer对应一个Shape的一个节点(Subshape), 表示一块逻辑buffer, 用pair<HloInstruction*, ShapeIndex>来索引一个LogicalBuffer。
-10- 作为Liveness分析的原材料-LogicalBuffer已经构造完毕, 存储下来, 并返回analysis, 准备进入下一阶段, 进行TuplePointsToAnalysis.
-11- 用存储有所有LogicalBuffer信息的LogicalBufferAnalysis实例构造TuplePointsToAnalysis实例, Tuple, 用来描述Buffer树状结构的方式, E = {%1, %2, {%3, %4}} 表示这样一个树:深为3, 深1的节点有一个, 树的根, 深为2的节点有3个, %1, %2, 没体现名字的节点暂且叫”X”, 深为3的节点有2个, %3, %4, 这两个节点是上一层中”X”的子节点. PointsTo, “指向”, 在前面的例子中, root的”PointsTo”就是%1, %2和”X”, “X”的”PointsTo”就是%3和%4, 所以TuplePointsToAnalysis就是分析整个计算图中的LogicalBuffer依赖关系并存储在PointsToSet中,后面会详细分析
-12- 执行分析逻辑.
-13- 对计算图中的LogicalBuffer的依赖关系分析完毕, 存储下来, 并返回analysis, 准备进入下一阶段, 进行BufferLiveness
-14- 用TuplePointsToAnalysis实例获取root的alias_buffer, 也就是潜在的需要传出的HloModule的Buffer, 存储在maybe_live_out_buffers_中.
-15- 返回liveness实例, TuplePointsToAnalysis实例会存储LogicalBuffer的依赖关系, 但BufferLiveness并不会存储每一个LogicalBuffer的”liveness”, 而是基于TuplePointsToAnalysis封装了一组判断特定LogicalBuffer的函数.
-16- 将BufferLiveness实例传入构造LogicalBuffer与BufferAllocation映射关系的BufferAssignment实例.

Continue reading

BERT Training By Example

BERT 模型的训练主要分为三个阶段:
1. Masked Language Model,基于对自然语言的语料打标签生成训练样本,将无监督学习转换为监督学习
2. Pretraining,模型训练的核心过程,生成能理解语言的模型
3. Fine-tuning,基于对模型的理解,完成特定的任务。

Masked Language Model

Masked Language Model 的任务是将无标签的语料转换为有标签的语料,进而将无监督学习转化为监督学习。这个过程需要准备两份材料:原始语料词汇表
原始语料就是一系列符合语法的句子或段落。
词汇表维持token-id的映射,除了一门语言自身的词汇,BERT还要求词汇表包含[PAD], [UNK], [CLS], [SEP]和[MASK]。

Pretraining

Continue reading

为Tensorflow 添加新Op

Tensorflow内部添加op有3种方式:
1. 不修改Tensorflow源码, 在运行环境里通过Tensorflow包中的.h和.so直接扩展, 无需Tensorflow编译环境, 二次开发首选
2. 修改Tensorflow源码, 利用源码中的User-supplied Op接口直接将新Op添加到Tensorflow安装包中, 关注Op实现, 无需调试Tensorflow的Bazel编译框架, 但个人认为比较鸡肋
3. 修改Tensorflow源码, 以Standard Op的方式将Op添加在Tensorflow包中, 将新Op提交到社区需要这种方式. 需要实现Op, 同时需要修改Tensorflow的Bazel编译框架, 由于Tensorflow使用了自动生成代码的技术, 后者调试比较困难,谨慎使用

本文主要介绍这3种添加TfOp的方式, 最后, 简要介绍添加XlaOp的方法. 对于需要修改源码的方法2和方法3, 在Tensorflow内部均需要修改python接口和tfop两层, 而每一种又有两种方式:User-supplied Op 和Standard Op, 所以一共是2*2 = 4 种, 简化起见, 本文分别用用user-supplied op实现了XZeroOutOp(输入清零后输出), 用Standard Op 实现了XDoNothingOp(输入透传给输出), 在实际操作中, tensorflow内部这两层使用自动生成代码技术进行了解耦, python层和tf层使用不同的方式完全没问题.

Customized Op

Horovod就是使用Cutomized Op这种方式的典型框架, 基于Tensorflow 接口, MPI和NCCL接口, Horovod在不修改Tensorflow源码的前提下实现了Ring-allreduce 的Op. 使用这种方式Op实现的方式与Tensorflow内部的Standard Op并无区别, 关键是编译的时候要include相关的头文件:”/usr/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/”, 以及链接相应的库:”/usr/lib/python2.7/site-packages/tensorflow/libtensorflow_framework.so”

此外需要注意, Tensorflow框架会自动将驼峰命名转化为下划线命名, 即名为”HorovodAllreduce”的Op, 在生成的so中, 对应”horovod_allreduce”, 感兴趣的同学可以参考Tensorflow自动生成py的代码.

User-supplied Op

Continue reading

Tensorflow-Horovod安装部署checklist

基础环境 GPU 检查多机中每个节点的GPU是否归属同一型号, horovod后台线程周期性处理提交的梯度, 如果由于GPU计算能力不同或网络延迟导致某次allreduce中某个提交者落后超过一个周期, 会发生使用错误的Tensor进行allreduce的情况, 通常, 这都会导致Tensor Shape Mismatch的错误, 笔者就层遇到过不小心同时使用M40和P40的情况, 在tensorflow1.8还能的运行, 但1.14就无法多机reduce, 定位了好久才发现是硬件不同
cuda 是否与驱动match:cuda兼容性检查下载, 随同一起安装的还有cupti等组件
cudnn 版本是否与cuda版本match: cudnn兼容性检查及下载
nccl 版本是否与cuda版本match: nccl兼容性检查及下载
mpi 其中libmpi.so要对外可见, 多机训练需要tensorflow编译的时候就要将MPI链接到lib中, 同时, MPI对于horovod也是必要依赖
Tensorflow tensorflow 选择的tensorflow 版本是否兼容系统cuda, 官方兼容列表
tensorboard tensorboard是否兼容tensorflow, tensorboard的版本号要与tensoflow保持一致
tensorflow-estimator tensorflow-estimator是否兼容了tensorflow, tensorflow 1.10.6 之后的特性, 是对底层API的封装, 需要与tensorflow版本号一致
bazel tensorflow对使用的bazel版本比较敏感, 如果没有修改bazelrc却出现bazelrc文件解析错误, 就要考虑下是不是bazel版本的问题, 参考tensorflow测试通过的bazel版本bazel下载
编译选项 tensorflow编译选项, 是否包含了cuda和nccl, 这两个缺省都是不安装的, 此处使用GPU, cuda为必须, 加之多机多卡, NCCL可以减少通信开销, 应该编译进tensorflow. 在1.13之前, nccl是随同config的过程一同配置, 1.14开始使用环境变量 + .bazelrc控制NCCL的安装, 具体的, 环境变量中需要export以下环境变量 同时, 修改tensorflow/.bazelrc文件
Horovod 环境变量 环境变量是否指定了libtensorflow_framework.so的路径以及NCCL路径信息

Tensorflow XLA Service 详解 II

本文主要介绍在XLA service阶段针对HloInstruction做的一些显存优化, 对于训练框架来说, 显存优化的工作至关重要, 主要是由于现阶段GPU+CUDA远没有CPU+Linux组合强大, 后者有完善的建立在虚拟内存基础上的内存管理机制, 内存的高效使用由linux kernel来负责, 即便物理内存不足, 还可以使用swap, 内存压缩等技术确保内存的高效供应, 而在GPU+CUDA里, 这方面的工作很大程度让渡给了程序员自己来搞定, GPU程序接触到的就是物理显存, 如果程序的显存申请超过显存容量, 整个程序就会直接coredump, 此外, 显存本身就集成在GPU板卡上, 无法像内存一样扩展, 而GPU本身造价昂贵, 最后, 在深度学习训练中, 大力出奇迹的现状下, 显存的消耗明显超过的摩尔定律, 这也加剧了显存供求关系的矛盾, 正式由于训练框架做了大量的优化, 才能让模型跑起来.

XLA Service的显存优化设计思想与tensorflow整体一样遵循”静态图”的设计: 先整体优化, 再落地实施. 其中, xla/service/buffer_assignment.cc 是整个显存优化的核心, 在1.14版本中, xla/service/支持两种后端: cpu和gpu, 纷纷针对两种backend有进一步的优化算法, 本文主要针对GPU的优化逻辑进行分析

核心文件

内存优化公共:
xla/service/buffer_assignment 内存优化核心文件
xla/service/buffer_liveness.cc  内存片生命周期分析
GPU相关:
xla/service/gpu/buffer_allocations.cc BufferAllocation的组合
xla/service/gpu/gpu_hlo_scheudle.cc   Hlo的处理顺序, 与显存的优化策略息息相关, 简单地说, 按照BFS并行执行的HloInstruction消耗的显存肯定大于所有的HloInstruction都顺序执行. 

核心抽象

HloSchedule

XxxSchedule是TF的代码风格, 类似的有更底层用于Thunk调度的ThunkSchedule, 以及Service提供的HloSchedule. HloSchedule中最重要的就是封装了HloOrdering.

HloSchedule

XLAService内存优化的本质是处理LoigicalBuffer和BufferAllocation之间的关系, 原则是使用尽可能少的BufferAllocation去承载尽可能多的LogicalBuffer, 而如何使用的更少, 就涉及到了对Hlo图的分析, 就涉及到了Ordering的问题, 使用不同策略生成Ordering, 直接影响两个LogicalBuffer之间的约束关系, 最简单的, 在图遍历中, 使用DFS和BFS的2种方式遍历会导致图上节点的内存依赖关系大有不同.

HloOrdering是描述HloInstruction加载序列的基类, 派生类有PredecessorHloOrdering, DependencyHloOrdering 和 SequentialHloOrdering, 其中, DependencyHloOrdering基于依赖关系, 所以可以并行, 性能更高, 但耗更多的内存, 而SequentialHloOrdering完全串行, 性能相对低, 但可以节约更多内存, 而 PredecessorHloOrdering 是个虚类, 需要子类进一步填充predecessors_, 这也是GPU后端使用的方式.不同的Ordering会影响内存的依赖关系, 进一步影响Launch到GPU后Kernel的执行序列. 

Continue reading

Tensorflow XLA Service 详解 I

compiler/aot/ 以AOT的方式将tf2xla/接入TF引擎
compiler/jit/以JIT的方式将tf2xla/接入TF引擎, 核心是9个优化器和3个tfop,其中XlaCompileOp调用tf2xla的“编译”入口完成功能封装,XlaRunOp调用xla/client完成“运行”功能。
compiler/tf2xla/对上提供xla_compiler.cc:XlaCompiler::CompileFunction()供jit:compile_fn()使用将cluster转化为XlaComputation。核心是利用xla/client提供的接口,实现对XlaOpKernel的“Symbolic Execution”功能。每个XlaOpKernel子类均做的以下工作: **从XlaOpKernelContext中取出XlaExpression或XlaOp, 调用xla/client/xla_buidler.h提供的方法完成计算, 将计算结果的XlaOp存入XlaKernelContext.**
compiler/xla/client/ 对上提供xla_builder.cc:Builder等供CompileFunction()使用,将Graph由Op表达转化为HloModuleProto:HloComputationProto:HloInstructionProto表达并保存在XlaComputation中。
对上提供local_client.cc:LocalClient::Compile(),作为编译入口供jit:BuildExecutable()使用,将已经得到的XlaComputation交给service并进一步编译为二进制。
对上提供local_client.cc:LocalExecutable::Run(),作为运行入口供jit/kernels/xla_ops.cc:XlaRunOp使用,通过Key找到相应的二进制交给service层处理
compiler/xla/service/ 对上提供local_service.cc:LocalService::BuildExecutable()供LocalClient::Compile()使用实现真正的编译,承接XlaComputation封装的HloProto, 将其转化为HloModule:HloComputation:HloInstruction表达, 对其进行优化之后, 使用LLVM后端将其编译为相应Executable后端的二进制代码
对上提供executable.cc:Executable::ExecuteOnStream()供LocalExecutable::Run()使用实现真正的执行二进制。

编译cubin

调用栈:

Continue reading

Tensorflow XLA Client | HloModuleProto 详解

compiler/aot/ 以AOT的方式将tf2xla/接入TF引擎
compiler/jit/以JIT的方式将tf2xla/接入TF引擎, 核心是9个优化器和3个tfop,其中XlaCompileOp调用tf2xla的“编译”入口完成功能封装,XlaRunOp调用xla/client完成“运行”功能。
compiler/tf2xla/对上提供xla_compiler.cc:XlaCompiler::CompileFunction()供jit:compile_fn()使用将cluster转化为XlaComputation。核心是利用xla/client提供的接口,实现对XlaOpKernel的“Symbolic Execution”功能。每个XlaOpKernel子类均做的以下工作: 从XlaOpKernelContext中取出XlaExpression或XlaOp, 调用xla/client/xla_buidler.h提供的方法完成计算, 将计算结果的XlaOp存入XlaKernelContext.**
compiler/xla/client/ 对上提供xla_builder.cc:Builder等供CompileFunction()使用,将Graph由Op表达转化为HloModuleProto:HloComputationProto:HloInstructionProto表达并保存在XlaComputation中。
对上提供local_client.cc:LocalClient::Compile(),作为编译入口供jit:BuildExecutable()使用,将已经得到的XlaComputation交给service并进一步编译为二进制。
对上提供local_client.cc:LocalExecutable::Run(),作为运行入口供jit/kernels/xla_ops.cc:XlaRunOp使用,通过Key找到相应的二进制交给service层处理
compiler/xla/service/ 对上提供local_service.cc:LocalService::BuildExecutable()供LocalClient::Compile()使用实现真正的编译,承接XlaComputation封装的HloProto, 将其转化为HloModule:HloComputation:HloInstruction表达, 对其进行优化之后, 使用LLVM后端将其编译为相应Executable后端的二进制代码
对上提供executable.cc:Executable::ExecuteOnStream()供LocalExecutable::Run()使用实现真正的执行二进制。

compiler/xla/client 向上为tf2xla/下的XlaOpKernel的实现提供支撑, 将上层请求转换为HloModule交给下层xla/service优化并编译.

接口上, client做上表中的三件事 , 实际上, 只有XlaOpKernel->HloProto在Client完成, 而对于另外两个编译和执行,都是在service中完成的,或者说,Client本质上是Service的Facade + Proxy,对繁琐的构造HloModuleProto的过程进行了封装,降低了XLA Service与上层模块的耦合

 client.h:Client    Client基类, 用于多态实现
 client_library.h:ClientLibarary  使用单例构造client_library对象,  用于检索/构造所需的Client实例
 lib/   同builder一起实现”Symbolic Execution”
 local_client.h:LocalClient, LocalExecutable JIT 使用的LocalClient定义, 是Service相关方法的Proxy
 xla_builder.h:XlaBuilder    提供接口用tf2xla使用实现”Symbolic Execution”, 是其中XlaBuilder::Build()是构造client构造HloModuleProto的核心方法。对于需要多个步骤完成初始化的类, 我们会使用Builder模式, 这就是个例子
 xla_computation.h:XlaComputationXlaComputation对象是对HloModuleProto的封装, 用于进一步二进制编译

职责1: 构造HloModuleProto

编译二进制之前首先要完成Graph表达方式的映射: Client之前的tf2xla的Graph由使用Op表达, Client之后的Service的Graph使用HloInstruction表达, Client负责完成这种转化, 具体地, 就是将Op转化为HloProto格式, 再交给Service解析为Hlo格式, 其中的HloProto就是封装在XlaComputation中. 所以, 这个过程可以看做是”编译”的准备工作. 在这个过程中, Graph, Cluster, XlaComputation, HloModuleProto, HloModule是一一对应的,都是”Program”的概念, 这一点从源码的ProgramShape等变量名中都可以体现

按照OOP的方式理解OOP的代码, 在说明Client是如果构造HloModuleProto之前, 来介绍几个概念, HloModuleProto-HloComputationProto- HloInstructionProto, 这3个概念是XLA Service提供给上层调用者(Here, XLA Client)的用于构造HloModule-HloComputation-HloInstruction的protobuf形式的”原材料”, Module-Computation-Instruction, 在概念上可以对应程序-函数-语句

  1. HloProto中的函数调用都是以”半inline“方式实现, inline的部分是指类似”a = add(x, y);”这样的语句, 会将add函数体整个copy到该语句的上下文, 在执行的时候调用之, 在HloProto中, 被调用的XlaComputation也会被copy到调用语句的上下文, “上下文”就是调用语句HloInstructionProto所处的XlaComputation(对应一个HloModuleProto)的XlaBuilder实例. “半”的部分是指被copy的XlaComputation并不是直接在调用语句处展开, 而是在调用语句的HloModuleProto中存储被调用的XlaComputation的Id, 执行的时候直接跳转执行.
  2. 类似多文件编程, 构造的过程可以有多个HloModuleProto, 但由于上一条所述, 最终所有的被调用函数均会被copy到根HloModuleProto的上下文, 最终只有也只需要根HloModuleProto就足以以描述整个程序的逻辑
  3. 一个C/C++程序有main作为入口, 对应到HloModuleProto就是entry_computation, 一个函数的第一条语句, 对应到HloComputationProto就是root_id

下面我们简要分析下代码, 既然Hlo是描述一个程序, 那么就需要解决一个基本的问题: 描述出函数->语句->函数->语句…这样的递归结构, OOP的Composite模式, 或者内核的kset-kobj结构, 都是解决递归的好方式, Hlo模块的实现也比较简单, 可以认为是kset-kobj的结构的C++版本:在HloComputationProto里记录HloInstructionProto,在HloInstructionProto里记录被调用的另一个HloCProto实例Id

Continue reading

Tensorflow XlaOpKernel | tf2xla 机制详解

compiler/aot/ 以AOT的方式将tf2xla/接入TF引擎
compiler/jit/以JIT的方式将tf2xla/接入TF引擎, 核心是9个优化器和3个tfop,其中XlaCompileOp调用tf2xla的“编译”入口完成功能封装,XlaRunOp调用xla/client完成“运行”功能。
compiler/tf2xla/对上提供xla_compiler.cc:XlaCompiler::CompileFunction()供jit:compile_fn()使用将cluster转化为XlaComputation。核心是利用xla/client提供的接口,实现对XlaOpKernel的“Symbolic Execution”功能。每个XlaOpKernel子类均做的以下工作: 从XlaOpKernelContext中取出XlaExpression或XlaOp, 调用xla/client/xla_buidler.h提供的方法完成计算, 将计算结果的XlaOp存入XlaKernelContext.
compiler/xla/client/ 对上提供xla_builder.cc:Builder等供CompileFunction()使用,将Graph由Op表达转化为HloModuleProto:HloComputationProto:HloInstructionProto表达并保存在XlaComputation中。
对上提供local_client.cc:LocalClient::Compile(),作为编译入口供jit:BuildExecutable()使用,将已经得到的XlaComputation交给service并进一步编译为二进制。
对上提供local_client.cc:LocalExecutable::Run(),作为运行入口供jit/kernels/xla_ops.cc:XlaRunOp使用,通过Key找到相应的二进制交给service层处理
compiler/xla/service/ 对上提供local_service.cc:LocalService::BuildExecutable()供LocalClient::Compile()使用实现真正的编译,承接XlaComputation封装的HloProto, 将其转化为HloModule:HloComputation:HloInstruction表达, 对其进行优化之后, 使用LLVM后端将其编译为相应Executable后端的二进制代码
对上提供executable.cc:Executable::ExecuteOnStream()供LocalExecutable::Run()使用实现真正的执行二进制。

从Kernel的视角, XLA并不会新增Op, 而是针对已有的Op, 新增了基于XLA的另一个版本的Kernel: XlaOpKerne。在TF引擎中, OpKernel在软件栈上已是底层, 即最终的计算过程都要在OpKernel中实现. 但在XLA中, XlaOpKernel只是编译的入口, 大量的实际工作都交给了更下层的XLA引擎去完成.XLA相关的代码在tensorflow/compiler中.

tf2xla/负责XlaOpKernel的构造, 注册. 虽然XLA与TF引擎不在一层, 但二者面临的问题有很多有相似之处, 比如都需要对Kernel和Device保持易扩展性, 都需要维持前驱/后继Kernel的数据流和控制流关系. 基于类似的种种原因, XLA内部实现的注册XlaOpKernel的接口与TF引擎中注册OpKernel的风格十分相似, 同时, 其内部实现又有本质的不同, 而这些”不同”, 正是我们需要关注的.

要理解XlaOpKernel与OpKernel的不同, 关键在于了解”Symbolic Execution“.
先来看TF引擎, 它的OpKernel::Compute()方法要: OpKernelContext.Input()取输入数据 ==> 计算 ==> OpKernelContext.SetOutput()存输出数据, 计算结果继续通过OpKernelContext流入后继Opkernel, 其中流动的是真正的训练数据, 暂且将这个过程称之为”Execution”.
对比之下, XLA中的”Symbolic Execution”中的”Symbolic”即是说, XlaOpKernel的设计目的不在于去处理训练数据, 而在于去生成能够正确的处理数据的代码. 整个JIT类似于python解释器,先将程序编译为二进制,再运行二进制,XlaOpKernel执行的”Symbolic Execution”就类似其中的”编译为二进制”的过程. 具体地, 在XlaOpKernel::Compile()中: XlaOpKernelContext.Input()以XlaOp形式取输入 ==> 调用xla/client/xla_buidler.h提供的方法实现Op该有的功能, 实际上是生成一组能处理数据的HloInstruction ==> XlaOpKernelContext.SetOutput()存储XlaOp形式的结果, 计算结果继续通过XlaOpKernelContext流入后继XlaOpkernel, 其中流动的都是以XlaOp表征的对训练数据的处理方法.

至于真正处理数据的时机, 就要交给XLA引擎, 它来负责后续的”编译”和”执行”, 具体地, 在JIT中, XlaCompileOp会在所有的XlaOpKernel::Compile()执行完毕之后, 继续调用xla/service中相应的方法将这些所有生成的HloInstruction编译生成二进制并进一步交给XlaRunOp来执行.

Continue reading

Tensorflow JIT 技术详解

compiler/aot/ 以AOT的方式将tf2xla/接入TF引擎
compiler/jit/以JIT的方式将tf2xla/接入TF引擎, 核心是9个优化器和3个tfop,其中XlaCompileOp调用tf2xla的“编译”入口完成功能封装,XlaRunOp调用xla/client完成“运行”功能。
compiler/tf2xla/对上提供xla_compiler.cc:XlaCompiler::CompileFunction()供jit:compile_fn()使用将cluster转化为XlaComputation。核心是利用xla/client提供的接口,实现对XlaOpKernel的“Symbolic Execution”功能。每个XlaOpKernel子类均做的以下工作: **从XlaOpKernelContext中取出XlaExpression或XlaOp, 调用xla/client/xla_buidler.h提供的方法完成计算, 将计算结果的XlaOp存入XlaKernelContext.**
compiler/xla/client/ 对上提供xla_builder.cc:Builder等供CompileFunction()使用,将Graph由Op表达转化为HloModuleProto:HloComputationProto:HloInstructionProto表达并保存在XlaComputation中。
对上提供local_client.cc:LocalClient::Compile(),作为编译入口供jit:BuildExecutable()使用,将已经得到的XlaComputation交给service并进一步编译为二进制。
对上提供local_client.cc:LocalExecutable::Run(),作为运行入口供jit/kernels/xla_ops.cc:XlaRunOp使用,通过Key找到相应的二进制交给service层处理
compiler/xla/service/ 对上提供local_service.cc:LocalService::BuildExecutable()供LocalClient::Compile()使用实现真正的编译,承接XlaComputation封装的HloProto, 将其转化为HloModule:HloComputation:HloInstruction表达, 对其进行优化之后, 使用LLVM后端将其编译为相应Executable后端的二进制代码
对上提供executable.cc:Executable::ExecuteOnStream()供LocalExecutable::Run()使用实现真正的执行二进制。

XLA基于编译技术将静态子图转换为二进制进而实现在某些网络下的加速,以BERT为例,P40单卡每batch性能从850ms提升到了700ms

JIT 是目前TF中两种XLA应用方式之一, 借助TF对计算图的先优化再执行的机制, JIT使用9个优化器+3个基于XLA模块实现的Tfop将Just In Time技术接入TF图计算引擎. 

9个Optimization如下图所示:

其中, MarkForCompilationPass, EncapsulateSubgraphsPass 和 BuildXlaOpsPass 最为关键.

3个Op如下, 经过之前的优化, 已经完成了”图->XlaCompileOp + XlaRunOp”的转化, 而JIT的编译过程就在替换了原始Op的XlaCompile中进行. XlaCompileOp编译得到的二进制直接送到紧接其后的XlaRunOp中执行, 由于, XlaCompileOp里有用于存储之前编译结果的Cache, 所以理想情况下(图不变,输入的shape也不变), 只有第一次会真正的编译, 之后的step中由于Cache hit, XlaCompileOp的成本就很低了, 这也是XLA你能够实现加速核心原因. 据此, 在特征识别等输入频繁变动的场景, 由于XlaCompileOp的Cache Miss的概率大大增加, 整体性能就会比常规的TF执行引擎差. 

Optimization

Continue reading

Tensorflow JIT/XLA UML

compiler/aot/ 以AOT的方式将tf2xla/接入TF引擎
compiler/jit/以JIT的方式将tf2xla/接入TF引擎, 核心是9个优化器和3个tfop,其中XlaCompileOp调用tf2xla的“编译”入口完成功能封装,XlaRunOp调用xla/client完成“运行”功能。
compiler/tf2xla/对上提供xla_compiler.cc:XlaCompiler::CompileFunction()供jit:compile_fn()使用将cluster转化为XlaComputation。核心是利用xla/client提供的接口,实现对XlaOpKernel的“Symbolic Execution”功能。每个XlaOpKernel子类均做的以下工作: **从XlaOpKernelContext中取出XlaExpression或XlaOp, 调用xla/client/xla_buidler.h提供的方法完成计算, 将计算结果的XlaOp存入XlaKernelContext.**
compiler/xla/client/ 对上提供xla_builder.cc:Builder等供CompileFunction()使用,将Graph由Op表达转化为HloModuleProto:HloComputationProto:HloInstructionProto表达并保存在XlaComputation中。
对上提供local_client.cc:LocalClient::Compile(),作为编译入口供jit:BuildExecutable()使用,将已经得到的XlaComputation交给service并进一步编译为二进制。
对上提供local_client.cc:LocalExecutable::Run(),作为运行入口供jit/kernels/xla_ops.cc:XlaRunOp使用,通过Key找到相应的二进制交给service层处理
compiler/xla/service/ 对上提供local_service.cc:LocalService::BuildExecutable()供LocalClient::Compile()使用实现真正的编译,承接XlaComputation封装的HloProto, 将其转化为HloModule:HloComputation:HloInstruction表达, 对其进行优化之后, 使用LLVM后端将其编译为相应Executable后端的二进制代码
对上提供executable.cc:Executable::ExecuteOnStream()供LocalExecutable::Run()使用实现真正的执行二进制。




Related:
Tensorflow XLA HLO I — BufferLiveness
Tensorflow XLA Service 详解 II
Tensorflow XLA Service 详解 I
Tensorflow XLA Client | HloModuleProto 详解
Tensorflow XlaOpKernel | tf2xla 机制详解
Tensorflow JIT 技术详解
Tensorflow JIT/XLA UML
Tensorflow OpKernel机制详解
Tensorflow Op机制详解
Tensorflow Optimization机制详解
Tensorflow 图计算引擎概述

Tensorflow XLA Debug/Profiling Methods

Tensorflow 调试方式主要有以下几种

  1. Timeline初步分析性能瓶颈
  2. Tensorboard,Graphviz等对图信息进行可视化
  3. 环境变量配置导出图信息
  4. 使用bazel编译选项输出内部信息
  5. 使用GDB进行细节问题定位

Timeline

可以直观的显示一个Step内Op粒度的执行时间, 为性能优化提供参考。Chrome中键入 chrome://tracing/ ,即可通过Load按钮加载time.json, 使之图形化显示

Timeline看起来很美好,但有几个问题一定要注意:

  1. 基本单位是Op。 Op内逻辑再简单,也一定会体现,Op内逻辑再复杂,也不会对其拆分。
  2. 是谁的执行流。Timeline体现的是CPU端启动Op的执行流,,上图中的GPU端的执行流看起来有多条,实际上到底有没有并行呢?这不一定,这张Timeline图本质只是对time.json进行了可视化,这里看到的GPU任务就对应time.json的 "name": "/job:localhost/replica:0/task:0/device:GPU:0 Compute" 而GPU的执行流编号就对应 "tid": 14, 所以,问题就变成了:这里的“tid”能否表示实际在GPU中的执行流?我理解是不能的,在使用GPU做backend的场景,HorovodAllreduce底层调用NcclAllReduce,而NcclAllReduce本质是DMA引擎在做实际工作,而以DMA的工作方式,怎么会占用GPU的Compute流呢,所以,这里应该比较准确的理解应该是,TODO:在某一时刻,不断开启新的执行流去执行一个GPU Op,至于这些执行流是否真的在GPU内并行,这里并不能表达,需要使用GPU专属的工具。
  3. 并行的欺骗性,本质上,这种误导是由于1和2导致的Timeline局限性,倒也称不上是错误。上图中使用Horovod进行分布式训练时,由于多个AllreduceOp并无相关性,也分属不同的执行流,看起来并行度很高,其实在其内部实现中,一方面Horovod会做融合,另一方面每个HorovodAllreduceOp只有一个后台线程真正的去做网络通信,最后,网络IO的能力上限取决于网卡,当网卡达到瓶颈时,更多的并行HorovodAllreduceOp是没有意义的。
  4. 和Step强相关,取一个Timeline的时候,通常会指定运行到某个Step来保存,这里的问题是,不同的Step可能执行的Timeline差异巨大,典型的问题就是开启XLA时,部分子图只会在Step1执行,同时,Step1进行的XlaCompile也会大大延长单Step时间,所以在测试的时候,使用的Step也是要小心设置的。

Tensorboard

Tensorboard用于传统TF Graph调试是比较容易理解的,这里讨论如何使用Tensorboard进行XLA的可视化。XLA可以将优化的各个阶段的子图导出为pbtxt或pb|dot,html等格式,如果是pbtxt,可以直接点击Tensorboard左侧Choose File 按钮上传。

执行”./bazel-bin/tensorboard/tensorboard –logdir=”./logs” –port=9000″,浏览器输入”http://localhost:9000/“,点击上传mark_for_compilation_annotated_1.pbtxt即可看到子图:

Dot + Graphvis

XLA的HLO PASS可以生成pb, html,dot等格式,这里的pb用import_pb_to_tensorboard或者自己写的脚本都没能成功转化成pbtxt,还没办法使用的Tensorboard查看,但相应的dot格式,显然可以使用Graphviz进行可视化。但鉴于graphviz的全量计算的方式,对于特别大神经网络,dot工具本身就会成为瓶颈。

Continue reading

Tensorflow 编译方法

Tensorflow使用SWIG自动根据so文件生成相应的python接口, 这就引起一系列反常识的问题, 比如, TF的编译环境同时要求TF部分可运行, 通常来说, 编译环境只需要配置编译依赖即可, 但在TF中行不通, 举个例子: 如果环境中只有CUDA, 没有GPU以及可用的Driver, 那么在编译的最后阶段, SWIG机制会去运行Tensorflow Runtime来生成Python接口, 但缺失GPU或Driver的环境是无法运行TF的, 也就无法完成位于编译最后阶段的python接口生成. 此外, 这也导致了如果在SWIG调用路径上的Runtime代码一旦有逻辑问题而不是语法问题, 在Runtime的编译过程不会报错, 编译最后生成Python接口也会出错导致最终编译失败, 进一步导致该类逻辑错误难以调试.

Debug信息保留

将编译命令调整如下

第一处--output_user_root="/.cache/bazel/" 主要解决编译环境root一般较小, 无法承载编译临时文件
第二处 --copt=-g  --strip=never  用于保留调试信息, 要配合python3-dev使用, 如果使用python2, 由于带有debug信息生成so过大, 会导致wheel制作失败(int32溢出), 当然, 相应的python代码也要适配python3进行修改

离线安装依赖

Tensorflow编译过程会从Internet上load所需依赖, 如果网络环境不允许, 可以通过其他路径现将依赖包下载copy到编译环境, 再修改tensorflow/tensorflow/workspace.bzl中的相关url, 即可规避联网编译问题

https://github.com/tensorflow/tensorflow/issues/5428#issuecomment-258732841

Tensorflow OpKernel机制详解

OpKernel是Op的具体实现, tf中已经实现的tfop的OpKernel在源码中的tensorflow/core/framework/kernel/中, OpKernel通过注册时使用相同的名字将自己和相应的Op联系到一起.

在tf中, OpKernel进一步可以分为两类, OpKernelAsyncOpKernel:

  1. OpKernel是同步执行的, 即"Compute()"返回即认为数据已经被正确处理, 注册OpKernel, 子类需要重写其Compute()方法.
  2. AsyncOpKernel是对OpKernel的封装, 顾名思义, AsyncOpKernel执行返回并不意味着数据已经被处理完毕, 数据的真正被处理完毕时通过回调的方式通知Op执行引擎, 注册一个AsyncOpKernel, 子类需要实现"AsyncCompute()"而不是Compute().

接口形式

无论是哪种OpKernel, 均使用"REGISTER_KERNEL_BUILDER()"注册到运行核心.

注册原理

注册机制的实现代码主要集中在tensorflow/core/framework/op_kernel.h(.cc). 与 Optimization以及Op在注册时直接构造一个static OptimizationPassRegistration(OpRegistrationData)对象的机制略有不同, OpKernel的通过一些trick实现了对OpKernel的延迟构造, 即"REGISTER_OP_KERNEL_BUILDER()"并没有直接构造一个"OpKernel"实例, 而是构造一个" static ::tensorflow::kernel_factory::OpKernelRegistrar "对象, 并借由该构造过程构造并注册一个"KernelRegistration" 对象到 global_regsitry, 该构造过程接受上层传入的, 用于new一个OpKernel的"[](::tensorflow::OpKernelConstruction* context) -> ::tensorflow::OpKernel* { return new __VA_ARGS__(context);}" 函数, 在上层真正需要这个OpKernel的时候, 才会通过一系列调用最终执行该"create_fn()/lambda"来构造一个实实在在的OpKernel对象. 即整体上不再是Registry->Registration(Optimization/Op对象), 而是 Registry->Registrar->Registration->在需要时create_fn()构造OpKernel对象.

构造一个OpKernelRegistrar:

Continue reading

Tensorflow Op机制详解

注册一个tfop分为两部分: OpOpKernel. 其中, Op是tfop的声明部分, 类似于函数的声明, 主要描述Op静态属性. OpKernel是tfop的实现部分, 同样类似于函数的实现, 主要描述OpKernel的具体计算逻辑. tf已经实现的Op相关源码在tensorflow/core/common_runtime/ops/中, OpKernel的部分在tensorflow/core/common_runtime/kernels/. 二者通过注册时使用的名字联系到一起.

接口形式

下面是一个注册Op的例子, 使用接口"REGISTER_OP()"‘注册一个Op,包括输入输出, Op属性等内容, 注册的Op最终会转换为OpDef, 即用protobuf格式存储的静态数据, 所以, 这里用到的属性也按照pb的格式以KV的形式编写. REGISTER_OP的实现在”tensorflow/core/framework/op.h OpDefBuilderWrapper” 中, 众多可设值中,"Attr"是一个允许自定义的值, 比如XLA引擎就根据自身需求提供了"XlaCompile", 如果一个Op将该值设置为true, 就会强制XLA引擎将其编译. 当然, 也可以设置一些无用的值, 就像函数声明里有一个并没有实际使用的参数, 除了浪费存储空间没有其他用途.

注册原理

在内部原理上, 上述代码可以进一步分解为两部分: 构造一个Op + 注册一个Op, 这部分代码主要在 “tensorflow/core/framework/ op.h(.cc) op_def_builder.h(.cc) ” 中

Continue reading

Tensorflow OptimizationPassRegistry机制详解

整个流水线大概是:Init Graph –> Grappler进行全图优化 –> 根据Device将Graph拆成 Partition –> OptimizationPassRegistry优化各个Partition –> 图执行. 同其他注册机制一样, OptimizationPassRegistry也使用的registry, registertion以及registerar等概念.

源码中, 使用REGISTER_OPTIMIZATION()注册一个优化器, 具体实现如下

-144-new了一个我们注册的optimization对象并用unique_ptr指向它, 这个unique_ptr就是registry管理的对象, 通过它间接管理相应的optimization. 注册的本质是返回一个静态的, 类型为'OptimazationPassRegistration'的, 名为'register_optimization_##ctr'的对象, 这里使用了C++预编译宏'__COUNTER__'生成唯一变量名

下面是一个OptimazationPassRegistration对象的构造过程, 可以看出, 就是将我们构造的'register_optimization_##ctr注册到全局的'global_optimization_registry'.

Continue reading

Tensorflow 图计算引擎概述

tensorflow的用户可以使用多种语言来构造的自己的图, 但各种语言的API最终都会经由C API 进入tensorflow 运行时. 可以说, 对于运行时代码, 其上边界就是C API. 比如, 通过python描述的一张网络, 就是通过类似下面的几个python-C接口进入运行时的.

tensorflow整体上可以看做一个”图语言的编译器”, 和所有编译器的优化以及翻译的功能类似, Graph在运行时中的处理过程, 可以分为 图构造->图优化->图执行 几个阶段. 其中, 图优化随同图构造一同被执行.

全图构造及其优化

Session初次构造时, 应用层代码中定义的数据流图转换成GraphDef格式, 经由C API传入DirectSession.Extend(), 参考调用栈如下

Continue reading

BERT 模型参数量估计

根据BERT论文, 其12层transformer结构有110M参数, 24层更是高达340M, 虽然google公开了这两个网络的预训练模型, 用户只需在后面加一层Full-connected 但如果自己去做BERT预训练, 到底要耗费多少显存呢?

tensorflow中使用tf.variable()在模型中生成训练参数, 同时, tf.dense()内部也含有weight和bias两类参数. 比如,

生成了shape是[vocab_size, embedding_size]的variable, 即 vocab_size * embedding_size 个参数,

内含 from_tensor_2d.shape[-1] * to_tensor_2d.shape[-1] 个weight 以及 to_tensor_2d.shape[-1] 个bias. 共计 (from_tensor_2d.shape[-1] + 1) * to_tensor_2d.shape[-1] 个参数, 这就是一个FC-Layer的参数量

有了上面的例子, BERT的参数量就不难分析, 整个BERT网络可以分为4个模块: embedding layer + transformer layers + pooled layer+ classifier layer

embedding-layer

Continue reading