[源碼解析] 深度學習流水線并行GPipe (2) ----- 梯度累積

梯度累積是一種增大訓練時 batch size的技術，在本地使用 micro-batch 多次進行正向和反向傳播積累梯度后，再進行梯度規約和優化器更新，這是用來均攤通信成本的一種常用策略。本文通過幾個框架/庫的實現對比，讓大家對這個技術有進一步的了解。

[源碼解析] 深度學習流水線并行GPipe (2) ----- 梯度累積

0x00 摘要

梯度累積是一種增大訓練時 batch size的技術，在本地使用 micro-batch 多次進行正向和反向傳播積累梯度后，再進行梯度規約和優化器更新，這是用來均攤通信成本的一種常用策略。本文通過幾個框架/庫的實現對比，讓大家對這個技術有進一步的了解。

本系列其他文章如下：

[源碼解析] 深度學習流水線并行Gpipe(1)---流水線基本實現

0x01 概述

1.1 前文回顧

前文提到，目前分布式模型訓練有幾個必要并行技術：

流水并行，尤其是如何自動設定流水；
梯度累加（Gradient Accumulation）；
后向重計算；
1F1B 策略（我們將采用PipeDream分析）；

在前文中，我們介紹了Gpipe如何實施流水線并行技術。本文我們介紹梯度累加（Gradient Accumulation）。

0x02 基本概念

梯度累積是一種用來均攤通信成本的一種常用策略。它在本地使用 micro-batch 多次進行正向和反向傳播積累梯度后，再進行梯度規約和優化器更新，相當于擴大了N倍的batch size。

2.1 背景知識

深度學習模型由許多相互連接的層組成，樣本在這些層中進行傳播，具體傳播包含兩個過程：前向（forward）過程與反向（backword）過程。

前向過程是從輸入計算得到輸出。樣本在每一步都通過前向傳播進行傳播，在通過所有層傳播后，網絡為樣本生成預測，然后計算每個樣本的損失值，損失值意味著 “對于這個樣本，本網絡錯了多少？”。
然后就是反向過程。神經網絡在此過程中計算這些損失值相對于模型參數的梯度。可以認為著就是一個梯度累積的過程。
最后，這些梯度用于計算各個模型參數的更新。

訓練中，每個樣本的大小由超參數batch size指定，此參數的大小會對最終的模型效果產生很大的影響。一定條件下，batch size設置的越大，模型就會越穩定。

2.2 產生原因

累加梯度顧名思義就是累加后的梯度值。為什么要累加呢？因為運行內存不夠用。

在訓練模型時，如果一次性將所有訓練數據輸入到模型，經常會造成內存不足，這時候就需要把一個大 Batch 拆分成若干小批次數據（專業術語為mini-batch）。分成小批次后，帶來一個問題，那就是本來應該是所有數據全部送入后計算梯度再更新參數，現在成了每個小批次都要計算梯度更新參數，為了不這么頻繁計算梯度，于是就引入了累加梯度。也就是說：

將整個dataset分成多個batch；
分別將每個batch分成多個小批次，將每個小批次喂給神經網絡；
每個小批次雖然計算梯度，但是在每次反向傳播后，先不進行優化器的迭代更新。
經過若干個小批次后（即一個batch中的所有小批次），用每個小批次計算的梯度的累積和去進行優化器迭代更新參數、梯度清零的操作。

這樣就跟把全部數據一次性送入模型進行訓練效果一樣了。

2.3 本質

梯度累加本質上就是累加 accumulation_steps 個 batch_size/accumulation_steps 的梯度，再根據累加的梯度來更新網絡參數，以達到真實梯度類似batch_size 的效果。在使用時，需要注意適當的擴大學習率。

也就是說：

首先將整個dataset分成多個batch，每個 batch size = 32，且假定 accumulation steps = 8 ；
因為 batch size = 32 ，太大了，單機顯卡無法跑，于是我們在前向傳播的時候以 batch_size = 32 / 8 = 4 來計算梯度；
這樣就再分別將每個batch分成多個batch size 為 4 的小批次，將每個小批次逐一喂給神經網絡；
每個小批次雖然計算梯度，但是在每次反向傳播（在反向傳播的時候，會將mean_loss也除以8）后，先不進行優化器的迭代更新。
經過 accumulation steps 個小批次后（即一個batch中的所有小批次），用每個小批次計算梯度的累積和去進行優化器迭代更新參數。
最后進行梯度清零的操作。
處理下一個batch。

這樣就跟把 32 batch size 一次性送入模型進行訓練效果一樣了。

具體如下，時間軸是由左自右：

                                     +-------------------+
                                     |    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 數據并行

micro-batch 跟數據并行有高度的相似性：

數據并行是空間上的，數據被拆分成多個 tensor，同時喂給多個設備并行計算，然后將梯度累加在一起更新。
micro-batch 是時間上的數據并行，數據被拆分成多個 tensor，這些 tensor 按照時序依次進入同一個設備串行計算，然后將梯度累加在一起更新。

當總的 batch size 一致，且數據并行的并行度和 micro-batch 的累加次數相等時，數據并行和 Gradient Accumulation 在數學上完全等價。

Gradient Accumulation 通過多個 micro-batch的梯度累加使得下一個 micro-batch 的前向計算不需要依賴上一個 micro-batch 的反向計算，因此可以暢通無阻的進行下去（當然在一個大 batch 的最后一次 micro-batch 還是會觸發這個依賴）。

2.5 解決問題

Gradient Accumulation 解決了很多問題：

在單卡下，Gradient Accumulation 可以將一個大的 batch size 拆分成等價的多個小 micro-batch ，從而達到節省顯存的目的。
在數據并行下，Gradient Accumulation 解決了反向梯度同步開銷占比過大的問題（隨著機器數和設備數的增加，梯度的 AllReduce 同步開銷也加大），因為梯度同步變成了一個稀疏操作，因此可以提升數據并行的加速比。
在流水線并行下， Gradient Accumulation 使得不同 stage 之間可以并行執行不同的 micro-batch，通過多個 micro-batch的梯度累加使得下一個 micro-batch 的前向計算不需要依賴上一個 micro-batch 的反向計算，因此從而讓各個階段的計算不阻塞，可以暢通無阻的進行下去（當然在一個大 batch 的最后一次 micro-batch 還是會觸發這個依賴），達到流水線的目的。

0x03 PyTorch 梯度累積

3.1 自動累積

PyTorch默認會對梯度進行累加。即，PyTorch會在每一次backward()后進行梯度計算，但是梯度不會自動歸零，如果不進行手動歸零的話，梯度會不斷累加.

至于為什么PyTorch有這樣的特點，https://discuss.pytorch.org/t/why-do-we-need-to-set-the-gradients-manually-to-zero-in-pytorch/4903/9 這里給出了一個解釋。我們結合其他的解釋大致得出如下：

從PyTorch的設計原理上來說，在每次進行前向計算得到預測值時，會產生一個用于梯度回傳的計算圖，這張圖儲存了進行反向傳播需要的中間結果，當調用了.backward()后，會從內存中將這張圖進行釋放。
利用梯度累加，可以在最多保存一張計算圖的情況下進行多任務的訓練。在多任務中，對前面共享的張量進行了多次計算操作后，調用不同任務的backward()，那些張量的梯度會自動累加。
另外一個理由就是在內存大小不夠的情況下疊加多個batch的grad作為一個大batch進行迭代，因為二者得到的梯度是等價的。
由于PyTorch的動態圖和autograd機制，導致并沒有一個確切的點知道何時停止前向操作，因為你不知道什么時候一個計算會結束以及什么時候又會有一個新的開始。所以自動設置梯度為 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: 輸入圖像和標簽，通過計算得到預測值，計算損失函數；
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）在module級別實現數據并行性。其使用torch.distributed包communication collectives來同步梯度，參數和緩沖區。并行性在單個進程內部和跨進程均有用。

在這種情況下，雖然gradient accumulation 也一樣可以應用，但是為了提高效率，需要做相應的調整。

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個步驟，所以有K次 all_reduce。但實際上，每次梯度累加循環中，optimizer.step()只有一次，這意味著我們這K次 loss.backward() 之中，其實只進行一次 all_reduce 即可，前面 K - 1 次 all_reduce 是沒有用的。

3.3.2 DDP如何加速

于是我們就思考，是否可以在 loss.backward() 之中有一個開關，使得我們在前面K-1次 loss.backward() 之中只做反向傳播，不做梯度同步（累積）。

DDP 已經想到了這個問題，它提供了一個暫時取消梯度同步的context函數 no_sync()。在no_sync()context之下，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 的 forward 方法中可以看到，只有在 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;
  }

  // 省略后續代碼

有點繞，我們梳理一下：

一個 step 有兩個操作：forward 和 backward。

forward 操作時候 ：require_backward_grad_sync = True 意味著 forward 時候
- 設置 require_forward_param_sync = True。
- 會調用reducer.prepare_for_forward() 和 reducer.prepare_for_backward
- reducer.prepare_for_backward 意味著會設置 expect_autograd_hooks_ = true，expect_autograd_hooks_是關鍵。
backward 操作時候 ：
- expect_autograd_hooks_ = true 意味著 backward 時候進行進行all-reduce操作。
- 否則直接返回，不做 all-reduce操作。

即如下圖，

上半部分是 forward 的邏輯，就是 forward()函數，
下半部分是 backward 邏輯，就是 Reducer::autograd_hook() 函數。
expect_autograd_hooks_ 是forward 和 backward 之間串聯的關鍵之處。

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。這樣，backward 時候就不會進行 All-Reduce 操作。

0x04 Tensorflow實現

在 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 的流水并行示例中，每個“時間點” 可以在多個階段（stage）上同時做不同的micro-batch，圖中每個方塊中的標號表示了第幾個 micro-batch；同一個 micro-batch 還是串行的經過所有的 stage，在這種情況下，每個設備的空閑時間只有 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 應用

DenseLm12kWide41BAdam16x16 中有如何使用 ShardedAdam。

@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

0xFF 參考

[原創][深度][PyTorch] DDP系列第三篇：實戰與技巧

posted @ 2021-08-26 08:27 羅西的思考閱讀(2875) 評論(0) 收藏舉報

刷新頁面返回頂部

[源碼解析] 深度學習流水線并行GPipe (2) ----- 梯度累積

[源碼解析] 深度學習流水線并行GPipe (2) ----- 梯度累積

0x00 摘要

0x01 概述

1.1 前文回顧

0x02 基本概念

2.1 背景知識

2.2 產生原因

2.3 本質

2.4 VS 數據并行

2.5 解決問題

0x03 PyTorch 梯度累積

3.1 自動累積

3.2 代碼示例

3.3 DistributedDataParallel 的梯度累積

3.3.1 單卡模型梯度累計

3.3.2 DDP如何加速

3.3.3 no_sync實現

0x04 Tensorflow實現

0x05 Gpipe實現

5.1 優化器

5.2 包裝器

5.3 應用

0xFF 參考

公告