[Анализ исходного кода] Конвейер глубокого обучения, параллельный GPipe (2) ----- накопление градиента

машинное обучение глубокое обучение

0x00 сводка

Накопление градиента — это метод увеличения размера пакета во время обучения, который используется локально.micro-batchПосле многократного накопления градиентов посредством прямого и обратного распространения выполняется уменьшение градиента и обновления оптимизатора, что является общей стратегией, используемой для амортизации затрат на связь. Эта статья позволит вам лучше понять эту технологию, сравнив реализацию нескольких фреймворков/библиотек.

Другие статьи из этой серии:

[Анализ исходного кода] Конвейер глубокого обучения, параллельный Gpipe(1) --- Базовая реализация конвейера

0x01 Обзор

1.1 Предыдущий обзор

Как упоминалось выше, существует несколько необходимых параллельных технологий для обучения распределенной модели:

  • Параллельный конвейер, особенно как автоматически установить конвейер;
  • накопление градиента;
  • обратный пересчет;
  • стратегия 1F1B (будем использовать анализ PipeDream);

В предыдущей статье мы рассказали, как Gpipe реализует конвейерный параллелизм. В этой статье мы представляем Gradient Accumulation.

0x02 Основные понятия

Накопление градиентаЭто распространенная стратегия, используемая для амортизации затрат на связь. используется локальноmicro-batchПосле многократного прямого и обратного распространения для накопления градиентов выполняется уменьшение градиента и обновление оптимизатора, что эквивалентно увеличению размера пакета в N раз.

2.1 Фоновые знания

Модель глубокого обучения состоит из множества взаимосвязанных слоев, в которых образцы распространяются через два процесса: прямой процесс и обратный процесс.

  • Прямой процесс вычисляет результат на основе ввода. Выборки распространяются посредством прямого распространения на каждом шаге, и после прохождения всех слоев сеть генерирует прогнозы для выборок, а затем вычисляет значение потерь для каждой выборки, что означает «насколько эта сеть ошибается в отношении этой выборки?» .
  • Затем идет обратный процесс. В процессе нейронная сеть вычисляет градиенты этих значений потерь по отношению к параметрам модели. Его можно рассматривать как процесс накопления градиента.
  • Наконец, эти градиенты используются для вычисления обновлений отдельных параметров модели.

Во время обучения размер каждой выборки определяется размером пакета гиперпараметров, и размер этого параметра будет иметь большое влияние на окончательную производительность модели. При определенных условиях, чем больше размер партии, тем более стабильной будет модель.

2.2 Причины

Накопленные градиенты, как следует из названия, представляют собой накопленные значения градиента. Зачем добавлять? Потому что оперативной памяти не хватает.

При обучении модели, если все обучающие данные вводятся в модель одновременно, это часто приводит к нехватке памяти.В это время необходимо разбить большой пакет на несколько небольших пакетов данных (профессиональный термин мини- партия). После деления на мелкие партии возникает проблема, то есть градиент должен рассчитываться после отправки всех данных, а затем параметры обновляются.Теперь необходимо рассчитать параметры обновления градиента для каждой мелкой партии. чтобы не так часто рассчитывать градиент, поэтому вводится кумулятивный градиент. Это:

  • Разделите весь набор данных на несколько пакетов;
  • Разделите каждую партию на несколько мини-партий и передайте каждую мини-партию в нейронную сеть;
  • Хотя градиент рассчитывается для каждого мини-пакета, итеративное обновление оптимизатора не выполняется после каждого обратного распространения.
  • После нескольких небольших пакетов (то есть всех небольших пакетов в пакете) кумулятивная сумма градиентов, рассчитанных в каждом небольшом пакете, используется для выполнения оптимизатором итеративного обновления параметров и обнуления градиентов.

Это то же самое, что отправить все данные в модель за один раз для обучения.

2.3 Суть

Накопление градиента по сути является накоплениемaccumulation_stepsКусокbatch_size/accumulation_stepsГрадиент , а затем обновить параметры сети в соответствии с накопленным градиентом для достижения аналогичного реального градиентаbatch_sizeЭффект. При его использовании необходимо обратить внимание на соответствующее расширение скорости обучения.

Это:

  • Сначала разделите весь набор данных на несколько пакетов, размер каждого пакета = 32, и предположим,accumulation steps = 8;
  • так какbatch size = 32, слишком большой, одномашинная видеокарта не может работать, поэтому мы используем прямое распространениеbatch_size = 32 / 8 = 4вычислить градиент;
  • Таким образом, каждый пакет делится на несколько пакетов с размером пакета 4, и каждый небольшой пакет передается в нейронную сеть один за другим;
  • Хотя градиент рассчитывается для каждой небольшой партии, итеративное обновление оптимизатора не выполняется после каждого обратного распространения (mean_loss также делится на 8 при обратном распространении).
  • проходить черезaccumulation stepsПосле каждого мини-пакета (т. е. всех мини-пакетов в пакете) каждый мини-пакет используется для вычисления кумулятивной суммы градиентов, чтобы оптимизатор итеративно обновлял параметры.
  • Наконец, выполняется операция очистки градиента.
  • Обработайте следующую партию.

Это то же самое, что отправить 32 размера партии в модель за один раз для обучения.

детали следующим образом:

                                     +-------------------+
                                     |    GLOBAL BATCH   +--------------------------+
                                     +-------------------+                          |
                                                                                    |
                                                                                    |
 +<---------------------------------------------------------------------------------+
 |
 |
 |    +--------------+     +--------------+     +--------------+     +--------------+
 +--> | MINI BATCH 0 +---->+ MINI BATCH 1 +---->+ MINI BATCH 2 +---->+ MINI BATCH 3 |
      +-----+--------+     +-------+------+     +------+-------+     +-------+------+
            |                      |                   |                     |
            |                      |                   |                     |
            |                      |                   |                     |
            v                      v                   v                     v
       +----+-----+          +-----+-----+       +-----+-----+          +----+-----+
       |  grad 0  |          |  grad 1   |       |  grad 2   |          |  grad 3  |
       +----+-----+          +-----+-----+       +-----+-----+          +----+-----+
            |                      |                   |                     |
            |                      |                   |                     |
            |                      |                   |                     |
            v                      v                   v                     v
     +------+----------------------+-------------------+---------------------+------+
     |                                                                              |
     |                              GLOBAL BATCHGRADIENTS                           |
     |                                                                              |
     +------------------------------------------------------------------------------+
​
​
+------------------------------------------------------------------------------------>
                                                                        Time
​

2.4 Параллелизм данных VS

Микропакет очень похож на параллелизм данных:

  • Параллелизм данных является пространственным, данные разбиваются на несколько тензоров, передаются на несколько устройств для параллельных вычислений, а затем градиенты накапливаются и обновляются вместе.
  • Микропакет — это параллелизм данных во времени, данные разбиваются на несколько тензоров, эти тензоры поступают в одно и то же устройство для последовательного расчета по временным рядам, а затем градиенты накапливаются и обновляются вместе.

Когда общий размер пакета одинаков, а параллелизм параллелизма данных и количество накоплений микропакетов равны, параллелизм данных и градиентное накопление математически эквивалентны.

Накопление градиента накапливает градиенты нескольких микропакетов, так что прямой расчет следующего микропакета не должен зависеть от обратного расчета предыдущего микропакета, поэтому он может продолжаться беспрепятственно (конечно, в последний раз). большой партии). Микропакет по-прежнему вызывает эту зависимость).

2.5 Решение проблем

Градиентное накопление решает многие проблемы:

  • На одной карте Gradient Accumulation может разделить большой пакет на несколько эквивалентных небольших микропакетов для экономии видеопамяти.
  • При параллелизме данных Gradient Accumulation решает проблему, заключающуюся в том, что накладные расходы на обратную синхронизацию градиента слишком велики (по мере увеличения количества машин и устройств накладные расходы на синхронизацию градиентов AllReduce также увеличиваются), поскольку градиентная синхронизация становится разреженной операцией, поэтому ускорение Параллелизм данных может быть улучшен.
  • В условиях конвейерного параллелизма накопление градиентов позволяет выполнять различные микропакеты параллельно между различными этапами Благодаря накоплению градиентов нескольких микропакетов прямой расчет следующего микропакета не должен зависеть от обратного расчета предыдущего микропакета, поэтому расчет каждого этапа не блокируется и может протекать беспрепятственно (конечно, эта зависимость все равно будет срабатывать в последнем микропакете большого пакета), чтобы достичь цели трубопровод.

0x03 Накопление градиента PyTorch

3.1 Автоматическое накопление

PyTorch по умолчанию накапливает градиенты. То есть PyTorch будетbackward()После выполнения вычисления градиента градиент не будет автоматически обнулен, если не выполнить ручное обнуление, градиент будет продолжать накапливаться.

Что касается того, почему у PyTorch есть такая функция,обсуждение.py torch.org/he/why-do-major…Здесь дается объяснение. В сочетании с другими объяснениями мы можем сделать примерно следующие выводы:

  • В соответствии с принципом разработки PyTorch каждый раз, когда для получения прогнозируемого значения выполняется прямой расчет, будет генерироваться расчетный график для возврата градиента. В этом графике хранятся промежуточные результаты, необходимые для обратного распространения. Когда после обратного() изображение будет освобожден из памяти.
  • Используя накопление градиента, можно выполнять многозадачное обучение, сохраняя при этом не более одного вычислительного графа. В многозадачности после выполнения нескольких вычислительных операций над ранее общими тензорами, вызывая reverse() различных задач, градиенты этих тензоров будут автоматически накапливаться.
  • Другая причина заключается в наложении градаций нескольких пакетов в виде большого пакета для итерации, когда объема памяти недостаточно, потому что градиенты, полученные двумя способами, эквивалентны.
  • Из-за динамического графа PyTorch и механизма автоградации нет точной точки, чтобы знать, когда остановить прямые операции, потому что вы не знаете, когда закончится вычисление и когда начнется новое. Поэтому автоматическая установка градиента на 0 сложна.

3.2 Примеры кода

Традиционный пример кода приведен ниже:

for i,(images,target) in enumerate(train_loader):
    # 1. input output
    images = images.cuda(non_blocking=True)
    target = torch.from_numpy(np.array(target)).float().cuda(non_blocking=True)
    outputs = model(images)
    loss = criterion(outputs,target)
  
    # 2. backward
    optimizer.zero_grad()   # reset gradient
    loss.backward()
    optimizer.step()

Затем приведите пример накопления градиента:

  • Получить потери: введите изображение и метку, получите прогнозируемое значение путем расчета и рассчитайте функцию потерь;
  • loss.backward()Обратное распространение, рассчитать текущий градиент;
  • Повторите шаги 1-2 несколько раз, не очищая градиент, чтобы градиент накапливался на существующем градиенте;
  • После накопления градиента определенное количество раз первыйoptimizer.step()Обновите параметры сети на основе накопленных градиентов, затемoptimizer.zero_grad()Очистите прошлые градиенты и подготовьтесь к следующей волне накопления градиентов;
for i, (images, target) in enumerate(train_loader):
    # 1. input output
    images = images.cuda(non_blocking=True)
    target = torch.from_numpy(np.array(target)).float().cuda(non_blocking=True)
    outputs = model(images) # 前向传播
    loss = criterion(outputs, target) # 计算损失
​
    # 2. backward
    loss.backward() # 反向传播，计算当前梯度
    
     # 3. update parameters of net
    if ((i+1)%accumulation)==0:
        # optimizer the net
        optimizer.step() # 更新网络参数
        optimizer.zero_grad() # reset grdient # 清空过往梯度

3.3 Градиентное накопление DistributedDataParallel

DistributedDataParallel (DDP) реализует параллелизм данных на уровне модулей. Он использует коллективы связи пакетов torch.distributed для синхронизации градиентов, параметров и буферов. Параллелизм полезен как внутри одного процесса, так и между процессами.

В этом случае, хотя накопление градиента также может быть применено, для повышения эффективности его необходимо соответствующим образом скорректировать.

3.3.1 Накопление градиента модели с одной картой

Сначала вспомним однокарточную модель, то есть как накапливать градиенты в обычных ситуациях.

# 单卡模式，即普通情况下的梯度累加
for data in enumerate(train_loader # 每次梯度累加循环
    optimizer.zero_grad()
    for _ in range(K):
        prediction = model(data / K)
        loss = loss_fn(prediction, label) / K
        loss.backward()  # 积累梯度，不应用梯度改变，执行K次
    optimizer.step()  # 应用梯度更新，更新网络参数，执行一次

В операторе loss.backward() DDP выполняет уменьшение градиента all_reduce.

Поскольку в каждом цикле накопления градиента есть K шагов, all_reduce выполняется K раз. Но на самом деле optimizer.step() встречается только один раз в каждом цикле накопления градиента, а это значит, что мыФактически среди K раз loss.backward() может быть выполнен только один all_reduce, а предыдущий K - 1 all_reduce бесполезен..

3.3.2 Как ускоряется DDP

Поэтому мы подумали о том, можем ли мы иметь переключатель в loss.backward(), чтобы мы выполняли обратное распространение только в предыдущем K-1 раз loss.backward() без градиентной синхронизации (накопления).

DDP подумал об этой проблеме, он предоставляет контекстную функцию, которая временно отменяет синхронизацию градиента.no_sync(). существуетno_sync()В контексте DDP не выполняет градиентную синхронизацию. Но в первый раз после завершения контекста no_sync()forward-backwardбудет синхронизироваться.

Окончательный код выглядит следующим образом:

model = DDP(model)
​
for data in enumerate(train_loader # 每次梯度累加循环
    optimizer.zero_grad()
    
    for _ in range(K-1):# 前K-1个step 不进行梯度同步（累积梯度）。
        with model.no_sync(): # 这里实施“不操作”
            prediction = model(data / K)
            loss = loss_fn(prediction, label) / K
            loss.backward()  # 积累梯度，不应用梯度改变
    
    prediction = model(data / K)
    loss = loss_fn(prediction, label) / K 
    loss.backward()  # 第K个step 进行梯度同步（累积梯度）
    optimizer.step() # 应用梯度更新，更新网络参数  

3.3.3 реализация no_sync

Код для no_sync выглядит следующим образом:

    @contextmanager
    def no_sync(self):
        r"""
        A context manager to disable gradient synchronizations across DDP
        processes. Within this context, gradients will be accumulated on module
        variables, which will later be synchronized in the first
        forward-backward pass exiting the context.
​
        Example::
​
            >>> ddp = torch.nn.parallel.DistributedDataParallel(model, pg)
            >>> with ddp.no_sync():
            >>>   for input in inputs:
            >>>     ddp(input).backward()  # no synchronization, accumulate grads
            >>> ddp(another_input).backward()  # synchronize grads
        """
        old_require_backward_grad_sync = self.require_backward_grad_sync
        self.require_backward_grad_sync = False
        try:
            yield
        finally:
            self.require_backward_grad_sync = old_require_backward_grad_sync

Как это использовать? мы вDistributedDataParallelКак вы можете видеть в прямом методе , только когда для require_backward_grad_sync установлено значение True, будут вызываться reducer.prepare_for_forward() и reducer.prepare_for_backward, а для require_forward_param_sync будет установлено значение True.

   def forward(self, *inputs, **kwargs):
    
        with torch.autograd.profiler.record_function("DistributedDataParallel.forward"):
            
            self.reducer.save_thread_local_state()
            if torch.is_grad_enabled() and self.require_backward_grad_sync:
                # True时候才会进入
                self.logger.set_runtime_stats_and_log()
                self.num_iterations += 1
                self.reducer.prepare_for_forward()
            
            # 省略部分代码
​
            if torch.is_grad_enabled() and self.require_backward_grad_sync:
                # True时候才会进入
                self.require_forward_param_sync = True
                if self.find_unused_parameters and not self.static_graph:
                    # Do not need to populate this for static graph.
                    self.reducer.prepare_for_backward(list(_find_tensors(output)))
                else:
                    self.reducer.prepare_for_backward([])
            else:
                self.require_forward_param_sync = False
​
            # 省略部分代码

Давайте взглянем на два метода Reducer.

prepare_for_forward просто выполняет статистическую работу и может быть проигнорирован.

void Reducer::prepare_for_forward() {
  std::lock_guard<std::mutex> lock(mutex_);
  num_iterations_++;
  if (should_collect_runtime_stats()) {
    record_forward_compute_start_time();
  }
}

prepare_for_backward выполнит сброс и подготовительную работу,С накоплением градиента связано expect_autograd_hooks_ = true.

void Reducer::prepare_for_backward(
    const std::vector<torch::autograd::Variable>& outputs) {
  std::lock_guard<std::mutex> lock(mutex_);
​
  // Reset accounting.
  expect_autograd_hooks_ = true; // 这里是关键
  reset_bucket_counting();
​
  // Reset unused parameter accounting.
  has_marked_unused_parameters_ = false;
  // Reset per iteration marked ready parameters.
  perIterationReadyParams_.clear();
​
  // If static graph is not set, search graph to detect unused parameters.
  // When static graph is set, unused_parameters_ will be detected and will
  // not change after 1st iteration.
  // If static_graph_ = false and find_unused_parameters_ is false,
  // we assume that autograd hooks for ALL variables will be called,
  // and we don't have to search the autograd graph for presence of these hooks.
  if (dynamic_graph_find_unused()) {
    unused_parameters_.clear();
    search_unused_parameters(outputs);
  }
}

expect_autograd_hooks_ = true как использовать? В Reducer::autograd_hook, если операция all-reduce не требуется, она вернется напрямую.

void Reducer::autograd_hook(VariableIndex index) {
    
  std::lock_guard<std::mutex> lock(this->mutex_);
​
  // Carry over thread local state from main thread. This allows for
  // thread-local flags such as profiler enabled to be configure correctly.
  at::ThreadLocalStateGuard g(thread_local_state_);
​
  // Ignore if we don't expect to be called.
  // This may be the case if the user wants to accumulate gradients
  // for number of iterations before reducing them.
  if (!expect_autograd_hooks_) { // 如果不需要进行all-reduce操作，则直接返回。
    return;
  }
​
  // 省略后续代码

Это немного запутанно, давайте разберемся:

Шаг имеет две операции: вперед и назад.

  • Во время движения вперед:require_backward_grad_sync = True означает, что вперед

    • Установите require_forward_param_sync = True.
    • вызовет редукторы.prepare_for_forward() и редукторы.prepare_for_backward
    • reducer.prepare_for_backward означает, что будет установлено expect_autograd_hooks_ = true, ключ — expect_autograd_hooks_.

  • При работе назад:

    • expect_autograd_hooks_ = true означает, что операции all-reduce выполняются в обратном порядке.
    • В противном случае вернитесь напрямую без операции all-reduce.

То есть, как показано на рисунке ниже,

  • Верхняя часть — это логика переадресации, то есть функция forward(),
  • Вторая половина — это обратная логика, то есть функция Reducer::autograd_hook().
  • expect_autograd_hooks_ — это ключ к конкатенации между прямым и обратным.
forward
+---------------------------------------------------------------------------------+
| forward()                                                                       |
|                                                                                 |
|                require_backward_grad_sync == True?? +---------+                 |
|                             +                                 |                 |
|                             |                                 |                 |
|                             | Yes                             |                 |
|                             |                                 | No              |
|                             v                                 |                 |
|                 reducer.prepare_for_forward                   |                 |
|                             +                                 |                 |
|                             |                                 |                 |
|                             |                                 |                 |
|                             v                                 |                 |
|                 reducer.prepare_for_backward                  |                 |
|                             +                                 |                 |
|                             |                                 |                 |
|                             |                                 |                 |
|                             v                                 v                 |
|                 expect_autograd_hooks_ = true    expect_autograd_hooks_ = false |
|                             +                                 +                 |
|                             |                                 |                 |
+---------------------------------------------------------------------------------+
                              |                                 |
+--------------------------------------------------------------------------------+
 backward                     |                                 |
                              |                                 |
 +--------------------------------------------------------------------------------+
 |                            |                                 |                 |
 | Reducer::autograd_hook()   |                                 |                 |
 |                            |                                 |                 |
 |                            |    +----------------------------+                 |
 |                            |    |                                              |
 |                            |    |                                              |
 |                            v    v                                              |
 |                 expect_autograd_hooks_ == True?? +------------+                |
 |                            +                                  |                |
 |                            | Yes                              |                |
 |                            |                                  |  No            |
 |                            v                                  v                |
 |                      Do All-Reduce                          Return             |
 |                                                                                |
 |                                                                                |
 +--------------------------------------------------------------------------------+
​

Операция no_sync означает установку require_backward_grad_sync = False, что в конечном итоге устанавливает expect_autograd_hooks_ = False. Таким образом, операция All-Reduce не будет выполняться при обратном направлении..

0x04 Реализация тензорного потока

В pytorch градиент накапливается по умолчанию до тех пор, пока он не обнуляется, поэтому реализовать указанную выше задачу несложно. Но в Tensorflow это не так просто.

Мы получаем пример кода из stackoverflow следующим образом:

## 定义优化器
opt = tf.train.AdamOptimizer()
​
## 得到你模型中的所有可训练变量
tvs = tf.trainable_variables()
​
# 用于记录每个变量的累积梯度，初始化为0s
accum_vars = [tf.Variable(tf.zeros_like(tv.initialized_value()), trainable=False) for tv in tvs]
# 定义清零操作
zero_ops = [tv.assign(tf.zeros_like(tv)) for tv in accum_vars]
​
## 使用优化器的compute_gradients来计算梯度
gvs = opt.compute_gradients(rmse, tvs)
​
## 将当前梯度累加在之前定义的变量上
accum_ops = [accum_vars[i].assign_add(gv[0]) for i, gv in enumerate(gvs)]
​
## 定义训练step，梯度下降，更新参数
train_step = opt.apply_gradients([(accum_vars[i], gv[1]) for i, gv in enumerate(gvs)])
​
## 训练循环
while ...:
    # 使用 zero_ops 初始化
    sess.run(zero_ops)
    # 使用accum_ops对accum_vars进行'n_minibatches'次梯度累积
    for i in xrange(n_minibatches):
        sess.run(accum_ops, feed_dict=dict(X: Xs[i], y: ys[i]))
    # 使用累积的梯度进行参数更新
    sess.run(train_step)

0x05 Реализация Gpipe

В примере параллельного конвейера GPipe каждый «момент времени» может выполнять разные микропакеты на нескольких этапах одновременно. Метки в каждом квадрате на рисунке указывают количество микропакетов; тот же самый микропакет также проходит через все этапы последовательно, в данном случае только около 25% времени простоя на устройство.

\

在这里插入图片描述

Конкретный код выглядит следующим образом:

5.1 Оптимизатор

В lingvo/core/optimizer.py есть конкретная реализация GradientAggregationOptimizer, ключевой код — apply_gradients, а логика такова:

  • Если _num_micro_batches равно 1, это означает, что накопление градиента не требуется, и применяется непосредственно apply_gradients;

  • Просматривайте список grads_and_vars и накапливайте градиенты;

  • accum_step — условие накопления градиента:

    • Если достигнуто количество итераций минипакета, вызывается _ApplyAndReset:

      • вызовите apply_gradients, чтобы применить градиенты;
      • Вызовите zero_op, чтобы очистить градиент;

    • В противном случае вызывается _Accum, что на самом деле является no_op и не операцией;

Конкретный код выглядит следующим образом:

  def apply_gradients(self, grads_and_vars, global_step=None, name=None):
    if self._num_micro_batches == 1:
      return self._opt.apply_gradients(grads_and_vars, global_step)
    global_step = global_step or py_utils.GetOrCreateGlobalStepVar()
    with tf.init_scope():
      self._create_slots([v for (_, v) in grads_and_vars])
​
    accums = []
    variables = []
​
    # 遍历，累积梯度
    for g, v in grads_and_vars:
      accum = self.get_slot(v, 'grad_accum')
      variables.append(v)
      # pytype: disable=attribute-error
      if isinstance(g, tf.IndexedSlices):
        scaled_grad = tf.IndexedSlices(
            g.values / self._num_micro_batches,
            g.indices,
            dense_shape=g.dense_shape)
      else:
        scaled_grad = g / self._num_micro_batches
      accum_tensor = accum.read_value()
      accums.append(accum.assign(accum_tensor + scaled_grad))
      # pytype: enable=attribute-error
​
    # 应用梯度，清零梯度
    def _ApplyAndReset():
      normalized_accums = accums
      if self._apply_crs_to_grad:
        normalized_accums = [
            tf.tpu.cross_replica_sum(accum.read_value()) for accum in accums
        ]
      apply_op = self._opt.apply_gradients(
          list(zip(normalized_accums, variables)))
      with tf.control_dependencies([apply_op]):
        zero_op = [tf.assign(accum, tf.zeros_like(accum)) for accum in accums]
      return tf.group(zero_op, tf.assign_add(global_step, 1))
​
    # 累积函数，其实是不做操作
    def _Accum():
      return tf.no_op()
​
    # 梯度累积条件，如果达到了小批次迭代数目，则应用梯度，清零梯度，否则就不做操作
    accum_step = tf.cond( 
        tf.equal(
            tf.math.floormod(self._counter + 1, self._num_micro_batches), 0),
        _ApplyAndReset,  # Apply the accumulated gradients and reset.
        _Accum)  # Accumulate gradients.
​
    with tf.control_dependencies([tf.group(accums)]):
      return tf.group(accum_step, tf.assign_add(self._counter, 1))

5.2 Обертки

ShardedAdam — это оболочка для GradientAggregationOptimizer и ShardedAdamOptimizer, и пользователи могут использовать ее напрямую.

class ShardedAdam(optimizer.Adam):
  """Adam optimizer wrapper that shards the slot variables."""
​
  @classmethod
  def Params(cls):
    params = super().Params()
    params.Define('num_micro_batches', 1, 'Number of accumulated batches.')
    return params
​
  def GetOptimizer(self, lr):
    p = self.params
    opt = ShardedAdamOptimizer(
        learning_rate=lr,
        beta1=p.beta1,
        beta2=p.beta2,
        epsilon=p.epsilon,
        name=p.name)
    if p.num_micro_batches > 1:
      tf.logging.info('Applying gradient aggregation.')
      
      opt = optimizer.GradientAggregationOptimizer( # 应用梯度累积
          opt, p.num_micro_batches, apply_crs_to_grad=True)
      self._cached_opt = opt
    return opt

5.3 Применение

Как использовать ShardedAdam в DenseLm12kWide41BAdam16x16.

@model_registry.RegisterSingleTaskModel
class DenseLm12kWide41BAdam16x16(DenseLm128B16x16):
  """41B params LM model with 2D split and ADAM optimizer on v3-512."""
​
  # Each layer has 1.6875B parameters.
  SEQUENCE_LENGTH = 2048
  NUM_DEVICES_PER_SPLIT = 512
  BATCH_DIM_PER_DEVICE = 0.5  # Total batch size 256
  DEVICE_MESH_SHAPE = [16, 32]
  DEVICE_MESH = gshard_utils.GetNonPod2dMesh(DEVICE_MESH_SHAPE, [16, 16, 2])
  NUM_TRANSFORMER_LAYERS = 24
  HIDDEN_DIM = 48 * 1024
  MODEL_DIM = 12 * 1024
  NUM_HEADS = 96
  ATTENTION_KEY_VALUE_DIM = 128
  GATED_GELU = False
  POSITIONAL_EMBEDDING = True
  NUM_MICRO_BATCHES = 1
​
  def Task(self):
    p = super().Task()
    
    # 使用ShardedAdam
    p.train.optimizer = ShardedAdam.Params().Set(
        beta1=0.9,
        beta2=0.999,
        epsilon=1e-6,
        num_micro_batches=self.NUM_MICRO_BATCHES)
    return p

0xEE Личная информация

★★★★★★Думая о жизни и технологиях★★★★★★

Публичный аккаунт WeChat:мысли Росси

ссылка 0xFF

[Оригинал][Глубина][PyTorch] Серия DDP, часть 3: реальный бой и навыки

Категории