Chenfan Blog

Do cool things that matter.

0%

TensorFlow 拆包（七）：Profiling 踩坑 & Benchmark

Posted on Edited on In

接上篇：

开始分析性能瓶颈了，本篇记录一下研究 TF 中自带的 Profiling 工具时遇到的几个坑点。

profiler

大概 17 年 5 月左右，/tensorflow/core/ 中新加了一个 profiler 的目录，里面是把原本在 contrib 中的 profiling 工具移过来了，大概正式 release 应该是在 1.6、1.7 里面。

关于生成 profiling 的 context 文件详见 tf.profiler 相关的内容，这里直接开始记录怎么用 tfprof 这个工具。

试了一下 pip 包里面应该是没有单独包含的，需要从源码手动编译：

1
bazel build //tensorflow/core/profiler

然后使用也是要从 bazel-bin 的目录中打开：

1
bazel-bin/tensorflow/core/profiler/profiler --profile_path=xxxxx

profiler_ui

profiler 的 README 中，示例代码除了 profiler 以外还有个 profiler_ui，基本上是一个类似 tensorboard 的网页前端，方便调用后端的 profiler 进行可视化查看用的。

这里虽然写着暂未开源，但是在 TensorFlow 的 github 总目录里面可以找到一个叫 profiler-ui 的项目，就是那个未完善开源的 ui 版了。

看了下，安装需要用到 go 以及 Google 自家的 pprof 工具，可能是因为耦合的其他部件比较多，所以暂时还没有并入 TF 的主代码中去。不过这里的 Installation 已经足够我们自己装上了。

装 pprof 的时候会有个坑点，CentOS 库中可以找到 gperftools 这个工具，也是 Google 提供的，yum 装上之后可执行文件的名字也叫 pprof ！！但是跟这里用到的 pprof 不是一个玩意！！

之后按照示例上的说明：

1
python ui.py --profile_context_path=xxxx

即可启用。

在我尝试使用它的时候，距离这个库上一次 git 的更新已经过去 1 个月左右了，不知道是 python 版本还是什么原因，直接运行可能会遇到找不到 server 的路径等等的 bug，直接在 ui.py 里面稍微改一下就好。

Profiling

运行 TF 时保存出来的 profiling 文件包含了大量信息，主要有几个方面：

  • scope：应该是 python 层代码中用 tf.name_scope() 包起来的视图
  • graph：TensorFlow 计算图的视图
  • op：把 TensorFlow 计算图再细化一层
  • code：Python 代码视图

默认会按列表把所选的视图中的一些信息给输出出来，另外用-output 选项可以指定输出成另外的格式：

1
2
3
4
5
6
7
8
9
10
tfprof>
xxx xxx -output timeline:outfile=xxxxx
# 把结果输出成 chrome 用的时间线 trace 文件，可以在 chrome 地址栏中输入 chrome://tracing 打开
# 只支持 graph、scope、code 这 3 种视图
xxx xxx -output pprof:outfile=xxxxx
# 把结果输出成 pprof 用的可视化文件（所以前面装 pprof 就是为了这个）
# 只支持 code 这种视图
# --------------------------------------------------------------------------
# pprof 可视化文件之后可以用 pprof 来变成图片（猜测大概是类似 GraphViz 的数据结构）
pprof -svg --nodecount=10000 --sample_index=1 xxxxxx.prof > xxxxxx.svg

profiler_ui 打开时的第一个页面就是 graph 视图生成的 timeline：

其中包含了计算图中每个 node 在卡上的情况，运行时间、数据流动依赖关系等等。（话说显示的太复杂了，事实上我觉得还是很难看）

然后默认的 scope 视图以及 code 视图得到的 timeline 我也感觉并没有什么用。

code 视图输出成的 pprof 图片倒是还可以看一下，但是感觉用处也不大

所以最后感觉还是不知道该怎么用好这套 profiling 工具

Options

在 tfprof 界面直接回车可以看到默认的选项，然后这里面的内容都是可以改的：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
tfprof>                             
-max_depth                  10      
-min_bytes                  0       
-min_peak_bytes             0       
-min_residual_bytes         0       
-min_output_bytes           0       
-min_micros                 0       
-min_accelerator_micros     0       
-min_cpu_micros             0       
-min_params                 0       
-min_float_ops              0       
-min_occurrence             0       
-step                       -1      
-order_by                   name    
-account_type_regexes       .*      
-start_name_regexes         .*      
-trim_name_regexes                  
-show_name_regexes          .*      
-hide_name_regexes                  
-account_displayed_op_only  false   
-select                     micros  
-output                     stdout:

稍微挑几个写一下：

-max_depth：指定显示前多少个 node（配合下面的 -order_by ？）

-step：profiling 记录的文件可能包含了很多个 step，用这个选项来指定当前分析哪个 step 的信息，默认 -1 是对所有 step 做平均

-order_by：打出来列表的时候，按照什么来排序：

  • name：node 的名称
  • depth：node 在节点树中的深度
  • bytes：占用的内存数
  • peak_bytes：占用的峰值内存数
  • residual_bytes：计算完成之后，还剩下不释放的内存数
  • output_bytes：输出的大小
  • micros：node 计算所花费的时间
  • accelerator_micros：node 计算所花费的加速卡时间（区别于 CPU 的其他设备）
  • cpu_micros：node 计算所花费的 CPU 时间
  • params：node 中包含的参数量
  • float_ops：node 所需要的浮点运算次数
  • occurrence：node 在图中出现的次数

-account_type_regexes：筛选出类型里面带有某些前缀的 node 有多少个

-start_name_regexes：筛选出名字中带某些前缀的 node

-trim_name_regexes：隐藏掉名字中带某些前缀的 node

-show_name_regexes：筛选出名字中带某些字符的 node

-hide_name_regexes：隐藏掉名字中带某些字符的 node

-select：选择视图中的哪些内容（有点像从数据库里面找东西的感觉），输出 timeline 的时候配合这个应该能够得到不同的数据：

  • bytes：占用的内存数
  • peak_bytes：占用的峰值内存数
  • residual_bytes：计算完成之后，还剩下不释放的内存数
  • output_bytes：输出的大小
  • micros：计算所花费的时间
  • accelerator_micros：计算所花费的加速卡时间
  • cpu_micros：计算所花费的 CPU 时间
  • params： 参数量
  • float_ops：浮点运算次数
  • occurrence：在计算图中出现的次数
  • tensor_value：tensor 数据的值（估计需要配合 checkpoint 用）
  • device：op 放在哪个设备上
  • op_types：op 类型
  • input_shapes：输入的形状

Trace

抛开上面那个目前还没有正式 Release 的 Profiling 接口不说，实际可以用来做分析的是一套生成 trace_file 的 API。

用法也很简单：

1
2
3
4
5
6
7
8
9
10
11
options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()

_ = sess.run(optimizer, options=options, run_metadata=run_metadata)

fetched_timeline = timeline.Timeline(run_metadata.step_stats)
chrome_trace = fetched_timeline.generate_chrome_trace_format()

with open(FLAGS.trace_file, 'w') as f:
    f.write(chrome_trace)
print('Chrome Trace File write in %s' % FLAGS.trace_file)

在 RunOptions 中设置好追踪的级别，然后作为参数一起参与 Session.run()，最后记录得到的每个 step 的追踪数据通过 run_metadata 的结构返回出来。通过对追踪结果的解析即可生成我们可以理解的图形数据了，这个用的是 chrome 支持的 json 格式，在 chrome 地址栏中输入 chrome://tracing/ 即可很方便地查看，timeline 最后出来的效果跟上面的是一致的。

应该说，前面这个 Profiling 的 API 应该底层封装的也是这套机制。

在 DirectSession 中可以非常容易地找到与 trace_level 相关的代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
  std::unique_ptr<DeviceTracer> tracer;
  if (run_options.trace_level() >= RunOptions::HARDWARE_TRACE) {
    tracer = CreateDeviceTracer();
    // tracer may be NULL on platforms without accelerators.
    if (tracer) {
      Status s = tracer->Start();
      if (!s.ok()) {
        run_state.executors_done.Notify();
        delete barrier;
        return s;
      }
    }
  }
...
  if (tracer) {
    TF_RETURN_IF_ERROR(tracer->Stop());
    TF_RETURN_IF_ERROR(tracer->Collect(run_state.collector.get()));
  }

其中 DeviceTracer 是一个预留给多种设备来方便进行性能分析的接口，可惜的是目前里面的实现只有 GPU 的，需要依靠 CUDA 提供的 CUPTI 库。所以大概追踪过程中得到的与 CPU 相关的信息应该也是 CUPTI 附带的，如果是纯 CPU 版本的 TensorFlow，CreateDeviceTracer() 直接返回的是一个空指针。

Distributed

由于分布式环境下的 Session 的执行模式与单机情况下有所不同，因而分布式下运行 trace 的工作方式也会有所区别。

MasterSession 中首次执行 PartialRun 时会初始化 PerStepState：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
// If this is the first partial run, initialize the PerStepState.
if (!run_state->step_started) {
  run_state->step_started = true;
  PerStepState pss;

  const auto count = run_state->count;
  pss.collect_timeline =
      req.options().trace_level() == RunOptions::FULL_TRACE;
  pss.collect_rpcs = req.options().trace_level() == RunOptions::FULL_TRACE;
  pss.report_tensor_allocations_upon_oom =
      req.options().report_tensor_allocations_upon_oom();

  // Build the cost model every 'build_cost_model_every' steps after skipping
  // an
  // initial 'build_cost_model_after' steps.
  const int64 build_cost_model_after =
      session_opts_.config.graph_options().build_cost_model_after();
  const int64 build_cost_model_every =
      session_opts_.config.graph_options().build_cost_model();
  pss.collect_costs =
      build_cost_model_every > 0 &&
      ((count + 1 - build_cost_model_after) % build_cost_model_every == 0);
  pss.collect_partition_graphs = req.options().output_partition_graphs();

  std::unique_ptr<ProfileHandler> ph = run_state->rcg->GetProfileHandler(
      run_state->step_id, count, req.options());
  if (ph) {
    pss.collect_timeline = true;
    pss.collect_rpcs = ph->should_collect_rpcs();
  }

  run_state->pss = std::move(pss);
  run_state->ph = std::move(ph);
}

这里会根据 trace_level 的值来设置一些标记。

下一步，pss 中的内容又会被写到 exec_opts 结构中：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
// Collect execution cost stats on a smoothly decreasing frequency.
ExecutorOpts exec_opts;
if (pss->report_tensor_allocations_upon_oom) {
  exec_opts.set_report_tensor_allocations_upon_oom(true);
}
if (pss->collect_costs) {
  exec_opts.set_record_costs(true);
}
if (pss->collect_timeline) {
  exec_opts.set_record_timeline(true);
}
if (pss->collect_rpcs) {
  SetRPCLogging(true);
}
if (pss->collect_partition_graphs) {
  exec_opts.set_record_partition_graphs(true);
}
if (pss->collect_costs || pss->collect_timeline) {
  pss->step_stats.resize(partitions_.size());
}

这个结构会被封装在 rpc 的 call 中发送给 WorkerService 来处理。

在 Worker 的运行结构中，可以看到这样的代码：

1
2
3
4
5
6
7
StepStatsCollector* collector = nullptr;
if (request->exec_opts().report_tensor_allocations_upon_oom() ||
    request->exec_opts().record_timeline() ||
    request->exec_opts().record_costs()) {
  collector = new StepStatsCollector(response->mutable_step_stats());
  // TODO(mrry,pbar): GPU tracing for distributed steps.
}

好吧，StepStatsCollector 都已经创建了，但是可惜后续具体的 GPU tracing 部分还没有往里面完善。

Benchmark

TensorFlow 官方的 Performance 页 和 Benchmarks 页 中给出了官方测性能用的 benchmark 脚本：

基本上把目前 TF 里最高效的 API 都用上了，并且包含了各种常见的多机多卡方案，很值得作为高效的样例脚本来参照。

有个问题是这个库目前没有定期 release，各种更新全都合并到 master 分支里面去了，然后随着 TF 版本的不断更新，它的 master 分支是跟着 TF 的 master 分支走的。

因此要想正常跑最新的 benchmarks，就需要装 tf-nightly-gpu 包或者源码编译一个比较新的 TensorFlow 分支。

简单分析一下这份脚本的结构。

从 tf_cnn_benchmarks.py 这个入口进去之后，核心的执行流程在 benchmark_cnn.py 中。

只测试前向走 BenchmarkCNN._eval_cnn()，测试训练全过程走 BenchmarkCNN._benchmark_cnn() 。

前面的 FLAG 解析什么的直接略过，从训练部分开始看。

_benchmark_cnn

首先构建计算图：(image_producer_ops, enqueue_ops, fetches) = self._build_model()

image_producer_ops 是处理输入数据的部分，enqueue_ops 涉及到计算图中的流水线队列，最后的 fetches 是等一下 sess.run() 中的目标 op。

设置 tf.summary 以及 tf.train.Saver 等等，Saver 中传入的是 variable_mgr.savable_variables()。

创建 tf.train.Supervisor 时同时完成变量初始化，初始化 op 组包含：

  • tf.local_variables_initializer 初始化本地变量
  • tf.tables_initializer 初始化用到的各种表（哈希表等等）
  • 本地变量初始化之后，执行variable_mgr.get_post_init_ops() 完成自定义的一些初始化执行动作，这个部分要根据不同的参数维护算法来定
  • 如果有同步用的队列 barrier，也一起在这里完成初始化

之后 sv.managed_session 开始真正的执行循环：

1
2
3
4
5
6
7
8
9
with sv.managed_session(...) as sess:
    如果使用的是真实数据，则往 enqueue_ops 中插入提取数据到队列的 op
    初始化 global_step
    创建一个 global_step_watcher，在新线程中监控 global_step 的变动情况
    while not done_fn():     这个函数与 global_step_watcher 相关，用于控制训练循环什么时候结束
        ...
        benchmark_one_step() 训练一个 step
        ...
    后续再处理一些收尾内容

_build_model

回到第一步看一下计算图的构建部分。

(image_producer_ops, image_producer_stages) = self._build_image_processing(shift_ratio=0) 创建输入数据。

1
2
3
4
对于当前设备上的每一块 GPU，variable_mgr.create_outer_variable_scope() 创建命名域：
    add_forward_pass_and_gradients()    添加网络的前向部分，并且计算得到梯度
    根据当前的任务是训练还是预测，处理准备网络中需要返回的内容
    从计算图中提取出 Batch Normalization 的更新部分，添加到 0 号卡的更新部分中，BatchNorm 只需要一块卡来计算

如果图中用了 staging_area 的数据组织方式，这里另外再添加一下，扩充 enqueue_ops。

fetches = self._build_fetches() 最终收集前面所有的信息，构建出等一下需要传入 sess.run() 中去的目标

完成前面的内容后，把 image_producer_ops，enqueue_ops，fetches 三部分内容返回给上一层的函数。

add_forward_pass_and_gradients

创建随机数据作为输入，或者处理传入的数据产生器。

logits, aux_logits = self.model.build_network() 构建完整的前向网络。

添加输出结果以及计算 loss 误差。

variable_mgr.trainable_variables_on_device() 获取当前 GPU 上所有的可训练参数。

如果当前是最后一块 GPU 卡，那么再额外计算 L2_loss，添加到前面的 loss 中去，L2_loss 只需要计算一次。

grads = tf.gradients(scaled_loss, params, aggregation_method=aggmeth) 根据前面收集的当前 GPU 上的可训练参数信息构建反向的梯度计算图，返回得到的是图中所有的梯度。

接下来再获取一次 variable_mgr.trainable_variables_on_device() ，然后把得到的参数与前面的梯度打包在一起返回回去，准备接下来的参数更新。需要注意的是，第一次调用 trainable_variables_on_device 时传入了一个 writable=False 的参数，这里传入的是 writable=True，在某些特别的多卡参数管理算法中，用于梯度计算和最终梯度更新写回的目标是不一样的。

所有前面的这些都封装在一个 results 的 dict 中返回回去。

_build_fetches

这里算是计算图构建的收尾部分了，传入的内容是所有 GPU 上计算图的合集。

variable_mgr.preprocess_device_grads() 预处理出需要在哪些设备上执行梯度更新操作。

1
2
3
4
5
6
对于梯度更新设备中的每一块 GPU：
    tf.reduce_mean()                          计算前面所有卡上梯度的平均值
    variable_mgr.get_gradients_to_apply()     获取有哪些梯度是要在当前设备上更新的
    get_learning_rate()                       计算学习率
    get_optimizer()                           获取梯度更新用的 Optimizer
    variable_mgr.append_apply_gradients_ops() 应用 Optimizer 进行梯度更新

把前面所有的东西打包在 fetches 这个 dict 中返回回去。

VariableMgr

可以看到上面有很多核心的操作都是通过 variable_mgr 结构完成的，这套脚本定义了一个 VariableMgr 类，想要自己修改参数管理、更新的算法只需要重写这里面的一些函数即可。

前面出现过的比较有用的几个接口函数：

  • def create_outer_variable_scope(self, device_num)

封装变量命名域，主要用于维护变量创建时要做的事情，一般情况下直接返回一个普通的 tf.variable_scope，需要对变量创建进行额外操作的话需要自己构造一个 custom_getter 作为参数传入tf.variable_scope。

  • def preprocess_device_grads(self, device_grads)

预处理出需要做梯度更新操作的设备，以及对应设备上的梯度和参数。

  • def get_gradients_to_apply(self, device_num, gradient_state)

与上一个函数对应使用，用于找出每个设备需要处理哪些参数更新任务。

  • def append_apply_gradients_ops(self, gradient_state, opt, grads, training_ops, loss_scale_params)

在设备上针对每一对需要更新的变量及其梯度，应用 apply_gradients 操作。

  • def get_post_init_ops(self)

用于额外附加一些在所有变量完成初始化之后，开始训练之前，需要执行的操作。

  • def get_devices(self)

返回当前节点中可用的 GPU 列表，在某些 PS-WORKER 的实现方式中，返回的是 tf.replica_device_setter 的封装。

  • def savable_variables(self)

返回哪些变量是需要被 tf.Saver 保存进检查点的。

  • def trainable_variables_on_device(self, rel_device_num, abs_device_num, writable=False)

返回当前设备上的可训练参数（即能计算梯度，可以进行反向更新的参数）。输入的两个 device_num 分别是 GPU 在当前节点中的编号以及在全局环境中的编号。

writable 用于标识需要被写回更新的参数，在有些情况下图中可能存在多份参数备份，writable 为 False 时返回的是图中用于求梯度以及构建反向数据通路用的参数，为 True 时返回的是等一下 apply_gradients 需要应用梯度更新操作的参数。

replica_device_setter & variable_scope-custom_getter

前面建图时用到的两个很重要的接口，用于额外处理 op 在设备上的分配操作。

benchmarks 脚本中的用法大概是这样：

1
2
3
4
5
6
for device_num in range(len(self.devices)):
    with self.variable_mgr.create_outer_variable_scope(device_num):
        # ...
        with tf.device(self.devices[rel_device_num]):
        # self.devices[] 里面是实现设好的 tf.replica_device_setter
            #...build_network

对节点中的每一块 GPU 卡，首先套上一个 variable_scope（里面可能会使用到 custom_getter），在构建 op 时再套一层 replica_device_setter。

tf.replica_device_setter 需要配合 tf.device 使用，作用范围是其 python 作用域以内的所有 op，这个函数简单地说就是对传入的 op 进行判断，如果是计算型的 op 就正常分配在运算设备上，如果是需要在 PS-WORKER 之间共享的参数型 op 则需要在参数服务器上。它的返回值是需要分配给的 device 的名字，所以直接用 tf.device 指定即可。

1
2
3
def replica_device_setter(ps_tasks=0, ps_device="/job:ps",
                          worker_device="/job:worker", merge_devices=True,
                          cluster=None, ps_ops=None, ps_strategy=None)

具体的源码实现上，主要是对新创建 op 的类型进行判断，如果在 ps_ops 包含的范围内（为 None 时会用一个 STANDARD_PS_OPS 来作为检查范围）则用某种 ps 分配策略放到参数服务器上，否则放到默认的计算设备上。

默认的 ps_strategy 不指定的话就是用的 round-robin，简单地说就是按顺序依次分。

tf.variable_scope 中 custom_getter 的作用范围就只限于作用域以内所有的 tf.get_variable 调用了（注意，必须是 tf.get_variable，这个对 tf.variable 是无效的）。前面 replica_device_setter 只是指定了参数存放的位置，这里则可以对参数创建进行更多的改动。

例如 StagedVariableGetter 做的事情就是把变量封装上一层 StagingArea，计算图中需要读取变量的时候返回一个 StagingArea.get，对于 apply_gradient 这种需要修改变量本身的操作，则返回参数本体（也就是前面看到的 writable 这个参数起作用的方式）。

VariableMgr instances

官方的 Benchmark 脚本中提供了 8 种内置的 VariableMgr 实例。

VariableMgrIndependent

不同卡之间完全不作数据交互，单纯用来测单机多卡的理论计算速度用。

不需要封装 custom_getter 和 replica_device_setter。

VariableMgrLocalFetchFromPS

多卡中的参数统一存储，不同卡在计算时直接从统一的 PS 中读取需要的数据。

不需要封装 custom_getter。

get_device 这里：

1
2
3
4
5
6
7
8
9
10
11
12
if self.benchmark_cnn.local_parameter_device_flag == 'gpu':
  return [
      variable_mgr_util.ParamServerDeviceSetter(d, raw_devices)
      for d in raw_devices
  ]
else:
  return [
      tf.train.replica_device_setter(
          worker_device=d,
          ps_device=self.benchmark_cnn.param_server_device,
          ps_tasks=1) for d in raw_devices
  ]

如果参数存放在 CPU 上，直接对每个 GPU 设备返回一个指定好 ps_device 的 replica_device_setter。

如果选择参数存放在 GPU 上，这里的做法是将所有参数均衡负载平分在各块卡上。

VariableMgrLocalFetchFromStagedPS

多卡中的参数统一存储，相比之前的增加了 StagingArea 的流水线操作。

custom_getter 中为每个变量额外创建了一个 StagingArea，计算图中需要读取变量的时候返回对应的 StagingArea.get。

trainable_variables_on_device 中 writable 为 True 时，返回变量本体，否则返回对应的 StagingArea.get。

其他部分与上一种方式相同。

VariableMgrLocalReplicated

每块卡上的计算图完全独立，各卡都是自己存储自己的参数，梯度更新的时候再采用某种 Allreduce 的算法对各卡上的参数做统一规约。

get_post_init_ops 在初始化完成后拷贝 GPU0 上的参数到其他卡上覆盖掉，保证所有卡的初始参数一致。

不需要封装 custom_getter 和 replica_device_setter。

preprocess_device_grads 中返回的梯度是调用某种规约算法去综合所有卡上的梯度值，之后再跟本地的参数一起交给 apply_gradient 去更新即可。

因此这里的计算流程是，初始所有卡上参数一致，训练完一步之后规约梯度，规约完成后所有卡上得到的梯度也都一致了，再 apply 更新到本地的卡上，这样下一步开始时所有卡上的参数仍然是一致的。

VariableMgrDistributedAllReduce

用于分布式。

这是脚本中唯一一种需要用到 single_session 的模式，基本上跟 Replicated 的方式一致，每块卡上都独立存数据，更新时全局规约，特殊点在于这种方式只需要由一个 python 进程来启动，所有 worker 上的图构造等等都是由一个 controller 的角色完成，其他所有的 worker 都像平时的 ps 一样 join_server 即可。

大体实现上跟上一种一致

VariableMgrDistributedFetchFromPS

用于分布式。

大体实现跟单节点的 FetchFromPS 一致。

custom_getter 使用了 OverrideCachingDevice，虽然由 replica_device_setter 指定好了所有参数都保存在 ps 上，但是在 worker 还可以做一次数据缓存。caching_device 这个参数与 tf.get_variable() 中的参数对应，即 worker 端的多块卡从远程 ps 获取数据只在第一块卡拉取数据时通过网络取一次，后续的几次直接从缓存中读取。缓存数据的分配方案跟单节点 ps 时 CPU/GPU 上存储参数的方案类似。

VariableMgrDistributedFetchFromStagedPS

用于分布式。

在上一种的基础上加上了 StagingArea。

VariableMgrDistributedReplicated

用于分布式。

计算流程其实跟 DistributedAllReduce 是一致的，大体上跟前面类似实现相一致。

Related Posts