[原始碼解析] 深度學習分散式訓練框架 horovod (7) --- DistributedOptimizer
目錄
0x00 摘要
Horovod 是Uber於2017年釋出的一個易於使用的高效能的分散式訓練框架，在業界得到了廣泛應用。
本系列將通過原始碼分析來帶領大家瞭解 Horovod。本文是系列第七篇，看看 Horovod 如何與 TensorFlow 融合。
前面幾篇連結如下：
[原始碼解析] 深度學習分散式訓練框架 Horovod (1) --- 基礎知識
[原始碼解析] 深度學習分散式訓練框架 horovod (2) --- 從使用者角度切入
[原始碼解析] 深度學習分散式訓練框架 horovod (3) --- Horovodrun背後做了什麼
[原始碼解析] 深度學習分散式訓練框架 horovod (4) --- 網路基礎 & Driver
[原始碼解析] 深度學習分散式訓練框架 horovod (5) --- 融合框架
[原始碼解析] 深度學習分散式訓練框架 horovod (6) --- 後臺執行緒架構
我們需要一些問題或者說是設計要點來引導分析，而且因為讀者可能沒有看過本系列其他文章，因此問題點會和其他文章有部分重複：
-
第一個技術難點是：Horovod 如何從 TF 的執行流程中獲取到 梯度（gradients）進行處理？
- 在 TensorFlow 1.x 中，深度學習計算過程被表示成為一個計算圖（graph），並且由 TensorFlow runtime 負責解釋和執行，所以 Horovod 為了獲得每個程式計算的梯度並且對於它們進行 AllReduce，就得用黑客的辦法進入到 TF 圖執行過程之中去獲取梯度。
-
第二個技術難點是：Horovod 可以自己定義 AllReduce操作, 但是它的AllReduce操作怎麼能夠嵌入到 TF 的處理流程之中？
- 因為 Horovod 自定義的這套HVD Operation 是跟TF OP 無關的，因此是無法直接插入到TF Graph之中進行執行，所以還需要有一個辦法來把HVD OP註冊到TF的OP之中。
0x01 背景概念
我們回憶一下背景概念。
1.1 深度學習框架
深度學習訓練的核心問題是過反向梯度計算來擬合f()，反向梯度計算的目的是計算梯度和更新引數。而計算梯度的方式則主要是通過鏈式求導。一次鏈式求導只是一次的前向和後向的計算結果。模型訓練的重點過程就是：前向傳播和反向傳播。
以簡單的深度神經網路為例，為了完成對損失的優化，我們把資料分成batch，不斷把資料送入模型網路中進行如下迭代過程，目的是使最終優化網路達到收斂：
- 一個batch的資料被送入網路進行前向傳播，前向傳播就是一系列的矩陣+啟用函式等的組合運算。
- 前向傳播輸出的預測值會同真實值 label 進行對比之後，使用損失函式計算出此次迭代的損失；
- 把這個損失進行反向傳播，送入神經網路模型中之前的每一層進行反向梯度計算，更新每一層的權值矩陣和bias；
深度學習框架幫助我們解決的核心問題之一就是反向傳播時的梯度計算和更新。如果不用深度學習框架，就需要我們自己寫方法以進行復雜的梯度計算和更新。
1.2 Tensorflow Optimizer
Tensorflow的底層結構是由張量組成的計算圖。計算圖就是底層的程式設計系統，每一個計算都是圖中的一個節點，計算之間的依賴關係則用節點之間的邊來表示。計算圖構成了前向/反向傳播的結構基礎。
給定一個計算圖, TensorFlow 使用自動微分 (反向傳播) 來進行梯度運算。tf.train.Optimizer允許我們通過minimize()函式自動進行權值更新，此時tf.train.Optimizer.minimize()做了兩件事：
- 計算梯度。即呼叫compute_gradients (loss, var_list ...) 計算loss對指定val_list的梯度，返回元組列表
list(zip(grads, var_list))。 - 用計算得到的梯度來更新對應權重。即呼叫 apply_gradients(grads_and_vars, global_step=global_step, name=None) 將 compute_gradients (loss, var_list ...) 的返回值作為輸入對權重變數進行更新；
將minimize()分成兩個步驟的原因是：可以在某種情況下對梯度進行修正，防止梯度消失或者梯度爆炸。
tensorflow也允許使用者自己計算梯度，在使用者做了中間處理之後，這個梯度會應用給權值進行更新，此時就會細分為以下三個步驟：
- 利用tf.train.Optimizer.compute_gradients計算梯度；
- 使用者對梯度進行自定義處理。這裡其實就是 Horovod 可以做手腳的地方；
- 對於使用者計算後的梯度，利用tf.train.Optimizer.apply_gradients更新權值；
0x02 總體架構
2.1 總體思路
Horovod 作業的每個程式都呼叫單機版 TensorFlow 做本地計算，然後收集梯度，並且通過 AllReduce 來匯聚梯度並且更新每個程式中的模型。
Horovod 需要從 TensorFlow 擷取梯度。
- TensorFlow 1.x
- 在 TensorFlow 1.x 中，深度學習計算是一個計算圖，由 TensorFlow 執行時負責解釋執行。
- Horovod 為了獲得每個程式計算的梯度並且可以對它們進行 AllReduce，就必須潛入圖執行的過程。為此，Horovod 通過對使用者Optimizer 進行封裝組合方式完成了對梯度的 AllReduce 操作，即， Horovod 要求開發者使用Horovod自己定義的 hvd.DistributedOptimizer 代替 TensorFlow 官方的 optimizer，從而可以在優化模型階段得到梯度。
- TensorFlow 2.0
- TensorFlow 2.0 的 eager execution模式 採用完全不同的計算方式。其前向計算過程把對基本計算單元（operator）的呼叫記錄在一個資料結構 tape 裡，隨後進行反向計算過程時候可以回溯這個 tape，以此呼叫 operator 對應的 gradient operator。Tape 提供一個操作讓使用者可以獲取每個引數的梯度。
- Horovod 呼叫 TensorFlow 2.0 API 可以直接獲取梯度。然後Horovod 通過封裝 tape 完成 AllReduce 呼叫。
3.2 總體呼叫關係
我們先給出總體呼叫關係：hvd.DistributedOptimizer繼承keras Optimizer，然後hvd.DistributedOptimizer在其過載的get_gradients中把獲取到的梯度傳給hvd.allreduce(gradients, ...)，從而實現整個horovod叢集的梯度集體歸併。
具體計算梯度的邏輯是：
- TF 呼叫 hvd.DistributedOptimizer 的 compute_gradients 方法：
- hvd.DistributedOptimizer 首先會利用 TF 官方 optimizer.compute_gradients 計算出本地梯度；
- 然後利用 AllReduce 來得到各個程式平均後的梯度；
- compute_gradients 返回一個(梯度，權值)對的列表。由apply_gradients使用；
- TF 呼叫 hvd.DistributedOptimizer 的 apply_gradients 方法：
- 呼叫 TF 官方 optimizer.apply_gradients 對傳入的引數進行處理，返回一個更新權值的op。TF 可以用這個返回值進行後續處理；
因為 TF 的版本問題，所以我們區分 1.x, 2.x 來分析。
0x04 TensorFlow 1.x
前面提到了，Horovod 要求開發者使用Horovod自己定義的 hvd.DistributedOptimizer 代替 TensorFlow 官方的 optimizer，從而可以在優化模型階段得到梯度，所以我們從_DistributedOptimizer進行分析。
4.1 _DistributedOptimizer
以 horovod/tensorflow/__init__.py 為例。
try:
# TensorFlow 2.x
_LegacyOptimizer = tf.compat.v1.train.Optimizer
except AttributeError:
try:
# TensorFlow 1.x
_LegacyOptimizer = tf.train.Optimizer
except AttributeError:
# Future TensorFlow versions
_LegacyOptimizer = None
可以看到，對於 TensorFlow 1.x，我們後續使用的基礎是 _LegacyOptimizer。
_DistributedOptimizer 就繼承了 _LegacyOptimizer。其封裝了另外一個tf.optimizer，在模型應用梯度之前使用allreduce操作收集梯度值並求其均值。這個被封裝的tf.optimizer就是使用者在使用時候指定的TF官方優化器。
具體可以回憶使用者如何使用：
# TF官方Optimizer
opt = tf.optimizers.Adam(scaled_lr)
# 把常規TensorFlow Optimizer通過Horovod包裝起來，進而使用 ring-allreduce 來得到平均梯度
opt = hvd.DistributedOptimizer(
opt, backward_passes_per_step=1, average_aggregated_gradients=True)
# 最後模型使用的是hvd.DistributedOptimizer
mnist_model.compile(loss=tf.losses.SparseCategoricalCrossentropy(),
optimizer=opt, metrics=['accuracy'],
experimental_run_tf_function=False)
opt 被傳給DistributedOptimizer的optimizer，在建構函式__init__.py中被賦值給了self._optimizer。
if _LegacyOptimizer is not None:
class _DistributedOptimizer(_LegacyOptimizer):
"""An optimizer that wraps another tf.Optimizer, using an allreduce to
combine gradient values before applying gradients to model weights."""
def __init__(self, optimizer, name=None, use_locking=False, device_dense='',
device_sparse='', compression=Compression.none,
sparse_as_dense=False, op=Average, gradient_predivide_factor=1.0,
backward_passes_per_step=1, average_aggregated_gradients=False,
groups=None):
self._optimizer = optimizer # 在建構函式中被賦值給了self._optimizer
self._allreduce_grads = _make_allreduce_grads_fn( # 設定歸併函式
name, device_dense, device_sparse, compression, sparse_as_dense, op,
gradient_predivide_factor, groups)
self._agg_helper = None
if backward_passes_per_step > 1:
# 可以先做本地梯度累積，再誇程式合併
self._agg_helper = LocalGradientAggregationHelper(
backward_passes_per_step=backward_passes_per_step,
allreduce_func=self._allreduce_grads,
sparse_as_dense=sparse_as_dense,
average_aggregated_gradients=average_aggregated_gradients,
rank=rank(),
optimizer_type=LocalGradientAggregationHelper._OPTIMIZER_TYPE_LEGACY,
)
4.2 compute_gradients
計算梯度的第一步是 呼叫 compute_gradients 計算loss對指定val_list的梯度，返回元組列表 list(zip(grads, var_list))。
每一個worker的 tensor 模型都會呼叫 compute_gradients，對於每個model來說，
gradients = self._optimizer.compute_gradients(*args, **kwargs) 就是本 model 本地計算得到的梯度。
DistributedOptimizer 重寫Optimizer類compute_gradients()方法。
_DistributedOptimizer初始化時候有配置self._allreduce_grads = _make_allreduce_grads_fn。這裡很重要。- compute_gradients()方法首先呼叫原始配置TF官方 optimizer 的 compute_gradients()。compute_gradients()返回值是一個元祖列表，列表的每個元素是 (gradient，variable)，gradient是每一個變數變化的梯度值；
- 如果設定了 _agg_helper，即 LocalGradientAggregationHelper，就呼叫 LocalGradientAggregationHelper 來做本地梯度累積（本地累積之後也會進行跨程式合併），否則呼叫 _allreduce_grads 計算，即直接跨程式合併（用MPI對計算出來的分散式梯度做allreduce）；
def compute_gradients(self, *args, **kwargs):
"""Compute gradients of all trainable variables.
See Optimizer.compute_gradients() for more info.
In DistributedOptimizer, compute_gradients() is overriden to also
allreduce the gradients before returning them.
"""
# _optimizer是原始配置的官方優化器，先呼叫其compute_gradients方法來計算所有訓練引數的梯度
# 官方優化器的compute_gradients()方法返回一個元組(gradient，variable)的列表
# gradients 被賦值為這個元組(gradient，variable)列表
gradients = self._optimizer.compute_gradients(*args, **kwargs)
grads, vars = zip(*gradients)
if self._agg_helper: # 是否本地先累積
avg_grads = self._agg_helper.compute_gradients(grads, vars)
else:
avg_grads = self._allreduce_grads(grads, vars)
return list(zip(avg_grads, vars))
邏輯如下：
+-----------------------------+
|_DistributedOptimizer |
| |
| | +---------------+
| self._optimizer +----------------> | tf.Optimizer |
| | | |
| | +---------------+
| |
| | +-------------------------+
| _allreduce_grads +----------------> |_make_allreduce_grads_fn |
| | +-------------------------+
| |
| |
| |
| |
| | +-------------------------------------------------+
| compute_gradients +-------------> |compute_gradients |
| | | |
+-----------------------------+ | |
| _optimizer.compute_gradients |
| + |
| | |
| | |
| v |
| _agg_helper.compute_gradients(grads, vars) |
| |
| _allreduce_grads(grads, vars) |
| + |
| | |
| | |
| v |
| list(zip(avg_grads, vars)) |
| |
+-------------------------------------------------+
4.3 LocalGradientAggregationHelper
前面提到，如果設定了 _agg_helper，即 LocalGradientAggregationHelper，就呼叫 LocalGradientAggregationHelper 來做本地累積梯度（本地累積之後也會進行跨程式合併）。所以我們講講 LocalGradientAggregationHelper。
LocalGradientAggregationHelper 會在本地更新梯度，但是因為在初始化時候，成員函式 self._allreduce_grads = allreduce_func 就是跨程式allreduce函式。所以 LocalGradientAggregationHelper 之中也會進行跨程式 allreduce。即每次 backward_passes_per_step 時候跨機器更新一次。
這裡需要注意的是：allreduce_func=self._allreduce_grads，其實 LocalGradientAggregationHelper 內部呼叫 self._allreduce_grads也是呼叫到了 _make_allreduce_grads_fn。
LocalGradientAggregationHelper(
backward_passes_per_step=backward_passes_per_step,
allreduce_func=self._allreduce_grads, # 就是_make_allreduce_grads_fn
sparse_as_dense=sparse_as_dense,
average_aggregated_gradients=average_aggregated_gradients,
rank=rank(),
optimizer_type=LocalGradientAggregationHelper._OPTIMIZER_TYPE_KERAS,
)
具體是呼叫了 LocalGradientAggregationHelper.compute_gradients 完成功能，其中：
- _init_aggregation_vars 函式會 遍歷 本地元組（gradient，variable）的列表，累積在 locally_aggregated_grads。
- allreduce_grads 會做一個遍歷 tensor & 應用 tensor 的操作，對於每個 tensor，_allreduce_grads_helper 函式會進行跨程式合併。
4.3.1 _init_aggregation_vars
_init_aggregation_vars 函式會 遍歷 本地元組（gradient，variable）的列表，累積在 locally_aggregated_grads。
def _init_aggregation_vars(self, grads):
"""
Initializes the counter that is used when to communicate and aggregate gradients
and the tensorflow variables that store the locally aggregated gradients.
"""
variable_scope_name = "aggregation_variables_" + str(self.rank)
with tf.compat.v1.variable_scope(variable_scope_name, reuse=tf.compat.v1.AUTO_REUSE):
self.counter = tf.compat.v1.get_variable(
"aggregation_counter", shape=(), dtype=tf.int32,
trainable=False, initializer=tf.compat.v1.zeros_initializer(),
collections=[tf.compat.v1.GraphKeys.LOCAL_VARIABLES],
)
# 遍歷本地的梯度
for idx, grad in enumerate(grads):
# Handle IndexedSlices.
# 如果是IndexedSlices，則轉換為張量
if self.sparse_as_dense and isinstance(grad, tf.IndexedSlices):
grad = tf.convert_to_tensor(grad)
elif isinstance(grad, tf.IndexedSlices):
raise ValueError(
"IndexedSlices are not supported when "
"`backward_passes_per_step` > 1 and "
"`sparse_as_dense` is False."
)
# Handle grads that are None.
# 如果為空，則跳過
if grad is None:
self.num_none_grad_updates += 1
continue
self.not_none_indexes[idx] = len(self.locally_aggregated_grads)
# Create shadow variable.
grad_aggregation_variable_name = str(idx)
zero_grad = tf.zeros(shape=grad.get_shape().as_list(), dtype=grad.dtype)
grad_aggregation_variable = tf.compat.v1.get_variable(
grad_aggregation_variable_name,
trainable=False,
initializer=zero_grad,
collections=[
tf.compat.v1.GraphKeys.LOCAL_VARIABLES,
"aggregating_collection"],
)
# 新增到本地累積變數 locally_aggregated_grads 之中
self.locally_aggregated_grads.append(grad_aggregation_variable)
assert len(self.locally_aggregated_grads) + \
self.num_none_grad_updates == len(grads)
# We expect to get a `sess` when we need to manually do a `sess.run(...)`
# for the variables to be initialized. This is the `tf.keras`
# optimizers.
# 遍歷locally_aggregated_grads的變數，如果需要則進行初始化
if self.optimizer_type == self._OPTIMIZER_TYPE_KERAS:
session = tf.compat.v1.keras.backend.get_session(op_input_list=())
vars_init_op = tf.compat.v1.variables_initializer(
[self.counter, *get_not_none_from_list(self.locally_aggregated_grads)]
)
session.run(vars_init_op)
4.3.2 compute_gradients
compute_gradients方法具體如下：
def compute_gradients(self, grads, vars):
"""
Applies the new gradient updates the locally aggregated gradients, and
performs cross-machine communication every backward_passes_per_step
times it is called.
"""
# 遍歷 本地元組（gradient，variable）的列表，累積在 locally_aggregated_grads
self._init_aggregation_vars(grads)
# Clear the locally aggregated gradients when the counter is at zero.
# 如果計數器為0，則清理本地累積梯度
clear_op = tf.cond(
pred=tf.equal(self.counter, 0),
true_fn=lambda: self._clear_grads(),
false_fn=tf.no_op
)
# Add new gradients to the locally aggregated gradients.
# 本地累積梯度
with tf.control_dependencies([clear_op]):
aggregation_ops_list = self._aggregate_grads(grads)
# Increment the counter once new gradients have been applied.
# 一旦本地梯度已經被應用，則把計數器加1
aggregation_ops = tf.group(*aggregation_ops_list)
with tf.control_dependencies([aggregation_ops]):
update_counter = self.counter.assign_add(tf.constant(1))
# 應用梯度
with tf.control_dependencies([update_counter]):
grads = get_not_none_from_list(grads)
assert len(grads) == len(self.locally_aggregated_grads)
# Allreduce locally aggregated gradients when the counter is equivalent to
# `backward_passes_per_step`. This the condition is true, it also resets
# the counter back to 0.
allreduced_grads = tf.cond(
tf.equal(self.counter, self.backward_passes_per_step),
lambda: self._allreduce_grads_helper(grads, vars),
lambda: grads,
)
# Handle case where there is only one variable.
if not isinstance(allreduced_grads, (list, tuple)):
allreduced_grads = (allreduced_grads,)
# Insert gradients that are None back in.
# 對於本地累積的梯度，進行跨程式合併，locally_aggregated_grads是本地累積的梯度
allreduced_grads = [
allreduced_grads[self.not_none_indexes[idx]] if idx in self.not_none_indexes else None
for idx in range(len(self.locally_aggregated_grads) + self.num_none_grad_updates)
]
# If gradients have not been allreduced this batch, we return the gradients
# that were submitted as the updates (the input).
return allreduced_grads # 返回跨程式合併之後的梯度
邏輯擴充如下，這裡需要注意的是 _agg_helper 或者 _allreduce_grads 選一個執行：
- 如果設定了 _agg_helper，即 LocalGradientAggregationHelper，就呼叫 _agg_helper 來計算梯度（本地累積之後也會進行跨程式合併）；
- 否則呼叫 _allreduce_grads，即 _make_allreduce_grads_fn 計算，即跨程式合併（用MPI來對計算出來的分散式梯度做allreduce操作）；
+-----------------------------+
|_DistributedOptimizer | +-----------------------------------------------------+
| | | LocalGradientAggregationHelper |
| | +---------------+ | |
| self._optimizer +----------------> | tf.Optimizer | | +---------------------------------------------+ |
| | | | | | compute_gradients | |
| | +---------------+ | | | |
| | | | | |
| | +------------------------------------------------------+ | | _init_aggregation_vars | |
| compute_gradients +-------------> |compute_gradients | | | + | |
| | | | | | | | |
| | | | | | | | |
| | | _optimizer.compute_gradients | | | v | |
| _allreduce_grads | | + | | | | |
| + | | | | | | _allreduce_grads_helper | |
| | | | | | | | + | |
+-----------------------------+ | v | | | | | |
| | _agg_helper.compute_gradients(grads, vars) +------------> | | | |
| | | | | v | |
| +--------------+ _allreduce_grads(grads, vars) | | | allreduced_grads | |
| | | + | | | | |
| | | | | | +---------------------------------------------+ |
| | | | | | |
| | | v | | allreduce_func |
| | | list(zip(avg_grads, vars)) | | + |
| | | | | | |
| | +------------------------------------------------------+ +-----------------------------------------------------+
| | |
| | |
v v |
+-------+-------------------+--------+ |
|_make_allreduce_grads_fn | |
| | <-----------------------------------------------------------------------+
| _allreduce_cond |
| |
| |
| |
+------------------------------------+
具體如下：
4.4 _make_allreduce_grads_fn
_make_allreduce_grads_fn 就是呼叫了 _make_cached_allreduce_grads_fn 完成功能。
def _make_allreduce_grads_fn(name, device_dense, device_sparse,
compression, sparse_as_dense, op,
gradient_predivide_factor, groups):
groups = vars_to_refs(groups) if isinstance(groups, list) else groups
return _make_cached_allreduce_grads_fn(name, device_dense, device_sparse,
compression, sparse_as_dense, op,
gradient_predivide_factor, groups)
_make_cached_allreduce_grads_fn 的作用是：
- 獲取所有grads；
- 遍歷元組(gradient，variable)的列表，對於每個grad，使用_allreduce_cond與其他worker進行同步；
- 最後返回同步好的梯度列表；
@_cache
def _make_cached_allreduce_grads_fn(name, device_dense, device_sparse,
compression, sparse_as_dense, op,
gradient_predivide_factor, groups):
groups = refs_to_vars(groups) if isinstance(groups, tuple) else groups
......
def allreduce_grads(grads, vars=None):
with tf.name_scope(name + "_Allreduce"): # 設定名稱空間
......
# 獲取所有的 grads
# 因為grads列表致為((grad0,var0),(grad1,var1)…)，裡面可能有很多None，所以提取出grad不為None的var進行梯度計算。
return [_allreduce_cond(grad,
device_dense=device_dense,
device_sparse=device_sparse,
compression=compression,
op=op,
prescale_factor=prescale_factor,
postscale_factor=postscale_factor)
if grad is not None else grad
for grad in grads]
if _executing_eagerly():
return _make_subgraph(allreduce_grads)
else:
return allreduce_grads
_allreduce_cond 函式中就是呼叫到 allreduce 進行集合通訊操作。
def _allreduce_cond(tensor, *args, **kwargs):
def allreduce_fn():
return allreduce(tensor, *args, **kwargs)
def id_fn():
return tensor
return tf.cond((size_op() > 1) if int(os.environ.get("HOROVOD_ELASTIC", 0)) else tf.convert_to_tensor(size() > 1),
allreduce_fn, id_fn)
4.5 allreduce
allreduce()方法之中，會依據所需要傳輸的張量型別是Tensor還是 IndexedSlices 做不同處理。
- 如果 tensor型別是IndexedSlices，則只需要做allgather操作，是否需要其他操作需要看具體附加配置。
- 因為對於分佈在不同worker上的IndexedSlices，其values和indices彼此沒有重複。
- 假設在 worker 1上分佈的indices是[1, 3, 5, 7, 9]，在worker 2上分佈的indices是[2, 4, 6, 8, 10]。只需要使用allgather方法將其收集彙總得到 [1,2,3,4,5,6,7,8,9,10] 即可，不需要做求和/平均的操作。
- 如果有附加操作，才需要進一步處理。
- 如果是 Tensor 型別，則需要呼叫_allreduce方法處理：先求張量的和，再取平均。
def allreduce(tensor, average=None, device_dense='', device_sparse='',
compression=Compression.none, op=None,
prescale_factor=1.0, postscale_factor=1.0,
name=None):
"""Perform an allreduce on a tf.Tensor or tf.IndexedSlices.
"""
op = handle_average_backwards_compatibility(op, average)
if isinstance(tensor, tf.IndexedSlices): # 對於IndexedSlices型別
# TODO: Need to fix this to actuall call Adasum
if op == Adasum:
with tf.device(device_sparse):
# For IndexedSlices, do two allgathers instead of an allreduce.
# 做兩個allgathers操作即可
horovod_size = tf.cast(size_op() if int(os.environ.get("HOROVOD_ELASTIC", 0)) else size(),
dtype=tensor.values.dtype)
values = allgather(tensor.values) # 一個 allgeathers對value進行處理
indices = allgather(tensor.indices) # 一個allgather對index進行處理
# To make this operation into an average, divide allgathered values by
# the Horovod size.
# 如果op是Average，則需要計算所有value的均值，否則不做操作
new_values = (values / horovod_size) if op == Average else values
return tf.IndexedSlices(new_values, indices,
dense_shape=tensor.dense_shape)
else: # 對於Tensor型別
average_in_framework = False
if rocm_built():
# For ROCm, perform averaging at framework level
average_in_framework = op == Average or op == Adasum
op = Sum if op == Average else op
with tf.device(device_dense):
# 首先，將size_op()結果的型別轉化為tensor的dtype型別
horovod_size = tf.cast(size_op() if int(os.environ.get("HOROVOD_ELASTIC", 0)) else size(),
dtype=tensor.dtype)
tensor_compressed, ctx = compression.compress(tensor)
# 定義了一個sum/壓縮操作: 將某張量和其他所有Horovod程式同名張量求和
summed_tensor_compressed = _allreduce(tensor_compressed, op=op,
prescale_factor=prescale_factor,
postscale_factor=postscale_factor,
name=name)
summed_tensor = compression.decompress(summed_tensor_compressed, ctx)
if op == Adasum: # 處理其他附加操作
if 'CPU' not in tensor.device and gpu_available('tensorflow'):
if nccl_built():
if not is_homogeneous:
elif not check_num_rank_power_of_2(int(size() / local_size())):
if rocm_built():
horovod_local_size = tf.cast(local_size_op() if int(os.environ.get("HOROVOD_ELASTIC", 0)) else local_size(),
dtype=tensor.dtype)
new_tensor = summed_tensor / horovod_local_size
else:
new_tensor = summed_tensor
else:
new_tensor = summed_tensor
else:
new_tensor = summed_tensor
else:
if rocm_built():
new_tensor = (summed_tensor / horovod_size) if average_in_framework else summed_tensor
else:
new_tensor = summed_tensor
return new_tensor
4.6 _allreduce
_allreduce方法 和 allgather方法在 horovod.tensorflow.mpi_ops.py 之中。
HorovodAllreduceOp和HorovodAllgatherOp這兩個方法是HVD自定義的與tensorflow相關的OP。_allreduce 和 allgather 分別與之對應。
- _allreduce使用名字“HorovodAllreduce”和HorovodAllreduceOp繫結，由 MPI_LIB.horovod_allreduce 做了中間轉換；
- allgather使用名字“HorovodAllgather”和HorovodAllgatherOp繫結，由 MPI_LIB.horovod_allgather 做了中間轉換；
結合前面的 _make_cached_allreduce_grads_fn 之中對於名字空間的配置，張量名稱大致為：DistributedAdam_Allreduce/cond_14/HorovodAllreduce_grads_5_0。
這樣就呼叫到了 MPI 對應操作。
def _allreduce(tensor, name=None, op=Sum, prescale_factor=1.0, postscale_factor=1.0,
ignore_name_scope=False):
"""An op which reduces an input tensor over all the Horovod processes. The
default reduction is a sum.
The reduction operation is keyed by the name of the op. The tensor type and
shape must be the same on all Horovod processes for a given name. The reduction
will not start until all processes are ready to send and receive the tensor.
Returns:
A tensor of the same shape and type as `tensor`, summed across all
processes.
"""
if name is None and not _executing_eagerly():
name = 'HorovodAllreduce_%s' % _normalize_name(tensor.name)
return MPI_LIB.horovod_allreduce(tensor, name=name, reduce_op=op,
prescale_factor=prescale_factor,
postscale_factor=postscale_factor,
ignore_name_scope=ignore_name_scope)
def allgather(tensor, name=None, ignore_name_scope=False):
"""An op which concatenates the input tensor with the same input tensor on
all other Horovod processes.
The concatenation is done on the first dimension, so the input tensors on the
different processes must have the same rank and shape, except for the first
dimension, which is allowed to be different.
Returns:
A tensor of the same type as `tensor`, concatenated on dimension zero
across all processes. The shape is identical to the input shape, except for
the first dimension, which may be greater and is the sum of all first
dimensions of the tensors in different Horovod processes.
"""
if name is None and not _executing_eagerly():
name = 'HorovodAllgather_%s' % _normalize_name(tensor.name)
return MPI_LIB.horovod_allgather(tensor, name=name,
ignore_name_scope=ignore_name_scope)
4.7 操作對映
Python世界中，呼叫 _allreduce 時傳遞了幾個引數，比如tensor和name。其中 op=Sum 最為重要。這個是被 C++ 內部用來確定 reduction具體操作。我們具體梳理下：
4.7.1 C++定義
在 C++中有：
enum ReduceOp {
AVERAGE = 0, // This value should never appear past framework code, as
// averaging is taken care of there.
SUM = 1,
ADASUM = 2
};
int horovod_reduce_op_sum() {
return ReduceOp::SUM;
}
4.7.2 Python獲取配置
在 python 的初始化程式碼中有：
class HorovodBasics(object):
"""Wrapper class for the basic Horovod API."""
def __init__(self, pkg_path, *args):
full_path = util.get_extension_full_path(pkg_path, *args)
self.MPI_LIB_CTYPES = ctypes.CDLL(full_path, mode=ctypes.RTLD_GLOBAL)
self.Average = self.MPI_LIB_CTYPES.horovod_reduce_op_average()
self.Sum = self.MPI_LIB_CTYPES.horovod_reduce_op_sum() # 在這裡聯絡起來
self.Adasum = self.MPI_LIB_CTYPES.horovod_reduce_op_adasum()
這樣，在呼叫 _allreduce 預設引數是 op=Sum，就對應了 C++ 的 ReduceOp::SUM。
4.7.3 建立聯絡
_allreduce 繼續呼叫：
MPI_LIB.horovod_allreduce(tensor, name=name, reduce_op=op
MPI_LIB.horovod_allreduce被轉換到了C++世界下面程式碼中
- 首先，通過OP_REQUIRES_OK的配置可以得到reduce_op_；
- 其次，ComputeAsync 之中通過 reduce_op_ 就可以確定具體需要呼叫那種操作；
因此，Python和C++世界就進一步聯絡起來。
class HorovodAllreduceOp : public AsyncOpKernel {
public:
explicit HorovodAllreduceOp(OpKernelConstruction* context)
: AsyncOpKernel(context) {
// 這裡會宣告，從 context 中得到reduce_op，賦值給reduce_op_
OP_REQUIRES_OK(context, context->GetAttr("reduce_op", &reduce_op_));
// 省略無關程式碼
}
void ComputeAsync(OpKernelContext* context, DoneCallback done) override {
OP_REQUIRES_OK_ASYNC(context, ConvertStatus(common::CheckInitialized()),
done);
// 省略無關程式碼
// 這裡會依據 reduce_op_，來確認C++內部呼叫何種操作
horovod::common::ReduceOp reduce_op = static_cast<horovod::common::ReduceOp>(reduce_op_);
// 省略無關程式碼
}
4.8 擴充流程
我們擴充目前流程圖如下：
+-----------------------------+
|_DistributedOptimizer | +-----------------------------------------------------+
| | | LocalGradientAggregationHelper |
| | +---------------+ | |
| self._optimizer +----------------> | tf.Optimizer | | +---------------------------------------------+ |
| | | | | | compute_gradients | |
| | +---------------+ | | | |
| | | | | |
| | +------------------------------------------------------+ | | _init_aggregation_vars | |
| compute_gradients +-------------> |compute_gradients | | | + | |
| | | | | | | | |
| | | | | | | | |
| | | _optimizer.compute_gradients | | | v | |
| _allreduce_grads | | + | | | | |
| + | | | | | | _allreduce_grads_helper | |
| | | | | | | | + | |
+-----------------------------+ | v | | | | | |
| | _agg_helper.compute_gradients(grads, vars) +------------> | | | |
| | | | | v | |
| +--------------+ _allreduce_grads(grads, vars) | | | allreduced_grads | |
| | | + | | | | |
| | | | | | +---------------------------------------------+ |
| | | | | | |
| | | v | | allreduce_func |
| | | list(zip(avg_grads, vars)) | | + |
| | | | | | |
| | +------------------------------------------------------+ +-----------------------------------------------------+
| | |
| | |
v v |
+-------+-------------------+--------+ |
|_make_allreduce_grads_fn | |
| | <-----------------------------------------------------------------------+
| |
| | +-----------------+ +----------------+ +----------------------------+
| _allreduce_cond +-------------------> | allreduce | | _allreduce | | MPI_LIB.horovod_allreduce |
| | | +----------------> | +---------------> | |
+------------------------------------+ | | | | | |
| | | | | |
+-----------------+ +----------------+ +----------------------------+
手機如下：
0x05 Tensorflow 2.x
5.1 Horovod 實施
對於 TF2.x，每行程式碼順序執行，不需要構建圖，也取消了control_dependency。Horovod 通過呼叫 TensorFlow 2.0 API 可以很直接地獲取梯度。所以 Horovod 梯度更新部分的實現並不是基於計算圖的實現，而是使用 hvd.DistributedGradientTape。
Worker 在訓練時候做如下操作：
- 使用 DistributedGradientTape 封裝 TF 官方的 Tape，配置 allreduce函式。
- 讀取一組訓練資料。
- 在本地模型呼叫前向傳播函式計算損失。
- 給定損失之後，worker 利用 TensorFlow eager execution 的 GradientTape 機制，呼叫基類函式得到梯度。
- 各個Worker 會呼叫 Allreduce 來同步梯度。
- 各個Worker 會依據最新梯度相應更新模型。
5.2 示例程式碼
首先，我們給出示例程式碼如下，下面省略部分非關鍵程式碼，具體可以參見注釋：
# Horovod: initialize Horovod.
hvd.init() # 初始化HVD
# Horovod: pin GPU to be used to process local rank (one GPU per process)
# 配置GPU
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
if gpus:
tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')
# 載入資料
(mnist_images, mnist_labels), _ = \
tf.keras.datasets.mnist.load_data(path='mnist-%d.npz' % hvd.rank())
# 把資料進行特徵切片
dataset = tf.data.Dataset.from_tensor_slices(
(tf.cast(mnist_images[..., tf.newaxis] / 255.0, tf.float32),
tf.cast(mnist_labels, tf.int64))
)
# 打亂資料，分批載入
dataset = dataset.repeat().shuffle(10000).batch(128)
mnist_model = tf.keras.Sequential([
tf.keras.layers.Conv2D(32, [3, 3], activation='relu'),
tf.keras.layers.Conv2D(64, [3, 3], activation='relu'),
tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
tf.keras.layers.Dropout(0.25),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dropout(0.5),
tf.keras.layers.Dense(10, activation='softmax')
])
# 損失函式
loss = tf.losses.SparseCategoricalCrossentropy()
# Horovod: adjust learning rate based on number of GPUs.
opt = tf.optimizers.Adam(0.001 * hvd.size())
@tf.function
def training_step(images, labels, first_batch):
with tf.GradientTape() as tape:
probs = mnist_model(images, training=True)
loss_value = loss(labels, probs)
# Horovod: add Horovod Distributed GradientTape.
# 呼叫 DistributedGradientTape，配置allreduce函式
tape = hvd.DistributedGradientTape(tape)
# 顯式得到梯度，其內部經過一系列操作後，會呼叫horovod的allreduce操作，最終是MPI_LIB.horovod_allreduce函式
grads = tape.gradient(loss_value, mnist_model.trainable_variables)
# 應用梯度，更新權重
opt.apply_gradients(zip(grads, mnist_model.trainable_variables))
# Horovod: broadcast initial variable states from rank 0 to all other processes.
# This is necessary to ensure consistent initialization of all workers when
# training is started with random weights or restored from a checkpoint.
#
# Note: broadcast should be done after the first gradient step to ensure optimizer
# initialization.
# 廣播變數
if first_batch:
hvd.broadcast_variables(mnist_model.variables, root_rank=0)
hvd.broadcast_variables(opt.variables(), root_rank=0)
return loss_value
# Horovod: adjust number of steps based on number of GPUs.
for batch, (images, labels) in enumerate(dataset.take(10000 // hvd.size())):
loss_value = training_step(images, labels, batch == 0)
5.3 _DistributedGradientTape
關鍵類_DistributedGradientTape 定義如下：
class _DistributedGradientTape(tf.GradientTape):
def __init__(self, tape, device_dense, device_sparse, compression, sparse_as_dense, op,
gradient_predivide_factor, groups, persistent=False,
watch_accessed_variables=True):
if hasattr(tape, '_watch_accessed_variables'):
super(self.__class__, self).__init__(persistent, watch_accessed_variables)
else:
super(self.__class__, self).__init__(persistent)
# 把TF官方tape儲存起來
self._tape = tape
# 配置allreduce函式
self._allreduce_grads = _make_allreduce_grads_fn(
'DistributedGradientTape', device_dense, device_sparse, compression,
sparse_as_dense, op, gradient_predivide_factor, groups)
# 使用者顯式的呼叫此函式，其內部使用_make_allreduce_grads_fn進行處理
def gradient(self, target, sources, output_gradients=None):
# 呼叫基類函式獲得梯度
gradients = super(self.__class__, self).gradient(target, sources, output_gradients)
return self._allreduce_grads(gradients, sources)
_make_allreduce_grads_fn 函式會進行一系列呼叫，最終呼叫到 MPI_LIB.horovod_allreduce，具體做如下工作：
- 修改name scope，加上字尾
_Allreduce； - 如果配置，則進行壓縮；
- 依據op型別，呼叫allreduce 或者 直接返回tensor；
- DistributedGradientTape 的 name scope 被改寫成了 DistributedGradientTape_Allreduce，名字被加上了 HorovodAllreduce_ 的字首。
- 呼叫MPI_LIB.horovod_allreduce函式；
@_cache
def _make_allreduce_grads_fn(name, device_dense, device_sparse,
compression, sparse_as_dense, op):
def allreduce_grads(grads):
with tf.name_scope(name + "_Allreduce"): # 修改name scope，加上字尾
if sparse_as_dense:
grads = [tf.convert_to_tensor(grad) # 壓縮
if grad is not None and isinstance(grad, tf.IndexedSlices)
else grad for grad in grads]
return [_allreduce_cond(grad,
device_dense=device_dense,
device_sparse=device_sparse,
compression=compression,
op=op)
if grad is not None else grad
for grad in grads]
def _allreduce_cond(tensor, *args, **kwargs):
def allreduce_fn():
return allreduce(tensor, *args, **kwargs)
def id_fn():
return tensor
return tf.cond(size_op() > 1, allreduce_fn, id_fn) # 不用的呼叫方法
def _allreduce(tensor, name=None, op=Sum):
"""An op which reduces an input tensor over all the Horovod processes. The
default reduction is a sum.
The reduction operation is keyed by the name of the op. The tensor type and
shape must be the same on all Horovod processes for a given name. The reduction
will not start until all processes are ready to send and receive the tensor.
Returns:
A tensor of the same shape and type as `tensor`, summed across all
processes.
"""
if name is None and not _executing_eagerly():
name = 'HorovodAllreduce_%s' % _normalize_name(tensor.name)
# # 呼叫HorovodAllreduceOp
return MPI_LIB.horovod_allreduce(tensor, name=name, reduce_op=op)
邏輯如下：
+-------------------------------+
| _DistributedGradientTape | +------------------------------------+
| | |_make_allreduce_grads_fn |
| | | |
| _allreduce_grads +---------------> | |
| | | |
| | | _allreduce_cond +---------+
| | | | |
+-------------------------------+ +------------------------------------+ |
|
|
+--------------------------------------------------------------------------+
|
|
|
|
| +----------------+ +----------------------------+
| | _allreduce | | MPI_LIB.horovod_allreduce |
+-------> | +---------------> | |
| | | |
| | | |
+----------------+ +----------------------------+
0x06 HorovodAllreduceOp
MPI_LIB.horovod_allreduce 呼叫的就是 HorovodAllreduceOp。MPI_LIB.horovod_allreduce 是 python 函式，HorovodAllreduceOp 是C++程式碼，這裡 TF 做了一個適配和轉換，讓我們可以從 python 函式直接呼叫到 C++ 函式。
HorovodAllreduceOp 繼承了AsyncOpKernel，是一種TF Async OP，而且被 REGISTER_KERNEL_BUILDER 註冊到 TF，因此就可以嵌入到 TF 流程之中。
TF 會呼叫到 HorovodAllreduceOp 所覆蓋的ComputeAsync方法，在ComputeAsync內部會把 張量的Allreduce操作加入Horovod後臺佇列，從而把 TF OP 和 Horovod OP 聯絡起來。
總結一下，HorovodAllreduceOp 繼承了TF AsyncOpKernel，因此可以嵌入到 TF 流程，同時用組合方式與 Horovod 後臺執行緒聯絡起來。
class HorovodAllreduceOp : public AsyncOpKernel { //派生了，所以可以嵌入到 TF流程之中
public:
explicit HorovodAllreduceOp(OpKernelConstruction* context)
: AsyncOpKernel(context) {
OP_REQUIRES_OK(context, context->GetAttr("reduce_op", &reduce_op_));
OP_REQUIRES_OK(context, context->GetAttr("prescale_factor", &prescale_factor_));
OP_REQUIRES_OK(context, context->GetAttr("postscale_factor", &postscale_factor_));
OP_REQUIRES_OK(context, context->GetAttr("ignore_name_scope", &ignore_name_scope_));
}
void ComputeAsync(OpKernelContext* context, DoneCallback done) override {
OP_REQUIRES_OK_ASYNC(context, ConvertStatus(common::CheckInitialized()),
done);
... // 省略一些變數驗證，初始化程式碼
// 將張量的Allreduce操作OP加入佇列
auto enqueue_result = EnqueueTensorAllreduce(
hvd_context, hvd_tensor, hvd_output, ready_event, node_name, device,
[context, done](const common::Status& status) {
context->SetStatus(ConvertStatus(status));
done();
}, reduce_op, (double) prescale_factor_, (double) postscale_factor_);
OP_REQUIRES_OK_ASYNC(context, ConvertStatus(enqueue_result), done);
}
private:
int reduce_op_;
// Using float since TF does not support double OP attributes
float prescale_factor_;
float postscale_factor_;
bool ignore_name_scope_;
};
從下文開始我們看看Horovod on Spark。
0xEE 個人資訊
★★★★★★關於生活和技術的思考★★★★★★
微信公眾賬號：羅西的思考
如果您想及時得到個人撰寫文章的訊息推送，或者想看看個人推薦的技術資料，可以掃描下面二維碼（或者長按識別二維碼）關注個人公眾號）。
0xFF 參考
tf.train.SyncReplicasOptimizer no synchronization among workers #11753
Synchronous distributed tensorflow training doesn’t synchronize among workers #9596
tf.train.SyncReplicasOptimizer
Slow and Stale Gradients Can Win the Race: Error-Runtime Trade-offs in Distributed SGD
當Spark遇上TensorFlow分散式深度學習框架原理和實踐
TensorFlow SyncReplicasOptimizer 解讀
https://blog.csdn.net/edward_zcl/article/details/90345318
TensorFlow學習（四）：梯度帶(GradientTape)，優化器(Optimizer)和損失函式(losses)
ElasticDL 深度學習框架簡化程式設計，提升叢集利用率和研發效率的祕訣
tensorflow分散式原始碼解讀4：AdamOptimizer