Month: August 2019

Tensorflow XLA Debug/Profiling Methods

1079  Jiang XIAO  
(last edited )

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

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

Timeline

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

python tf_cnn_benchmarks.py --trace_file="time.json"  --save_summaries_steps=10 --summary_verbosity=1 --train_dir "./tmp"  --num_gpus=1 --batch_size=1 --model=alexnet
mpirun --allow-run-as-root -np 4 -H 10.54.32.198:4 python tf_cnn_benchmarks.py --xla="True" --trace_file="time.json" --num_gpus=1 --batch_size=32 --model=alexnet --variable_update=horovod

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 按钮上传。

python tf_cnn_benchmarks.py --trace_file="time.json" --xla="True" --save_summaries_steps=10 --summary_verbosity=1 --train_dir "./tmp"  --num_gpus=1 --batch_size=1 --model=alexnet

执行”./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工具本身就会成为瓶颈。

foo $ dot -Tsvg -o module_0002.after_optimizations.svg module_0002.after_optimizations.dot
Continue reading →

Tensorflow 编译方法

1200  Jiang XIAO  
(last edited )

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

Debug信息保留

将编译命令调整如下

bazel --output_user_root="/.cache/bazel/" build --config=cuda --config=nccl --copt=-g  --strip=never //tensorflow/tools/pip_package:build_pip_package

第一处`–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

tf_http_archive(
    name = "eigen_archive",
    build_file = clean_dep("//third_party:eigen.BUILD"),
    patch_file = clean_dep("//third_party/eigen3:gpu_packet_math.patch"),
    sha256 = "0dde8fb87f5dad2e409c9f4ea1bebc54e694cf4f3b633081b0d51a55c00f9c9f",
    strip_prefix = "eigen-eigen-a0d250e79c79",
    urls = [
        "http://mirror.tensorflow.org/bitbucket.org/eigen/eigen/get/a0d250e79c79.tar.gz",
        #"https://bitbucket.org/eigen/eigen/get/a0d250e79c79.tar.gz",
        "file:////tmp/sources/tensorflow/a0d250e79c79.tar.gz",
    ],
)

Tensorflow OpKernel机制详解

1033  Jiang XIAO  
(last edited )

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

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

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

接口形式

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

//tensorflow/core/common_runtime/kernels/nccl_ops.cc
#include "third_party/nccl/nccl.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/nccl/nccl_manager.h"
namespace tensorflow {
class NcclAllReduceOpKernel : public AsyncOpKernel {
 public:
  explicit NcclAllReduceOpKernel(OpKernelConstruction* c)
      : NcclReduceOpBase(c) {}

  void ComputeAsync(OpKernelContext* c, DoneCallback done) override {
    //...
  }
};
REGISTER_KERNEL_BUILDER(Name("NcclAllReduce").Device(DEVICE_GPU),
                        NcclAllReduceOpKernel);  
}

注册原理

注册机制的实现代码主要集中在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机制详解

1026  Jiang XIAO  
(last edited )

注册一个tfop分为两部分: Op和OpKernel. 其中, 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引擎将其编译. 当然, 也可以设置一些无用的值, 就像函数声明里有一个并没有实际使用的参数, 除了浪费存储空间没有其他用途. 

//tensorflow/core/common_runtime/nccl_ops.h
#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"

namespace tensorflow {

using shape_inference::InferenceContext;
using shape_inference::ShapeHandle;

REGISTER_OP("NcclAllReduce")
    .Input("input: T")
    .Output("data: T")
    .Attr("reduction: {'min', 'max', 'prod', 'sum'}")
    .Attr("T: {half, float, float64, int32, int64}")
    .Attr("num_devices: int")
    .Attr("shared_name: string")
    .Attr("XlaCompile: bool=true")
    .SetIsStateful()
    .SetShapeFn(shape_inference::UnchangedShape);
}

注册原理

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

Continue reading →