[原始碼解析] 深度學習分散式訓練框架 horovod (2) --- 從使用者角度切入

羅西的思考發表於2021-06-10

[原始碼解析] 深度學習分散式訓練框架 horovod (2) --- 從使用者角度切入

0x00 摘要

Horovod 是Uber於2017年釋出的一個易於使用的高效能的分散式訓練框架,在業界得到了廣泛應用。

本系列將通過原始碼分析來帶領大家瞭解 Horovod。系列大約有15 ~ 18 篇,本文是系列第二篇,從使用者角度切入 Horovod。

前一篇參見如下:

[原始碼解析] 深度學習分散式訓練框架 Horovod — (1) 基礎知識

0x01 Horovod 簡介

Horovod 是Uber於2017年釋出的一個易於使用的高效能的分散式訓練框架,支援TensorFlow,Keras,PyTorch和MXNet。Horovod 的名字來自於俄國傳統民間舞蹈,舞者手牽手圍成一個圈跳舞,與分散式 TensorFlow 流程使用 Horovod 互相通訊的場景很像。

因為各個機器學習框架對於底層集合通訊庫( nccl,openmpi,gloo 等等)的利用水平可能各不相同,使得他們無法充分利用這些底層集合通訊庫的威力。因而,hovorod 就整合這些框架,提供一個易用高效的解決方案。

Uber的工程師就是根據FaceBook的一篇paper:“Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour”和百度的一篇“Bringing HPC Techniques to Deep Learning” 改進併發布了開源框架Horovod。

Horovod 相比於百度的工作,並無學術上的貢獻。但是 Horovod 紮實的工程實現,使得它受到了更多的關注。它最大的優勢在於對 RingAllReduce 進行了更高層次的抽象,使其支援多種不同的框架。同時引入了 Nvidia NCCL,對 GPU 更加友好。

Horovod依賴於Nvidia的 NCCL2 做 All Reduce,依賴於MPI做程式間通訊,簡化了同步多 GPU 或多節點分散式訓練的開發流程。由於使用了NCCL2,Horovod也可以利用以下功能:NVLINK,RDMA,GPUDirectRDMA,自動檢測通訊拓撲,能夠回退到 PCIe 和 TCP/IP 通訊。

我們需要幾個問題來引導分析:

  • Hovorod 怎麼進行資料分割?
  • Hovorod 怎麼進行訓練程式碼分發?
  • Hovorod 啟動時候,python 和 C++ 都做了什麼?
  • 如何確保 Hovorod 啟動時候步驟一致;

0x02 Hovorod 機制概述

2.1 Horovod 機制

Horovod使用資料並行化策略在GPU上分配訓練。

在資料並行化中,作業中的每個GPU都會接收其自己的資料批處理的獨立切片,即它的“批處理切片”。 每個GPU都使用自己分配到的資料來獨立計算,進行梯度更新。

假如使用兩個GPU,批處理大小為32,則第一個GPU將處理前16條記錄的正向傳播和向後傳播,以及第二個GPU處理後16條記錄的正向傳播和向後傳播。然後,這些梯度更新將在GPU之間平均在一起,最後應用於模型。

每一個迭代的操作方法如下:

  1. 每個 worker 將維護自己的模型權重副本和自己的資料集副本。

  2. 收到執行訊號後,每個工作程式都會從資料集中提取一個不相交的批次,並計算該批次的梯度。

  3. Workers 使用ring all-reduce演算法來同步彼此的梯度,從而在本地所有節點上計算同樣的平均梯度。

    1. 將每個裝置上的梯度 tensor 切分成長度大致相等的 num_devices 個分片,後續每一次通訊都將給下一個鄰居傳送一個自己的分片(同時從上一個鄰居接受一個新分片)。

    2. ScatterReduce 階段:通過 num_devices - 1 輪通訊和相加,在每個 device 上都計算出一個 tensor 分片的和,即每個 device 將有一個塊,其中包含所有device 中該塊中所有值的總和;具體如下:

    3. AllGather 階段:通過 num_devices - 1 輪通訊和覆蓋,將上個階段計算出的每個 tensor 分片的和 廣播到其他 device;最終所有節點都擁有所有tensor分片和。具體如下:

    4. 在每個裝置上合併分片,得到梯度和,然後除以 num_devices,得到平均梯度;

  4. 每個 worker 將 梯度更新 應用於其模型的本地副本。

  5. 執行下一個batch。

0x03 示例程式碼

3.1 摘要程式碼

我們此處給出官網示例程式碼部分摘要,具體分析參見下面程式碼中的註釋。

import tensorflow as tf
import horovod.tensorflow.keras as hvd

# Horovod: initialize Horovod.
hvd.init() # 初始化 Horovod,啟動相關執行緒和MPI執行緒

# Horovod: pin GPU to be used to process local rank (one GPU per process)
# 依據 local rank 為不同的程式分配不同的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.Dense(10, activation='softmax')
])

# Horovod: adjust learning rate based on number of GPUs.
scaled_lr = 0.001 * hvd.size() # 根據Worker的數量增加學習率的大小
opt = tf.optimizers.Adam(scaled_lr)

# Horovod: add Horovod DistributedOptimizer.
# 把常規TensorFlow Optimizer通過Horovod包裝起來,進而使用 ring-allreduce 來得到平均梯度
opt = hvd.DistributedOptimizer(
    opt, backward_passes_per_step=1, average_aggregated_gradients=True)

# Horovod: Specify `experimental_run_tf_function=False` to ensure TensorFlow
# uses hvd.DistributedOptimizer() to compute gradients.
mnist_model.compile(loss=tf.losses.SparseCategoricalCrossentropy(),
                    optimizer=opt, metrics=['accuracy'],
                    experimental_run_tf_function=False)

callbacks = [
    hvd.callbacks.BroadcastGlobalVariablesCallback(0), # 廣播初始化,將模型的引數從第一個裝置傳向其他裝置,以保證初始化模型引數的一致性
    hvd.callbacks.MetricAverageCallback(),
    hvd.callbacks.LearningRateWarmupCallback(initial_lr=scaled_lr, warmup_epochs=3, verbose=1),
]

# Horovod: save checkpoints only on worker 0 to prevent other workers from corrupting them. # 只有裝置0需要儲存模型引數作為checkpoint
if hvd.rank() == 0:
    callbacks.append(tf.keras.callbacks.ModelCheckpoint('./checkpoint-{epoch}.h5'))

# Horovod: write logs on worker 0.
verbose = 1 if hvd.rank() == 0 else 0

# Train the model.
# Horovod: adjust number of steps based on number of GPUs.
mnist_model.fit(dataset, steps_per_epoch=500 // hvd.size(), callbacks=callbacks, epochs=24, verbose=verbose)

3.2 horovodrun

Horovod訓練指令碼 作為Python指令碼啟動。 例如,您不能使用python train.py執行此培訓指令碼。 需要採用特殊的CLI命令 horovodrun 來啟動(訓練程式碼 train.py 需要手動拷貝到各個節點上,且目錄相同):

$ horovodrun -np 4 -H localhost:4 python train.py

0x04 執行邏輯

我們按照順序梳理,看看在程式初始化過程背後都做了什麼。

4.1 引入python檔案

如下程式碼會引入各種相關python檔案。

import tensorflow as tf
import horovod.tensorflow.keras as hvd

4.2 初始化 in python

python 世界的初始化位於 horovod-master/horovod/mxnet/mpi_ops.py

4.2.1 引入SO庫

4.2.1.1 SO庫

horovod/tensorflow/mpi_ops.py 之中會引入SO庫。

比如 dist-packages/horovod/tensorflow/mpi_lib.cpython-36m-x86_64-linux-gnu.so。

SO庫 就是 horovod 中 C++ 程式碼編譯出來的結果

def _load_library(name):
    """Loads a .so file containing the specified operators.
    """
    filename = resource_loader.get_path_to_datafile(name)
    library = load_library.load_op_library(filename)
    return library

# Check possible symbol not found error from tensorflow version mismatch
try:
    MPI_LIB = _load_library('mpi_lib' + get_ext_suffix())
except Exception as e:
    check_installed_version('tensorflow', tf.__version__, e)
    raise e
else:
    check_installed_version('tensorflow', tf.__version__)
4.2.2.2 SO作用

引入庫的作用是獲取到 C++ 的函式,並且用 python 封裝一下,這樣就可以在 python 世界使用 C++程式碼了。

由下文可以看出來,python 的 _allreduce 函式就會把功能轉發給 C++,由 MPI_LIB.horovod_allreduce 完成。

def _allreduce(tensor, name=None, op=Sum, prescale_factor=1.0, postscale_factor=1.0,
               ignore_name_scope=False):
    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)

4.2.2 初始化配置

我們摘錄了主要部分,就是初始化 _HorovodBasics,然後從 _HorovodBasics 內獲取各種函式,變數和配置,比如是否編譯了mpi,gloo等等。

from horovod.common.basics import HorovodBasics as _HorovodBasics

_basics = _HorovodBasics(__file__, 'mpi_lib')

# import basic methods
init = _basics.init
size = _basics.size
local_size = _basics.local_size
rank = _basics.rank
local_rank = _basics.local_rank
mpi_built = _basics.mpi_built
gloo_enabled = _basics.gloo_enabled
......

4.2.3 hvd.init() 初始化

首先需要用 hvd.init() 來初始化,horovod 管理的所有狀態都會傳到 hvd 物件中。

# Horovod: initialize Horovod.
hvd.init()

此處呼叫的是 HorovodBasics 中的函式,我們看看做了什麼。

可以看到,這部分會一直深入到 C++世界,呼叫了大量的 MPI_LIB_CTYPES 函式,所以我們接下來就要進入到 C++的世界看看。

def init(self, comm=None):
    """A function that initializes Horovod.
    """
    atexit.register(self.shutdown)

    if not isinstance(comm, list):
        mpi_built = self.MPI_LIB_CTYPES.horovod_mpi_built()

        from mpi4py import MPI
        if MPI._sizeof(MPI.Comm) == ctypes.sizeof(ctypes.c_int):
            MPI_Comm = ctypes.c_int
        else:
            MPI_Comm = ctypes.c_void_p
            self.MPI_LIB_CTYPES.horovod_init_comm.argtypes = [MPI_Comm]

        comm_obj = MPI_Comm.from_address(MPI._addressof(comm))
        self.MPI_LIB_CTYPES.horovod_init_comm(comm_obj)
    else:
        comm_size = len(comm)
        self.MPI_LIB_CTYPES.horovod_init(
            (ctypes.c_int * comm_size)(*comm), ctypes.c_int(comm_size))

目前邏輯如下圖:

           Import python files

                    +
                    |
                    |
                    v

           Import C++ SO files
                    |
                    |
                    |
                    v

           Create _HorovodBasics
                    +
                    |
                    |
                    v
                hvd.init()
                    +
Python              |
+------------------------------------------+
C++                 |
                    |
                    v

4.3 初始化 in C++

4.3.1 horovod_init_comm

在初始化的時候,Horovod 會:

  • 呼叫 MPI_Comm_dup 獲取一個 Communicator,這樣就有了和 MPI 協調的基礎。
  • 然後呼叫 InitializeHorovodOnce。
void horovod_init_comm(MPI_Comm comm) {
  MPI_Comm_dup(comm, &mpi_context.mpi_comm);
  InitializeHorovodOnce(nullptr, 0);
}

4.3.2 InitializeHorovodOnce

InitializeHorovodOnce 是初始化的主要工作,主要是:

  • 依據是否編譯了 mpi 或者 gloo,對各自的 context 進行處理,為 globalstate 建立對應的 controller;
  • 啟動了後臺執行緒 BackgroundThreadLoop 用來在各個worker之間協調;
void horovod_init(const int* ranks, int nranks) {
  InitializeHorovodOnce(ranks, nranks);
}

void InitializeHorovodOnce(const int* ranks, int nranks) {
  // Ensure background thread is only started once.
  if (!horovod_global.initialize_flag.test_and_set()) {
    horovod_global.control_operation = ParseControllerOpsFromEnv();
    horovod_global.cpu_operation = ParseCPUOpsFromEnv();
    
#if HAVE_MPI // 依據是否編譯了MPI進行處理
    // Enable mpi is it's used either in cpu data transfer or controller
    if (horovod_global.cpu_operation == LibType::MPI ||
        horovod_global.control_operation == LibType::MPI) {
      mpi_context.Enable();
    }

    if (horovod_global.control_operation == LibType::MPI){
      // 建立一個 MPIController 物件
      horovod_global.controller.reset(new MPIController(
          horovod_global.response_cache,
          horovod_global.tensor_queue, horovod_global.timeline,
          horovod_global.parameter_manager, horovod_global.group_table,
          mpi_context));
      horovod_global.controller->SetRanks(ranks, nranks);
    }
#endif

#if HAVE_GLOO // 依據是否編譯了 GLOO 進行處理
    // Enable gloo is it's used either in cpu data transfer or controller
    if (horovod_global.cpu_operation == LibType::GLOO ||
        horovod_global.control_operation == LibType::GLOO) {
      gloo_context.Enable();
    }

    if (horovod_global.control_operation == LibType::GLOO) {
      horovod_global.controller.reset(new GlooController(
          horovod_global.response_cache,
          horovod_global.tensor_queue, horovod_global.timeline,
          horovod_global.parameter_manager, horovod_global.group_table,
          gloo_context));
    }
#endif
    // Reset initialization flag
    // 啟動後臺執行緒
    horovod_global.initialization_done = false;
    horovod_global.background_thread = std::thread(
        BackgroundThreadLoop, std::ref(horovod_global));
  }

  // Wait to ensure that the background thread has finished initializing MPI.
  while (!horovod_global.initialization_done) {
    std::this_thread::sleep_for(std::chrono::milliseconds(1));
  }
}

4.3.3 HorovodGlobalState

在 C++ 世界,HorovodGlobalState 起到了集中管理各種全域性變數的作用。

HorovodGlobalState 在 horovod 中是一個全域性變數,其中的元素可以供不同的執行緒訪問。HorovodGlobalState 在載入 C++ 的程式碼時候就已經建立了,同時建立的還有各種 context(mpi_context, nccl_context, gpu_context)。

Horovod 主要會在backgroundThreadLoop 中完成 HorovodGlobalState 不同元素初始化,比較重要的有:

  • controller 管理總體通訊控制流;
  • tensor_queue 會處理從前端過來的通訊需求(allreduce,broadcast 等);
// All the Horovod state that must be stored globally per-process.
HorovodGlobalState horovod_global;

#if HAVE_MPI
MPIContext mpi_context;
#endif

#if HAVE_GLOO
GlooContext gloo_context;
#endif

....

std::unique_ptr<OperationManager> op_manager;

HorovodGlobalState 摘要如下:

struct HorovodGlobalState {

  // Background thread running MPI communication.
  std::thread background_thread; // 後臺執行緒,用來在各個worker之間協調

  ParameterManager parameter_manager; // 維護後臺總體引數配置

  // Encapsulates the fusion buffers, handles resizing and auto-tuning of buffer
  // size.
  FusionBufferManager fusion_buffer; // 融合tensor,以便縮減通訊開銷

  std::shared_ptr<Controller> controller; //管理總體通訊控制流

  TensorQueue tensor_queue; //處理從前端過來的通訊需求(allreduce,broadcast 等)

  // Pointer to shared buffer for allgather
  void* shared_buffer = nullptr;

  // LRU cache of Responses
  ResponseCache response_cache;

  // Information on registered groups.
  GroupTable group_table;

  ~HorovodGlobalState() {
    // Make sure that the destructor of the background thread is safe to
    // call. If a thread is still joinable (not detached or complete) its
    // destructor cannot be called.
    if (background_thread.joinable()) {
      shut_down = true;
      background_thread.join();
    }
  }
};

目前具體邏輯如下:

           Import python files

                    +
                    |
                    |
                    v

           Import C++ SO files
                    |
                    |
                    |
                    v

           Create _HorovodBasics
                    +
                    |
                    |
                    v
                hvd.init()
                    +
Python              |
+-------------------------------------------------------------------------------------------------------------+
                    |
c++                 |
                    v                                                          +-----------------------------+
                                                                               |  HorovodGlobalState         |
              horovod_init_comm                                                |                             |
                    +                             +------------------+         |                             |
                    |                             | horovod_global +---------> |        TensorQueue          |
                    |                             |                  |         |                             |
                    v                             |                  |         |        background_thread    |
                                                  | mpi_context      |         |                             |
           InitializeHorovodOnce   +------------> |                  |         |        ParameterManager     |
                    +                             |                  |         |                             |
                    |                             | gloo_context     |         |        FusionBufferManager  |
                    |                             |                  |         |                             |
                    |                             |                  |         |        Controller           |
                    v                             | op_manager       |         |                             |
             background_threa                     |                  |         |        ResponseCache        |
                                                  +------------------+         |                             |
                                                                               |        shared_buffer        |
                                                                               +-----------------------------+

手機如下:

至此,horovod 已經初始化完成,使用者程式碼可以使用了。

4.4 hvd 概念

在使用者程式碼中,接下來是rank概念。

hvd.local_rank()

hvd.rank()

我們介紹下幾個相關概念:

  • Horovod為裝置上的每個GPU啟動了該訓練指令碼的一個副本。**local rank **就是分配給 某一臺計算機 上每個執行訓練 的唯一編號(也可以認為是程式號或者GPU裝置的ID號),範圍是 0 到 n-1,其中 n 是該計算機上GPU裝置的數量。
  • rank 可以認為是代表分散式任務裡的一個執行訓練的唯一全域性編號(用於程式間通訊)。Rank 0 在Horovod中通常具有特殊的意義:它是負責此同步的裝置。
    • 在百度的實現中,不同 Rank 的角色是不一樣的,Rank 0 會充當 coordinator 的角色。它會協調來自其他 Rank 的 MPI 請求,是一個工程上的考量。這一設計也被後來的 Horovod 採用。
    • Rank 0 也用來把引數廣播到其他程式 & 儲存 checkpoint。
  • world_size:程式總數量,會等到所有world_size個程式就緒之後才會開始訓練。

hvd.init 這部分的目的就是讓並行程式們可以知道自己被分配的 rank / local rank 等資訊,於是後續可以根據 local rank(所在節點上的第幾張 GPU 卡) 來設定所需的視訊記憶體分配

4.5 資料處理

接下來是資料處理。

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)

這裡有幾點需要說明:

  • 首先,訓練的資料需要放置在任何節點都能訪問的地方。

  • 其次,Horovod 需要對資料進行分片處理,需要在不同機器上按Rank進行切分,以保證每個GPU程式訓練的資料集是不一樣的。

  • 資料集本體需要出於資料並行性的需求而被拆分為多個分片,Horovod的不同工作節點都將分別讀取自己的資料集分片。

從 PyTorch 示例指令碼看得更加清楚。

# Horovod: use DistributedSampler to partition the training data.
train_sampler = torch.utils.data.distributed.DistributedSampler(
    train_dataset, num_replicas=hvd.size(), rank=hvd.rank())
train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=args.batch_size, sampler=train_sampler, **kwargs)
  • DataLoader的取樣器元件從要繪製的資料集中返回可迭代的索引。 PyTorch中的預設取樣器是順序的,返回序列0, 1, 2, …, n 。 Horovod使用其DistributedSampler覆蓋了此行為,該DistributedSampler處理跨計算機的資料集分割槽。 DistributedSampler本身接受兩個引數作為輸入: hvd.size() (GPU的總數,例如16)和hvd.rank() (從總體列表中分配給該裝置的ID,例如0…15)。

  • Pytorch使用的是資料分散式訓練,每個程式實際上是獨立載入資料的,所以需要載入相同資料集後用一定的規則根據rank來順序切割獲取不同的資料子集,DistributedSampler就是用來確保dataloader只會load到整個資料集的一個特定子集的做法(實際上不用Pytorch提供的DistributedSampler工具,自己做載入資料後切分word_size個子集按rank順序拿到子集效果也是一樣)。

  • 同時為了能夠按順序劃分資料子集,拿到不同部分資料,所以資料集不能夠進行隨機打散,所以用了引數 'shuffle': False。

4.6 廣播初始化變數

以下程式碼完成廣播初始化的功能。

hvd.callbacks.BroadcastGlobalVariablesCallback(0)

這句程式碼保證的是 rank 0 上的所有引數只在 rank 0 初始化,然後廣播給其他節點,即變數從第一個流程向其他流程傳播,以實現引數一致性初始化

下面就介紹下 Horvod 之中廣播的使用。

4.6.1 廣播定義

廣播的具體實現是:

class BroadcastGlobalVariablesCallbackImpl(object):
    def __init__(self, backend, root_rank, device='', *args):
        super(BroadcastGlobalVariablesCallbackImpl, self).__init__(*args)
        self.backend = backend
        self.root_rank = root_rank
        self.device = device
        self.broadcast_done = False

    def on_batch_end(self, batch, logs=None):
        if self.broadcast_done:
            return

        with tf.device(self.device):
            if hvd._executing_eagerly() and hasattr(self.model, 'variables'):
                # TensorFlow 2.0 or TensorFlow eager
                hvd.broadcast_variables(self.model.variables,
                                        root_rank=self.root_rank)
                hvd.broadcast_variables(self.model.optimizer.variables(),
                                        root_rank=self.root_rank)
            else:
                bcast_op = hvd.broadcast_global_variables(self.root_rank)
                self.backend.get_session().run(bcast_op)

        self.broadcast_done = True

4.6.2 broadcast_variables

broadcast_variables 呼叫了 _make_broadcast_group_fn 完成功能,可以看到對於 執行圖 的每個變數,呼叫了 broadcast。

def broadcast_variables(variables, root_rank):
    """Broadcasts variables from root rank to all other processes.

    Arguments:
        variables: variables for broadcast
        root_rank: rank of the process from which global variables will be broadcasted
                   to all other processes.
    """
    broadcast_group = _make_broadcast_group_fn()
    return broadcast_group(variables, root_rank)

以及

@_cache
def _make_broadcast_group_fn():
    if _executing_eagerly():
        # Eager mode will parallelize independent control flow
        def broadcast_group(variables, root_rank):
            for var in variables:
                var.assign(broadcast(var, root_rank))

        return _make_subgraph(broadcast_group)
    else:
        # Graph mode requires an Op
        def broadcast_group(variables, root_rank):
            return tf.group(*[var.assign(broadcast(var, root_rank))
                              for var in variables])

        return broadcast_group

4.6.3 呼叫 MPI

broadcast 就是呼叫了 MPI 函式真正完成了功能。

def broadcast(tensor, root_rank, name=None, ignore_name_scope=False):
    """An op which broadcasts the input tensor on root rank to the same input tensor
    on all other Horovod processes.

    The broadcast 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 broadcast
    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`, with the value broadcasted
      from root rank.
    """
    if name is None and not _executing_eagerly():
        name = 'HorovodBroadcast_%s' % _normalize_name(tensor.name)
    return MPI_LIB.horovod_broadcast(tensor, name=name, root_rank=root_rank,
                                     ignore_name_scope=ignore_name_scope)

4.6.4 同步引數

在後臺程式中,會根據情況定期同步引數。

bool RunLoopOnce(HorovodGlobalState& state) {
	// 業務邏輯功能

  if (state.parameter_manager.IsAutoTuning()) {
    bool should_sync =
        state.parameter_manager.Update(tensor_names, total_tensor_size);
    // 看看是否需要同步,如果需要,就同步。
    if (should_sync) {
      state.controller->SynchronizeParameters();
    }
  }
  ......
}

同步引數程式碼也是呼叫了 Bcast 功能完成。

void Controller::SynchronizeParameters() {
  ParameterManager::Params param;
  if (is_coordinator_) { // rank 0 執行操作
    param = parameter_manager_.GetParams();
  }

  void* buffer = (void*)(&param);
  size_t param_size = sizeof(param);
  Bcast(buffer, param_size, 0, Communicator::GLOBAL);

  if (!is_coordinator_) { // worker 執行操作
    parameter_manager_.SetParams(param);
  }
}

4.7 DistributedOptimizer

最後需要配置DistributedOptimizer,這就是關鍵點之一。

# Horovod: add Horovod DistributedOptimizer.
opt = hvd.DistributedOptimizer(
    opt, backward_passes_per_step=1, average_aggregated_gradients=True)

TF Optimizer 是模型訓練的關鍵API,可以獲取到每個OP的梯度並用來更新權重。HVD 在原始 TF Optimizer的基礎上包裝了hvd.DistributedOptimizer

DistributedOptimizer包裝器將原始優化器作為輸入,將梯度計算委託給它。 即DistributedOptimizer會呼叫原始優化器進行梯度計算。這樣,在叢集中每臺機器都會用原始優化器得到自己的梯度(Local Gradient)。

Horovod DistributedOptimizer接下來會使用all-reduce或all-gather來完成全域性梯度歸併,然後將這些平均梯度應用於所有裝置。

我們梳理下其中的呼叫關係:

  • hvd.DistributedOptimizer繼承 keras Optimizer,在計算時候,依然由傳入的原始優化器做計算。
  • 在得到計算的梯度之後,呼叫 hvd.allreduce 或者 hvd.allgather 來計算。
  • 最後實施這些平均之後的梯度。從而實現整個叢集的梯度歸併操作。

具體後文會詳細介紹。

4.8 未來可能

Horovod 目前架構的基礎是:機器學習的模型引數在一張 GPU 上可以存下。

未來是否可以把模型分片結合進來,是一個很大的看點。

另外,如果模型的全連線層較多,則全連線層的強耦合性結合 allreduce 類似 bsp 的同步機制,還是會讓網路通訊時間成為瓶頸。因此,在 ring-allreduce 環境下,同步協議的改造,比如利用 SSP 來替換 BSP,或者利用梯度壓縮來加快 allreduce 程式也是值得探索的方向。

0x05 總結

針對文初提出的幾個問題,我們現在回答如下:

  • Hovorod 怎麼進行資料分割?
    • 答案:有的框架可以自動做資料分割。如果框架不提供,則需要使用者自己進行資料分割,以保證每個GPU程式訓練的資料集是不一樣的。
  • Hovorod 怎麼進行模型分發?
    • 使用者需要手動拷貝訓練程式碼到各個節點上。
  • Hovorod 啟動時候,python 和 C++ 都做了什麼?
    • 答案:python 會引入 C++庫,初始化各種變數和配置。C++部分會對 MPI,GLOO上下文進行初始化,啟動後臺程式處理內部通訊。
  • 如何確保 Hovorod 啟動時候步驟一致;
    • 答案: rank 0 上的所有引數只在 rank 0 初始化,然後廣播給其他節點,即變數從第一個流程向其他流程傳播,以實現引數一致性初始化。

下一篇文章將深入到python世界看看。

0xEE 個人資訊

★★★★★★關於生活和技術的思考★★★★★★

微信公眾賬號:羅西的思考

如果您想及時得到個人撰寫文章的訊息推送,或者想看看個人推薦的技術資料,敬請關注。

在這裡插入圖片描述

0xFF 參考

瞭解Pytorch 分散式訓練,這一篇足夠了!

horovod使用_用horovod進行分散式模型訓練

Spark新願景:讓深度學習變得更加易於使用

Scaling model training in PyTorch using distributed data parallel

使用分散式資料並行在PyTorch中進行縮放模型訓練

A developer-friendly guide to mixed precision training with PyTorch

開發人員友好的PyTorch混合精度培訓指南

It’s 2020, why isn’t deep learning 100% on the cloud yet?

到了2020年,為什麼還不可以在雲上進行100%的深度學習?

帶你瞭解當紅炸子雞Horovod分散式訓練框架

在 Amazon SageMaker 管道模式下使用 Horovod 實現多 GPU 分散式訓練

kubernetes 培訓_在Kubernetes上使用horovod進行分散式深度學習培訓

Horovod-基於TensorFlow分散式深度學習框架

Paracel十問

PARACEL:讓分散式機器學習變得簡單

Spark on Angel大規模分散式機器學習平臺介紹

分散式TensorFlow入門教程

引數伺服器——分散式機器學習的新殺器

NCCL--GPU的collective communication通訊技術

飛槳異構引數伺服器架構

談分散式機器學習系統中的網路相關問題

騰訊大規模分散式機器學習系統無量是如何進行技術選型的

如何理解Nvidia英偉達的Multi-GPU多卡通訊框架NCCL?

百度將高效能運算引入深度學習:可高效實現模型的大規模擴充套件

tensorflow分散式原始碼解讀4:AdamOptimizer

機器學習中的平行計算

分散式機器學習(上)-平行計算與機器學習

分散式機器學習(中)-平行計算與機器學習

分散式機器學習(下)-聯邦學習

[Distributed ML] Parameter Server & Ring All-Reduce

深度學習的分佈和並行處理系統

並行和分散式深度學習

分散式機器學習、聯邦學習(Shusen Wang)

相關文章