0x00 сводка
Horovod — это простая в использовании высокопроизводительная распределенная платформа обучения, выпущенная Uber в 2017 году и широко используемая в отрасли.
Эта серия поможет вам понять Horovod с помощью анализа исходного кода. В этой статье, седьмой в серии, рассматривается, как Horovod интегрируется с TensorFlow.
Предыдущие ссылки выглядят следующим образом:
Нам нужны некоторые вопросы или вопросы дизайна, чтобы направить анализ, и, поскольку читатели, возможно, не видели других статей в этой серии, пункты вопроса будут частично повторяться в других статьях:
-
Первая техническая трудность: как Horovod получает градиенты из потока выполнения TF для обработки?
- В TensorFlow 1.x процесс расчета глубокого обучения представлен в виде вычислительного графа (графа), который интерпретируется и исполняется средой выполнения TensorFlow, поэтому Хороводу приходится использовать хаки, чтобы получить рассчитанные каждым процессом градиенты и выполнить AllReduce. на них Метод входит в процесс выполнения графа TF для получения градиента.
-
Вторая техническая трудность: Horovod может сам определять операции AllReduce, но как его операции AllReduce могут быть встроены в поток обработки TF?
- Поскольку операция HVD, настроенная Horovod, не имеет ничего общего с TF OP, ее нельзя напрямую вставить в TF Graph для выполнения, поэтому должен быть способ зарегистрировать HVD OP в TF OP.
0x01 Фоновая концепция
Вспомним понятие фона.
1.1 Фреймворк глубокого обучения
Основная проблема обучения глубокому обучению состоит в том, чтобы подобрать f () с помощью вычисления обратного градиента, Цель вычисления обратного градиента — вычислить градиенты и обновить параметры. Способ расчета градиента в основном заключается в построении цепочки. Вывод цепочки — это всего лишь одно прямое и обратное вычисление. Ключевым процессом обучения модели является: прямое распространение и обратное распространение.
Взяв в качестве примера простую глубокую нейронную сеть, чтобы завершить оптимизацию потерь, мы разделяем данные на пакеты и непрерывно отправляем данные в модельную сеть для следующего итеративного процесса, чтобы сделать окончательную сеть оптимизации сходятся:
- Пакет данных отправляется в сеть для прямого распространения, которое представляет собой комбинацию ряда матриц + функций активации.
- После того, как прогнозируемое значение выхода прямого распространения сравнивается с меткой реального значения, функция потерь используется для расчета потерь этой итерации;
- Эта потеря распространяется обратно и отправляется на каждый предыдущий уровень в модели нейронной сети для вычисления обратного градиента, а весовая матрица и смещение каждого слоя обновляются;
Одной из основных проблем, которую помогают нам решить фреймворки глубокого обучения, является вычисление и обновление градиента во время обратного распространения. Без фреймворка глубокого обучения нам нужно писать собственные методы для сложных вычислений и обновлений градиента.
1.2 Tensorflow Optimizer
Базовая структура Tensorflow представляет собой вычислительный граф, состоящий из тензоров. Граф вычислений — это базовая система программирования, каждое вычисление — это узел в графе, а зависимости между вычислениями представлены ребрами между узлами. Вычислительные графы составляют структурную основу для прямого/обратного распространения.
Учитывая вычислительный граф, TensorFlow использует автоматическое дифференцирование (обратное распространение) для выполнения градиентных операций. tf.train.Optimizer позволяет нам автоматически обновлять веса с помощью функции минимизации().На данный момент tf.train.Optimizer.minimize() делает две вещи:
- Рассчитать градиент. То есть, вызовите calculate_gradients(loss, var_list...), чтобы вычислить градиент потерь до указанного val_list и вернуть список кортежей.
list(zip(grads, var_list)). - Соответствующие веса обновляются вычисленными градиентами. То есть вызовите apply_gradients(grads_and_vars, global_step=global_step, name=None), чтобы обновить переменную веса возвращаемым значением calculate_gradients (loss, var_list...) в качестве входных данных;
Причина разделения функции минимизации() на два этапа заключается в том, что градиент можно скорректировать при определенных обстоятельствах, чтобы предотвратить исчезновение или взрыв градиента.
Tensorflow также позволяет пользователям вычислять градиент самостоятельно.После того, как пользователь выполнит промежуточную обработку, градиент будет применен к весам для обновления.В это время он будет разделен на следующие три шага:
- Используйте tf.train.Optimizer.compute_gradients для расчета градиентов;
- Пользовательская обработка градиентов. Именно здесь Хоровод может проделывать трюки;
- Для градиента, рассчитанного пользователем, используйте tf.train.Optimizer.apply_gradients для обновления весов;
0x02 Общая архитектура
2.1 Общая идея
Каждый процесс задания Horovod вызывает автономную версию TensorFlow для локальных вычислений, затем собирает градиенты и использует AllReduce для объединения градиентов и обновления модели в каждом процессе.
Хороводу нужно перехватывать градиенты из TensorFlow.
- TensorFlow 1.x
- В TensorFlow 1.x вычисление глубокого обучения представляет собой вычислительный граф, который интерпретируется и выполняется средой выполнения TensorFlow.
- Horovod должен погрузиться в процесс выполнения графа, чтобы получить градиенты, вычисляемые каждым процессом, и иметь возможность их AllReduce. к этому концу,Horovod завершает операцию AllReduce с градиентами, инкапсулируя и комбинируя пользовательский оптимизатор.То есть Horovod требует от разработчиков использовать собственный определенный Horovod hvd.DistributedOptimizer вместо официального оптимизатора TensorFlow, чтобы градиент можно было получить на этапе оптимизации модели.
- TensorFlow 2.0
- Активный режим выполнения TensorFlow 2.0 использует совершенно другой подход к вычислениям. Процесс прямого вычисления записывает вызов базовой вычислительной единицы (оператора) в ленту структуры данных, а затем ленту можно проследить назад, когда выполняется процесс обратного вычисления, чтобы вызвать оператор градиента, соответствующий оператору. Лента обеспечивает операцию, которая позволяет пользователю получить градиент каждого параметра.
- Horovod может напрямую получать градиенты, вызывая API TensorFlow 2.0. потомХоровод завершает вызовы AllReduce, обматывая их лентой.
3.2 Общие отношения вызова
Сначала мы задаем общее отношение вызова: hvd.DistributedOptimizer наследует keras Optimizer, а затем hvd.DistributedOptimizer передает полученный градиент в hvd.allreduce(gradients, ...) в своем перегруженном get_gradients, чтобы реализовать весь кластер horovod. коллективное слияние.
Конкретная логика для расчета градиента:
- TF вызывает метод calculate_gradients hvd.DistributedOptimizer:
- hvd.DistributedOptimizer сначала будет использовать официальный Optimizer.compute_gradients TF для расчета локального градиента;
- Затем используйте AllReduce, чтобы получить средний градиент каждого процесса;
- calculate_gradients возвращает список пар (градиент, вес). используется apply_gradients;
- TF вызывает метод apply_gradients hvd.DistributedOptimizer:
- Вызовите официальный Optimizer.apply_gradients TF, чтобы обработать входящие параметры и вернуть операцию, которая обновляет веса. TF может использовать это возвращаемое значение для последующей обработки;
Из-за проблем с версиями TF для анализа мы различаем версии 1.x и 2.x.
0x04 TensorFlow 1.x
Как упоминалось ранее, Horovod требует от разработчиков использовать собственный определенный Horovod hvd.DistributedOptimizer вместо официального оптимизатора TensorFlow, чтобы градиент можно было получить на этапе оптимизации модели, поэтому мы анализируем его из _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 в конструкторе__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
Первым шагом в вычислении градиента является вызов calculate_gradients для вычисления градиента потерь до указанного val_list и возврата списка кортежей.list(zip(grads, var_list)).
Каждая тензорная модель рабочего процесса будет вызывать calculate_gradients для каждой модели,
gradients = self._optimizer.compute_gradients(*args, **kwargs)Это градиент, рассчитанный локально по этой модели.
DistributedOptimizer переопределяет метод calculate_gradients() класса оптимизатора.
-
_DistributedOptimizerНастраивается при инициализацииself._allreduce_grads = _make_allreduce_grads_fn. Важно здесь. - Метод calculate_gradients() сначала вызывает calculate_gradients() официального оптимизатора исходной конфигурации TF. Возвращаемое значение calculate_gradients() — это список кортежей, каждый элемент списка — (градиент, переменная), а градиент — это значение градиента для каждого изменения переменной;
- Если установлен _agg_helper, то есть LocalGradientAggregationHelper, вызовите LocalGradientAggregationHelper для локального накопления градиента (цель локального накопления — уменьшить количество кросс-процессов, а кросс-процессное слияние будет выполняться только после определенного этапа), в противном случае вызовите расчет _allreduce_grads, то есть прямое слияние по процессам (allreduce рассчитанных распределенных градиентов с MPI);
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. То есть он обновляется на всех машинах каждый раз back_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 перебирает список локальных кортежей (градиент, переменная), накопленный в locally_aggregated_grads.
- allreduce_grads будет выполнять операцию обхода тензора и применения тензора, для каждого тензора функция _allreduce_grads_helper будет выполнять слияние между процессами.
4.3.1 _init_aggregation_vars
Функция _init_aggregation_vars перебирает список локальных кортежей (градиент, переменная), накопленный в 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
Метод calculate_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), #判断是否可以allreduce
lambda: self._allreduce_grads_helper(grads, vars), # 如果counter等于backward_passes_per_step，就对本地累积的梯度来应用跨进程allreduce
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:
- получить все оценки;
- Пройдитесь по списку кортежей (градиент, переменная), для каждого града используйте _allreduce_cond для синхронизации с другими воркерами;
- Наконец, верните список синхронизированных градиентов;
@_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.
- Если тип тензора — IndexedSlices, требуется только операция allgather, а необходимость других операций зависит от конкретной дополнительной конфигурации.
- Потому что для IndexedSlices, распределенных по разным воркерам, их значения и индексы не дублируют друг друга.
- Предположим, что индексы, распределенные по рабочему 1, равны [1, 3, 5, 7, 9], а индексы, распределенные по рабочему 2, равны [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 — это OP, связанные с тензорным потоком, настроенные HVD. _allreduce и allgather соответствуют им соответственно.
- _allreduce использует имя HorovodAllreduce для связывания с HorovodAllreduceOp и выполняет промежуточное преобразование с помощью MPI_LIB.horovod_allreduce;
- allgather связан с именем HorovodAllgather и HorovodAllgatherOp, а промежуточное преобразование выполняется с помощью MPI_LIB.horovod_allgather;
Таким образом вызывается соответствующая операция MPI.
В сочетании с конфигурацией пространства имен в предыдущем _make_cached_allreduce_grads_fn имя тензора будет примерно таким:DistributedAdam_Allreduce/cond_14/HorovodAllreduce_grads_5_0.
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 вызывается с несколькими аргументами, такими как тензор и имя. Среди них op=Sum является наиболее важным. Это используется внутри C++ для определения конкретной операции сокращения. Разберем подробно:
4.7.1 Определения 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, что соответствует ReduceOp::SUM C++.
4.7.3 Установление соединений
_allreduce продолжает вызывать:
MPI_LIB.horovod_allreduce(tensor, name=name, reduce_op=op
MPI_LIB.horovod_allreduce преобразуется в следующий код в мире C++
- Во-первых, reduce_op_ можно получить через конфигурацию OP_REQUIRES_OK;
- Во-вторых, в 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 Реализация Хоровода
Для TF2.x каждая строка кода выполняется последовательно, нет необходимости строить граф, а также отменяется control_dependency. Horovod может напрямую получать градиенты, вызывая API TensorFlow 2.0. Таким образом, реализация части Horovod, посвященной обновлению градиента, основана не на реализации графа расчета, а на использованииhvd.DistributedGradientTape.
Во время обучения работник выполняет следующие действия:
- Используйте DistributedGradientTape для инкапсуляции официальной ленты TF и настройки функции allreduce.
- Прочитайте набор обучающих данных.
- Убыток рассчитывается путем вызова функции прямого прохода в локальной модели.
- После заданной потери рабочий процесс использует механизм GradientTape нетерпеливого выполнения TensorFlow, чтобы вызвать функцию базового класса для получения градиента.
- Каждый рабочий вызовет Allreduce для синхронизации градиентов.
- Каждый рабочий будет обновлять модель соответственно в соответствии с последним градиентом.
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, который выполняет следующую работу:
- Измените область имени и добавьте суффикс
_Allreduce; - Сжатие, если настроено;
- В зависимости от типа операции вызовите allreduce или верните тензор напрямую;
- Область имени DistributedGradientTape была переписана на 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 сделал адаптацию и преобразование, чтобы мы могли напрямую вызывать функцию C++ из функции python.
HorovodAllreduceOp наследует AsyncOpKernel, является асинхронным OP TF и зарегистрирован в TF с помощью REGISTER_KERNEL_BUILDER, поэтому его можно встроить в процесс TF.
TF вызовет метод ComputeAsync, охватываемый HorovodAllreduceOp.Внутри 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_;
};
Начнем с рассмотрения Хоровода на Спарке.
0xEE Личная информация
★★★★★★Думая о жизни и технологиях★★★★★★
Публичный аккаунт WeChat:мысли Росси
ссылка 0xFF
tf.train.SyncReplicasOptimizer no synchronization among workers #11753
Синхронное обучение распределенного тензорного потока не синхронизируется между рабочими # 9596
tf.train.SyncReplicasOptimizer
Slow and Stale Gradients Can Win the Race: Error-Runtime Trade-offs in Distributed SGD
MPI, OpenMPI и глубокое обучение
Когда Spark встречается с принципом и практикой распределенной среды глубокого обучения TensorFlow
Анализ исходного кода TensorFlowOnSpark
Интерпретация TensorFlow SyncReplicasOptimizer
Метод обновления веса TensorFlow
Различные градиенты обработки градиента в Tensorflow
blog.CSDN.net/Edward_Make it out/…
Подробное описание tf.GradientTape: инструмент решения градиента
Обучение TensorFlow (4): GradientTape, оптимизатор и функция потерь (потери)
Интерпретация распределенного исходного кода tensorflow 4: AdamOptimizer
[TensorFlow] Анализ исходного кода оптимизатора AdamOptimizer