[原始碼解析] PyTorch分散式優化器(2)----資料並行優化器

model=torch.nn.DaraParallel(model);
optimizer = torch.optim.SGD(model.parameters(), args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)

for batch_idx, (data, label) in pbar:   
    if args.cuda:
        data,label= data.cuda(),label.cuda(); # 資料放到了預設GPU
    data_v = Variable(data)
    target_var = Variable(label)
    
    prediction= model(data_v,target_var,args) # 多執行緒並行前向傳播
    criterion = nn.CrossEntropyLoss()
    loss = criterion(prediction,target_var) # 在預設GPU之上計算loss
    
    optimizer.zero_grad()      
    loss.backward()  # 多執行緒並行後向傳播    
    optimizer.step() # 更新引數

我們給出一個簡化的圖示如下，每個thread進行梯度計算，最後把梯度歸併到GPU 0，在GPU 0之上進行優化：

             Forward                                                    Backward
      +-------------------+                                       +------------------+
  +-->+ Thread 0 on  GPU0 +--+                                +-->+ Thread 1 on GPU0 +-+
  |   +-------------------+  |          GPU 0                 |   +------------------+ |
  |   +-------------------+  | output +---------------+ loss  |   +------------------+ |
+---->+ Thread 1 on  GPU1 +---------> |  Compute Loss +---------->+ Thread 2 on GPU1 +---+
| |   +-------------------+  |        +---------------+       |   +------------------+ | |
| |   +-------------------+  |                                |   +------------------+ | |
| +-->+ Thread 2 on  GPU2 +--+                                +-->+ Thread 3 on GPU2 +-+ |
|     +-------------------+                                       +------------------+   |
|                                                                                        |
|                                                                                        |
|                                    GPU 0                                               |
|   Model                  +-------------------------+    gradient                       |
+--------------------------+     optimizer.step      |  <--------------------------------+
                           +-------------------------+

0x03 DDP 之中的優化器

下圖來自快手八卦的論文，圖中羅列了原生訓練過程與DDP/Horovod的對比。

上面的 vanilla 就是原生訓練過程，其中 U 部分對應的就是優化器過程。常規優化器主要功能就是根據梯度來更新模型當前引數： w.data -= w.grad * lr。
下面部分就是DDP/Horovod優化過程，可以看到，其後向計算和歸併梯度是部分並行的。

3.1 流程

DDP 之中，依然使用的是普通優化器，但採用的是多程式方式，每個程式都完成訓練的全部流程，只是在後向計算時候需要使用 all-reduce 來歸併梯度。每個程式擁有自己獨立的優化器，優化器也是常規優化器。

這裡有兩個特點：

每個程式維護自己的優化器，並在每次迭代中執行一個完整的優化步驟。雖然這可能看起來是多餘的，但由於梯度已經聚合（gather）並跨程式平均，因此梯度對於每個程式都是相同的，這意味著不需要引數廣播步驟，減少了在節點之間傳輸張量所花費的時間。
All-Reduce 操作是在後向傳播之中完成的。
- 在 DDP 初始化時候會生成一個Reducer，其內部會註冊 autograd_hook。
- autograd_hook 在反向傳播時候進行梯度同步。

DDP 選擇了在 PyTorch 核心角度修改，在 DistributedDataParallel 模型的初始化和前向操作中做了處理。

具體邏輯如下：

DDP 使用多程式並行載入資料，在 host 之上，每個worker程式都會把資料從硬碟載入到 page-locked memory。分散式 minibatch sampler 保證每個程式載入到的資料是彼此不重疊的。
不需要廣播資料，而是並行把 minibatch 資料從 page-locked memory 載入到每個GPU，每個GPU都擁有模型的一個副本，所以也不需要拷貝模型。
在每個GPU之上執行前向傳播，計算輸出，每個GPU都執行同樣的訓練，不需要有主 GPU。
在每個GPU之上計算損失，執行後向傳播來計算梯度，在計算梯度同時對梯度執行all-reduce操作。
更新模型引數。因為每個GPU都從完全相同的模型開始訓練，並且梯度被all-reduced，因此每個GPU在反向傳播結束時最終得到平均梯度的相同副本，所有GPU上的權重更新都相同，這樣所有 worker 上的模型都一致，也就不需要模型同步了。

因為也是在模型的前向後向操作之中進行修改，所以優化器也不需要修改，每個worker分別在自己本地程式之中進行優化。

3.2 優化器狀態

這裡要留意的是，如何保證各個程式的優化器狀態相同？

DDP 與優化器實際上沒有關聯，DDP不對此負責，所以需要使用者協同操作來保證各程式之間的優化器狀態相同。這就圍繞著兩個環節：

優化器引數初始值相同。
- 優化器初始值相同由 "使用者在DDP模型建立後才初始化optimizer" 來確保。
優化器引數每次更新值相同。
- 每次更新的梯度都是all-reduce過的，所以各個優化器拿到的梯度delta數值是一樣的。

3.3 使用

其示例如下：

model = ToyModel().to(rank)
# 構造DDP model
ddp_model = DDP(model, device_ids=[rank])

loss_fn = nn.MSELoss()
# 優化器要在構造DDP model之後，才能初始化。
optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

optimizer.zero_grad()
outputs = ddp_model(torch.randn(20, 10))
labels = torch.randn(20, 5).to(rank)
loss_fn(outputs, labels).backward()
optimizer.step()

圖示如下：

+--------------------------------------------------------------------------------------+
| Process 1 on GPU 1                                                                   |
|                              +------------------------------+                        |
|                              | Backward                     |                        |
|                              |                              |                        |
| Forward +---->  Loss +-----> |  Compute  +---->  ALL+REDUCE | +---->  Optimizer.step |
|                              |                     ^        |                        |
|                              |                     |        |                        |
|                              +------------------------------+                        |
|                                                    |                                 |
|                                                    |                                 |
+--------------------------------------------------------------------------------------+
                                                     |
                                                     |
                                                     |
                                                     |
                                                     +
                                                 SYNC GRADS
                                                     +
                                                     |
                                                     |
                                                     |
+--------------------------------------------------------------------------------------+
| Process 2 on GPU 2                                 |                                 |
|                                                    |                                 |
|                              +------------------------------+                        |
|                              | Backward            |        |                        |
|                              |                     v        |                        |
| Forward +---->  Loss +-----> |  Compute  +---->  ALL+REDUCE | +---->  Optimizer.step |
|                              |                              |                        |
|                              |                              |                        |
|                              +------------------------------+                        |
|                                                                                      |
+--------------------------------------------------------------------------------------+

0x04 Horovod 的優化器

Horovod 並沒有對模型 fw/bw 進行修改（可能因為沒有Facebook自己修改那麼順手），而是對優化器進行了修改，實現了一個 DistributedOptimizer。

我們以 horovod/torch/optimizer.py 為例。

An optimizer that wraps another torch.optim.Optimizer, using an allreduce to
combine gradient values before applying gradients to model weights.

Allreduce operations are executed after each gradient is computed by ``loss.backward()``
in parallel with each other. The ``step()`` method ensures that all allreduce operations are
finished before applying gradients to the model.

DistributedOptimizer 包裝了另一個torch.optim.optimizer，其作用是：

在worker 並行執行loss.backward()計算出每個梯度之後，在 "將梯度應用於模型權重之前“ 這個時間點使用allreduce來合併梯度。
使用step()方法來確保所有allreduce操作在將梯度應用於模型之前會完成。

其具體實現是 _DistributedOptimizer，而_DistributedOptimizer對於梯度的歸併有兩個途徑，一個是通過 hook，一個是顯性呼叫了 synchronize 函式，我們接下來逐一介紹。

4.1 hook 同步梯度

hook 就是採用了 PyTorch 的 hook 方法，和 DDP 的思路非常類似，即在梯度計算函式之上註冊了hook，其作用是在計算完梯度之後呼叫hook，這樣all-reduce 就是在計算梯度過程中自動完成的，不需要等待 step 方法顯式呼叫來完成（類似 DP 那樣），具體來說就是：

在每個GPU之上計算損失，執行後向傳播來計算梯度，在計算梯度同時對梯度執行all-reduce操作。
更新模型引數。因為每個GPU都從完全相同的模型開始訓練，並且梯度被all-reduced，因此每個GPU在反向傳播結束時最終得到平均梯度的相同副本，所有GPU上的權重更新都相同，也就不需要模型同步了。

注：程式碼主要分為兩部分，處理 groups 相關和普通情況。

groups 是 PyTorch 的相關配置，作用是把梯度 allreduce 操作放在一起進行，因為程式碼比較複雜並且與本文主體邏輯不相關，所以我們略過這部分，只看普通非分組情況。

groups: The parameter to group the gradient allreduce ops. Accept values is a
        non-negative integer or a list of list of tf.Variable.
        If groups is a non-negative integer, it is the number of groups to assign
        gradient allreduce ops to for explicit grouping.
        If groups is a list of list of tf.Variable. Variables in the same
        inner list will be assigned to the same group, while parameter that does
        not appear in any list will form a group itself.
        Defaults as None, which is no explicit groups.

4.1.1 註冊 hooks

Hook 功能分為兩步驟，第一部分是註冊 hooks。

    def _register_hooks(self):

        if self._groups is not None: # groups，有興趣同學可以自行研究，可以理解為把梯度分組
            p_list = []
            # Get list of parameters with grads
            for param_group in self.param_groups:
                for p in param_group['params']:
                    if p.requires_grad:
                        p_list.append(p)

            # To ensure parameter order and group formation is consistent, broadcast p_list order
            # from rank 0 and use for every worker
            p_list_names = [self._parameter_names.get(p) for p in p_list]
            p_list_names = broadcast_object(p_list_names, root_rank=0)
            p_list = sorted(p_list, key=lambda p : p_list_names.index(self._parameter_names.get(p)))

            # Form groups
            if isinstance(self._groups, list):
                p_groups = []
                grouped_id = set()
                p_list_ids = [id(p) for p in p_list]
                for group in self._groups:
                    p_groups.append([p for p in group if id(p) in p_list_ids])
                    for p in p_groups[-1]:
                        grouped_id.add(id(p))
                for p in p_list:
                    if id(p) not in grouped_id:
                        p_groups.append([p])
            else:
                p_groups = split_list(p_list, self._groups)

            p_groups = [tuple(p) for p in p_groups]
            for group in p_groups:
                for p in group:
                    self._p_to_group[p] = group
                self._group_counts[group] = 0

        # 註冊hooks
        for param_group in self.param_groups: # 遍歷組
            for p in param_group['params']: # 遍歷組中的引數
                if p.requires_grad: # 如果需要計算梯度
                    p.grad = p.data.new(p.size()).zero_()
                    self._requires_update.add(p)
                    p_tmp = p.expand_as(p)
                    grad_acc = p_tmp.grad_fn.next_functions[0][0] # 獲取梯度函式
                    grad_acc.register_hook(self._make_hook(p)) # 註冊hook到梯度函式之上
                    self._grad_accs.append(grad_acc)

_make_hook 會構建 hooks，返回了 hook 函式，該函式會在反向傳播時候被呼叫，其內部執行了all-reduce。

def _make_hook(self, p):
    def hook(*ignore):
        # 省略部分程式碼
        handle, ctx = None, None
        self._allreduce_delay[p] -= 1
        if self._allreduce_delay[p] == 0:
            if self._groups is not None: # 處理 groups 相關部分，我們略過
                group = self._p_to_group[p]
                self._group_counts[group] += 1
                if self._group_counts[group] == len(group):
                    handle, ctxs = self._grouped_allreduce_grad_async(group) # 被呼叫時候會進行all-reduce
                    self._handles[group] = (handle, ctxs)
                    # Remove any None entries from previous no-op hook calls
                    for gp in group:
                        self._handles.pop(gp, None)
                    self._group_counts[group] = 0
                    return
            else:
                handle, ctx = self._allreduce_grad_async(p) # 被呼叫時候會進行all-reduce
        self._handles[p] = (handle, ctx) # 把handle註冊到本地，後續會使用
        
    return hook

4.1.2 歸併梯度

第二個階段是歸併，就是在反向傳播階段呼叫了 hook 函式，進行 all-reduce。

def _allreduce_grad_async(self, p):
    name = self._parameter_names.get(p)
    tensor = p.grad
    tensor_compressed, ctx = self._compression.compress(tensor)

    if self.op == Average:
       # Split average operation across pre/postscale factors
       # C++ backend will apply additional 1 / size() factor to postscale_factor for op == Average.
        prescale_factor = 1.0 / self.gradient_predivide_factor
        postscale_factor = self.gradient_predivide_factor
    else:
        prescale_factor = 1.0
        postscale_factor = 1.0

    # 呼叫 allreduce_async_ 完成 MPI 呼叫    
    handle = allreduce_async_(tensor_compressed, name=name, op=self.op,
                              prescale_factor=prescale_factor,
                              postscale_factor=postscale_factor)
    return handle, ctx

def _grouped_allreduce_grad_async(self, ps):
    name = self._parameter_names.get(ps[0])
    tensors_compressed, ctxs = zip(*[self._compression.compress(p.grad) for p in ps])

    handle = grouped_allreduce_async_(tensors_compressed, name=name, op=self.op)
    return handle, ctxs

4.1.2.1 MPI 函式

具體 MPI 函式位於 horovod/torch/mpi_ops.py

這裡要點是：allreduce_async_ 返回了一個 handle，後續可以控制，比如 poll 或者 synchronize。

def allreduce_async_(tensor, average=None, name=None, op=None,
                     prescale_factor=1.0, postscale_factor=1.0):
    """
    A function that performs asynchronous in-place averaging or summation of the input
    tensor over all the Horovod processes.

    The reduction operation is keyed by the name. If name is not provided, an incremented
    auto-generated name is used. 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.

    Arguments:
        tensor: A tensor to reduce.
        average:
            .. warning:: .. deprecated:: 0.19.0

                Use `op` instead. Will be removed in v0.21.0.

        name: A name of the reduction operation.
        op: The reduction operation to combine tensors across different ranks. Defaults to
            Average if None is given.
        prescale_factor: Multiplicative factor to scale tensor before allreduce.
        postscale_factor: Multiplicative factor to scale tensor after allreduce.

    Returns:
        A handle to the allreduce operation that can be used with `poll()` or
        `synchronize()`.
    """
    op = handle_average_backwards_compatibility(op, average)
    return _allreduce_async(tensor, tensor, name, op, prescale_factor, postscale_factor)

_allreduce_async 位於 horovod/torch/mpi_ops.py，其從 MPI 庫之中提取函式進行處理。

def _allreduce_async(tensor, output, name, op, prescale_factor, postscale_factor):
    # Set the divisor for reduced gradients to average when necessary
    if op == Average:
        if rocm_built():
            # For ROCm, perform averaging at framework level
            divisor = size()
            op = Sum
        else:
            divisor = 1

    elif op == Adasum:
        if tensor.device.type != 'cpu' and gpu_available('torch'):
            if nccl_built():
                if rocm_built():
                    # For ROCm, perform averaging at framework level
                    divisor = local_size()
                else:
                    divisor = 1
            else:
                divisor = 1
        else:
            divisor = 1
    else:
        divisor = 1

    function = _check_function(_allreduce_function_factory, tensor)
    try:
        handle = getattr(mpi_lib, function)(tensor, output, divisor,
                                            name.encode() if name is not None else _NULL, op,
                                            prescale_factor, postscale_factor)
    except RuntimeError as e:
        raise HorovodInternalError(e)
    _handle_map[handle] = (tensor, output)
    return handle

4.1.2.2 原理圖

這個圖和DDP類似，因此略過。

4.2 step 同步梯度

step 是另外一個進行all-reduce 操作的途徑。

step函式定義如下，可以看到，如果需要強制同步，就呼叫self.synchronize()，否則就呼叫基類的 step 函式來更新引數。

    def step(self, closure=None):
        if self._should_synchronize:
            if self._synchronized:
                warnings.warn("optimizer.step() called without "
                              "optimizer.skip_synchronize() context after "
                              "optimizer.synchronize(). This can cause training "
                              "slowdown. You may want to consider using "
                              "optimizer.skip_synchronize() context if you use "
                              "optimizer.synchronize() in your code.")
            self.synchronize()
        self._synchronized = False
        return super(self.__class__, self).step(closure)

4.2.1 synchronize

上面提到了 synchronize，我們下面就仔細研究一下。

從註釋中可以瞭解，synchronize() 是用來強制allreduce 操作完成，這對於梯度裁剪（gradient
clipping）或者其他有 in place 梯度修改的操作特別有用，這些操作需要在step()之前完成。

synchronize() 需要和 optimizer.skip_synchronize()一起合作。

DistributedOptimizer exposes the ``synchronize()`` method, which forces allreduce operations
to finish before continuing the execution. It's useful in conjunction with gradient
clipping, or other operations that modify gradients in place before ``step()`` is executed.
Make sure to use ``optimizer.skip_synchronize()`` if you're calling ``synchronize()``
in your code.

4.2.2 梯度裁剪

首先要了解什麼是梯度爆炸，梯度爆炸指的是在模型訓練過程之中，因為梯度變得太大而使得模型不穩定，容易直接跳過最優解。梯度裁剪（gradient clipping）就是一種解決梯度爆炸的技術：如果梯度變得太大，那麼就調節它使其保持較小的狀態，這樣可以避免模型越過最優點。

為了和梯度裁剪協同，需要在 step 之前呼叫 synchronize 以強制 all-reduce 完成。原始碼中的例子如下：

    output = model(data)
    loss = F.nll_loss(output, target)
    loss.backward()
    optimizer.synchronize()
    torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip)
    with optimizer.skip_synchronize():
        optimizer.step()

4.2.3 實現

我們接下來看看 synchronize 的實現。這裡最重要的是 outputs = synchronize(handle) 呼叫了 horovod.torch.mpi_ops.synchronize 完成了同步操作，這地方很容易讓新手誤解，因為名字相同，容易誤會成遞迴。

from horovod.torch.mpi_ops import synchronize

def synchronize(self):
    completed = set()
    for x in self._handles.keys():
      completed.update(x) if isinstance(x, tuple) else completed.add(x)
    missing_p = self._requires_update - completed # 找到目前沒有計算完畢的梯度
    
    for p in missing_p:
        handle, ctx = self._allreduce_grad_async(p) # 對於沒有計算完畢的，顯式進行all-reduce
        self._handles[p] = (handle, ctx) # 記錄下來本次計算的handle

    for p, (handle, ctx) in self._handles.items():
        if handle is None: # 如果沒有記錄呼叫過all-reduce
            handle, ctx = self._allreduce_grad_async(p)  # 進行all-reduce
            self._handles[p] = (handle, ctx)
            
    for p, (handle, ctx) in self._handles.items(): # 最後統一進行同步！
        if isinstance(p, tuple):
            # This was a grouped result, need to unpack
            outputs = synchronize(handle) # 呼叫 mpi 同步操作
            for gp, output, gctx in zip(p, outputs, ctx):
                self._allreduce_delay[gp] = self.backward_passes_per_step
                gp.grad.set_(self._compression.decompress(output, gctx))
        else:
            output = synchronize(handle) # 呼叫 mpi 同步操作
            self._allreduce_delay[p] = self.backward_passes_per_step
            p.grad.set_(self._compression.decompress(output, ctx))
            
    self._handles.clear()

    self._synchronized = True

4.2.4 MPI 同步操作

程式碼位於 horovod/torch/mpi_ops.py，直接呼叫了MPI 庫函式，有興趣同學可以自己深入研究。

def synchronize(handle):
    """
    Synchronizes an asynchronous allreduce, allgather or broadcast operation until
    it's completed. Returns the result of the operation.

    Arguments:
        handle: A handle returned by an allreduce, allgather or broadcast asynchronous
                operation.

    Returns:
        An output tensor of the operation.
    """
    if handle not in _handle_map:
        return

    try:
        mpi_lib.horovod_torch_wait_and_clear(handle)
        output = _handle_map.pop(handle)[-1]
        return output
    except RuntimeError as e:
        raise HorovodInternalError(e)

4.2.5 圖示

目前邏輯如下圖所示：

+---------------------------------------------------------------------------------+
| Process 1 on GPU 1                                                              |
|                                                 +----------------------------+  |
|                                                 | Optimizer                  |  |
|                                                 |                            |  |
| Forward +---->  Loss +----->  Backward  +---->  |     ALL-REDUCE +----> step |  |
|                                                 |                            |  |
|                                                 |            ^               |  |
|                                                 |            |               |  |
|                                                 +----------------------------+  |
|                                                              |                  |
+---------------------------------------------------------------------------------+
                                                               |
                                                               |
                                                               |
                                                               |
                                                               |
                                                           SYNC|GRADS
                                                               |
                                                               |
                                                               |
                                                               |
+----------------------------------------------------------------------------------+
| Process 2 on GPU 2                                           |                   |
|                                                              |                   |
|                                                 +-----------------------------+  |
|                                                 | Optimizer  |                |  |
|                                                 |            |                |  |
| Forward +---->  Loss +----->   Backward  +----> |            v                |  |
|                                                 |     ALL-REDUCE +----> step  |  |
|                                                 |                             |  |
|                                                 +-----------------------------+  |
|                                                                                  |
+----------------------------------------------------------------------------------+

至此，資料並行優化器分析完畢，下一篇我們介紹PyTorch 分散式優化器，敬請期待。