作者 | 金色旭光

一、背景介紹

我是一名Python開發(fā)，就職于一家AI公司，負(fù)責(zé)開發(fā)迭代一個(gè)深度學(xué)習(xí)的模型訓(xùn)練平臺(tái)。模型訓(xùn)練平臺(tái)主要是給算法工程師訓(xùn)練模型，開發(fā)語言是Python，Web框架為Fastapi。模型訓(xùn)練使用Pytorch框架，封裝成Docker運(yùn)行。我負(fù)責(zé)除Pytorch之外平臺(tái)功能開發(fā)，有一位算法工程師負(fù)責(zé)Pytorch開發(fā)，封裝成容器提供給我。

目前這個(gè)訓(xùn)練平臺(tái)是單機(jī)版，支持多顯卡訓(xùn)練，也就是所謂的單機(jī)多卡的訓(xùn)練模式。隨著公司業(yè)務(wù)的發(fā)展，模型訓(xùn)練需要的GPU越來越多。單臺(tái)服務(wù)器支持顯卡數(shù)量再多也會(huì)有一個(gè)上限，這時(shí)就需要能夠使用多臺(tái)GPU服務(wù)器上的多個(gè)顯卡，也就是多機(jī)多卡的訓(xùn)練模式。

在這樣的背景下，我需要將單機(jī)的訓(xùn)練平臺(tái)升級(jí)為分布式的訓(xùn)練平臺(tái)。只有我一桿槍，一個(gè)配合的算法工程師，一個(gè)前端，一個(gè)測(cè)試。經(jīng)過將近兩個(gè)月的開發(fā)，完成了這個(gè)任務(wù)。

開發(fā)過程遇到的非常多的問題，折磨了我一次又一次。好在最后基本都解決問題了。本篇就從需求說明、實(shí)現(xiàn)方案、踩坑經(jīng)歷等方面來介紹這一段特殊而難得的開發(fā)經(jīng)歷。

二、需求說明

1.單機(jī)版架構(gòu)

首先介紹一下單機(jī)版模型訓(xùn)練平臺(tái)。模型訓(xùn)練簡(jiǎn)單來說就是用訓(xùn)練容器讀取數(shù)據(jù)集，跑模型訓(xùn)練，最終生成一個(gè)模型文件。

圖片

模型訓(xùn)練主要有兩個(gè)步驟：

• 準(zhǔn)備數(shù)據(jù)集
• 啟動(dòng)訓(xùn)練容器跑訓(xùn)練

單機(jī)版顧名思義數(shù)據(jù)集、訓(xùn)練容器、生成的模型等所有流程都在一臺(tái)服務(wù)器上完成。

2.分布式版架構(gòu)

分布式的訓(xùn)練平臺(tái)是將訓(xùn)練任務(wù)分發(fā)到多個(gè)訓(xùn)練節(jié)點(diǎn)上，讓多臺(tái)服務(wù)器的GPU互相通信，算力統(tǒng)一起來使用。

圖片

想要將這樣一個(gè)單機(jī)架構(gòu)的平臺(tái)升級(jí)成分布式平臺(tái)需要實(shí)現(xiàn)的功能有三個(gè)：

• 每個(gè)訓(xùn)練節(jié)點(diǎn)都要讀取全量的數(shù)據(jù)集，需要將數(shù)據(jù)集復(fù)制到各個(gè)訓(xùn)練節(jié)點(diǎn)
• 多個(gè)訓(xùn)練節(jié)點(diǎn)要支持跨節(jié)點(diǎn)的模型訓(xùn)練
• 容器啟動(dòng)命令要能夠下發(fā)到訓(xùn)練節(jié)點(diǎn)

其中第2個(gè)功能Pytorch模型訓(xùn)練框架已經(jīng)支持了分布式的訓(xùn)練模式，并且當(dāng)前系統(tǒng)做分布式也是基于這個(gè)能力才有可能開發(fā)完成。

DistributedDataParallel（DDP）是一個(gè)支持多機(jī)多卡、分布式訓(xùn)練的深度學(xué)習(xí)工程方法。Pytorch現(xiàn)已原生支持DDP，可以直接通過torch.distributed使用。

讓Pytorch訓(xùn)練容器支持ddp是由算法工程師去完成的，對(duì)于我來說，只需要在訓(xùn)練節(jié)點(diǎn)1和2上執(zhí)行不同的容器啟動(dòng)命令即可。

三、實(shí)現(xiàn)方案

針對(duì)實(shí)現(xiàn)功能1、3，技術(shù)方案設(shè)計(jì)如下：

1.數(shù)據(jù)集復(fù)制

因?yàn)閿?shù)據(jù)集只能從主節(jié)點(diǎn)上傳到平臺(tái)，所以要想將數(shù)據(jù)集移動(dòng)到訓(xùn)練節(jié)點(diǎn)有兩個(gè)方案，分別是：NFS共享目錄、文件同步 。

NFS

NFS不用過多介紹了，就是本地掛載一塊遠(yuǎn)端機(jī)器的目錄，將遠(yuǎn)端目錄當(dāng)做本地目錄使用。

優(yōu)點(diǎn)：NFS 的優(yōu)點(diǎn)是內(nèi)核直接支持，部署簡(jiǎn)單、運(yùn)行穩(wěn)定，協(xié)議簡(jiǎn)單。

缺點(diǎn)：通過網(wǎng)絡(luò)讀取數(shù)據(jù)集，IO速度會(huì)成為數(shù)據(jù)加載的瓶頸。

NFS的缺點(diǎn)是網(wǎng)絡(luò)傳輸速度慢，我們的環(huán)境只有千兆帶寬，在模型訓(xùn)練時(shí)通過千兆帶寬分布式進(jìn)程通信會(huì)讓整體的訓(xùn)練速度慢一個(gè)等級(jí)。最優(yōu)解是IB網(wǎng)，IB網(wǎng)是轉(zhuǎn)為大規(guī)模數(shù)據(jù)中心設(shè)計(jì)的網(wǎng)絡(luò)架構(gòu)，帶寬能達(dá)到50G，但是我們沒有，客戶大概率也用不上成本飆升的IB網(wǎng)。

文件同步

通過文件同步可以將數(shù)據(jù)集分發(fā)到訓(xùn)練節(jié)點(diǎn)。比較了常規(guī)文件同步使用的技術(shù)，最后選擇了lsyncd這款工具。

rsync 是Linux系統(tǒng)上一款開源的快速的可實(shí)現(xiàn)全量及增量遠(yuǎn)程數(shù)據(jù)同步備份的優(yōu)秀工具。lysncd 是lua語言封裝了 inotify 和 rsync 工具，采用了 Linux 內(nèi)核里的 inotify 事件觸發(fā)機(jī)制，然后通過rsync同步差異，達(dá)到實(shí)時(shí)的效果。

優(yōu)點(diǎn)：支持?jǐn)帱c(diǎn)續(xù)傳；同步數(shù)據(jù)集，能夠滿足模型訓(xùn)練需要的IO速度

缺點(diǎn)：同一份文件會(huì)復(fù)制出多份，存在冗余，增加存儲(chǔ)的壓力；文件同步是基于時(shí)間間隔或累計(jì)文件數(shù)據(jù)量，非嚴(yán)格意義的實(shí)時(shí)。

方案選擇

因?yàn)镹FS的缺點(diǎn)比較致命，而lsyncd的缺點(diǎn)通過邏輯可以克服。所以數(shù)據(jù)集的最終解決辦法是使用lsyncd同步數(shù)據(jù)集，同時(shí)也需要將訓(xùn)練節(jié)點(diǎn)生成的模型等文件同步到主服務(wù)上，即雙向同步（這里有坑，下文會(huì)說）

2.遠(yuǎn)程執(zhí)行命令

單機(jī)版運(yùn)行時(shí)直接在本機(jī)通過命令啟動(dòng)訓(xùn)練容器，命令類似：

nvidia-docker run--gpus '"device=0"' -name train_container_1 7055fe2b9719

分布式訓(xùn)練需要選擇一臺(tái)服務(wù)器做主節(jié)點(diǎn)，多臺(tái)服務(wù)器做訓(xùn)練節(jié)點(diǎn)，需要在不同的服務(wù)器上啟動(dòng)多條命令。所以就需要一個(gè)能遠(yuǎn)程執(zhí)行命令的功能。

能夠遠(yuǎn)程執(zhí)行的技術(shù)選型還有：

• http 請(qǐng)求
• rpc 請(qǐng)求
• ssh 遠(yuǎn)程執(zhí)行
• socket 網(wǎng)絡(luò)
• 中間件發(fā)布訂閱

經(jīng)過對(duì)比，最終選擇rpc來實(shí)現(xiàn)這個(gè)功能，理由如下：

• rpc 可以同步響應(yīng)
• rpc 基于IP地址調(diào)用，符合當(dāng)前整體架構(gòu)
• rpc 服務(wù)端方便記錄調(diào)用的日志
• 沒有中間件，組件少，技術(shù)簡(jiǎn)單

Python 中rpc相關(guān)的庫有很多，如：

• python自帶的庫 xmlRPC
• google開源可跨語言的 grpc
• 第三方庫 zerorpc
• 第三方庫 jsonrpclib

經(jīng)過比較選擇zerorpc，原因是靈活、輕量級(jí)、高性能。zerorpc的demo如下：

服務(wù)端：

import zerorpc

class HelloRPC(object):
    def hello(self, name):
        return "Hello, %s" % name

s = zerorpc.Server(HelloRPC())
s.bind("tcp://0.0.0.0:4242")
s.run()

客戶端：

import zerorpc

c = zerorpc.Client()
c.connect("tcp://127.0.0.1:4242")
print(c.hello("RPC"))

四、開發(fā)過程記錄

編碼時(shí)間大概5個(gè)星期左右，時(shí)間是蠻久，只怪咱只有一個(gè)人。進(jìn)度流水賬如下：

1.部署lsyncd同步工具，讓數(shù)據(jù)集能夠從主節(jié)點(diǎn)同步到訓(xùn)練節(jié)點(diǎn)

2.開發(fā)rpc的服務(wù)端和客戶端，訓(xùn)練命令下發(fā)到選中的節(jié)點(diǎn)

3.通過rpc獲取訓(xùn)練節(jié)點(diǎn)GPU信息，讓頁面支持選擇不同機(jī)器的GPU

4.調(diào)試遠(yuǎn)程訓(xùn)練單機(jī)單卡，調(diào)通數(shù)據(jù)集分發(fā)和訓(xùn)練命令下發(fā)

5.和算法工程師確定多機(jī)多卡訓(xùn)練容器的啟動(dòng)命令

6.調(diào)試多機(jī)多卡訓(xùn)練，發(fā)現(xiàn)ddp啟動(dòng)會(huì)阻塞，解決問題花費(fèi)一個(gè)星期

7.發(fā)現(xiàn)rpc有問題，替換zerorpc為grpc

8.數(shù)據(jù)集同步和執(zhí)行訓(xùn)練命令之間有先后依賴關(guān)系，解決同步問題

9.模型訓(xùn)練、推理、驗(yàn)證三個(gè)主要功能完成

經(jīng)過5個(gè)星期的開發(fā)，最終完成了模型訓(xùn)練的基本功能，包括模型訓(xùn)練、模型驗(yàn)證、模型推理。由于架構(gòu)的變化應(yīng)該可能潛在一些未發(fā)現(xiàn)的bug。對(duì)于bug來說，發(fā)現(xiàn)它是測(cè)試同事的工作，而是我的任務(wù)就是送它去見測(cè)試同事。所以，就轉(zhuǎn)測(cè)了。

五、遇到的問題

經(jīng)過三輪的測(cè)試，在修復(fù)了很多bug之后，最終完成了分布式功能版本的開發(fā)。

在這兩個(gè)月中，我遇到了非常多問題，我登記在冊(cè)的問題是13個(gè)，實(shí)際上還有一些未上榜的，主要原因是從單節(jié)點(diǎn)到分布式涉及到存儲(chǔ)、通信等變化讓系統(tǒng)復(fù)雜。

限于篇幅挑選幾個(gè)講講，給后續(xù)使用相關(guān)技術(shù)的人一個(gè)避坑的提醒。

1.zerorpc 服務(wù)端不支持多線程并發(fā)請(qǐng)求

調(diào)研zerorpc時(shí)，我關(guān)注的點(diǎn)包括是否滿足功能要求、代碼復(fù)雜性、模塊的活躍度、github代碼提交時(shí)間等。從我關(guān)注點(diǎn)出發(fā)，zerorpc是比較完美符合我要求的，但是完成相關(guān)功能開發(fā)之后才發(fā)現(xiàn)zerorpc竟然不支持并發(fā)請(qǐng)求。真是廁所里跳高——過分。

現(xiàn)象：多個(gè)客戶端請(qǐng)求達(dá)到服務(wù)端，請(qǐng)求會(huì)變成串行執(zhí)行

原因：zerorpc是基于協(xié)程庫gevent實(shí)現(xiàn)的并發(fā)，而我們的技術(shù)棧不是協(xié)程，這就導(dǎo)致zerorpc不支持并發(fā)操作。

解決辦法：rpc服務(wù)端肯定需要支持并發(fā)請(qǐng)求，將zerorpc換成了grpc。

之所以開始沒有選擇grpc，是因?yàn)間rpc使用略復(fù)雜，需要先寫proto文件，編譯，再分別實(shí)現(xiàn)客戶端和服務(wù)端。但實(shí)事證明雖然繁瑣了一些，grpc還是值得信賴的。

2.數(shù)據(jù)集篩選報(bào)錯(cuò)

數(shù)據(jù)集是決定模型質(zhì)量的一個(gè)重要因素，所以對(duì)數(shù)據(jù)集會(huì)有合并、過濾、篩選等操作，每次操作都會(huì)生成一份新的數(shù)據(jù)集文件。

現(xiàn)象：測(cè)試發(fā)現(xiàn)篩選數(shù)據(jù)集時(shí)偶爾會(huì)報(bào)錯(cuò)，大概篩選10次以上就會(huì)出現(xiàn)。

原因：非必現(xiàn)的問題是最頭疼的問題。這個(gè)問題我排查了3天，最后發(fā)現(xiàn)是lsyncd雙向同步的問題，也就是我在技術(shù)選型中提到的坑。

數(shù)據(jù)集篩選會(huì)生成一個(gè)新的文件夾，這個(gè)文件夾由主節(jié)點(diǎn)一邊生成一邊同步給訓(xùn)練節(jié)點(diǎn)，而訓(xùn)練節(jié)點(diǎn)在同步時(shí)間到來時(shí)也會(huì)給反向同步給主節(jié)點(diǎn)，這就會(huì)導(dǎo)致覆蓋掉主節(jié)點(diǎn)原本文件夾的目錄，從而破壞了原數(shù)據(jù)集。

解決辦法：根據(jù)規(guī)則關(guān)閉到從節(jié)點(diǎn)到主節(jié)點(diǎn)的同步，避免反向同步。

這個(gè)問題想要解決除非更換lsyncd工具，否則沒有完美的方法?？梢試L試使用git大文件同步方案來替換lsyncd，當(dāng)然最好的方案還是基于IB網(wǎng)的NFS，既能滿足速度要求，又能避免冗余。

3.主節(jié)點(diǎn)異常退出，從節(jié)點(diǎn)不能退出

Pytorch實(shí)現(xiàn)多機(jī)多卡的分布式訓(xùn)練時(shí)會(huì)啟動(dòng)一個(gè)主節(jié)點(diǎn)和多個(gè)從節(jié)點(diǎn)。

現(xiàn)象：在多機(jī)多卡訓(xùn)練時(shí)，主節(jié)點(diǎn)異常退出時(shí)，從節(jié)點(diǎn)不能正常退出。主節(jié)點(diǎn)可能是因?yàn)樽x取數(shù)據(jù)集失敗或者GPU顯存不夠等原因退出，從節(jié)點(diǎn)會(huì)一直阻塞，并且顯示占用GPU顯存。

原因：在基于Pytorch的分布式中，使用nccl作為后端通信機(jī)制時(shí)，是沒有超時(shí)功能的。如果主服務(wù)阻塞，那么從服務(wù)會(huì)一直等待。

解決辦法：設(shè)置一個(gè)環(huán)境變量，NCCL_ASYNC_ERROR_HANDLING=1，然后給ddp進(jìn)程組設(shè)置超時(shí)時(shí)間

import torch.distributed as dist

dist.init_process_group(
    …
    backend="nccl",
    timeout=timedelta(secnotallow=60)
)

六、一種解決疑難問題的套路

遇到的這么多問題，如果全都是靠蠻力解決，那我頭上的頭發(fā)也保不住了。在這個(gè)過程中，我使用自己總結(jié)的一種解決疑難雜癥的思路去分析問題，解決問題，我把它叫做解決疑難問題的套路。

解決疑難問題的套路包含了分析和解決，簡(jiǎn)單來說分為三步：?jiǎn)栴}的現(xiàn)象是什么？已知內(nèi)容是什么？列出合理的猜測(cè)。

下面分別介紹每一步做什么。

1.問題的現(xiàn)象

想要解決一個(gè)問題首先要非常清楚問題是什么，所以第一步就是要搞清楚問題的現(xiàn)象是什么。以主節(jié)點(diǎn)異常退出，從節(jié)點(diǎn)不能退出為例，這個(gè)問題的現(xiàn)象就是當(dāng)主節(jié)點(diǎn)訓(xùn)練容器exit之后，從節(jié)點(diǎn)繼續(xù)運(yùn)行，不會(huì)退出。

有時(shí)看到的還不一定是真正的現(xiàn)象，需要稍作分析判斷，找出真正的現(xiàn)象，否則可能會(huì)南轅北轍。

2.已知的內(nèi)容

在知道問題的現(xiàn)象之后，列出已經(jīng)掌握肯定的、準(zhǔn)確無誤的線索。這些線索是解決問題的基礎(chǔ)、靈感、出發(fā)點(diǎn)。比如可以是一些計(jì)算機(jī)基礎(chǔ)知識(shí)，也可以是在這個(gè)場(chǎng)景下反復(fù)實(shí)驗(yàn)得到的結(jié)論。以主節(jié)點(diǎn)異常退出，從節(jié)點(diǎn)不能退出為例，已知的內(nèi)容是主節(jié)點(diǎn)和訓(xùn)練節(jié)點(diǎn)之間網(wǎng)絡(luò)肯定是互通的，排除網(wǎng)絡(luò)不達(dá)的可能。

列出已知內(nèi)容，能夠收縮猜想的范圍，排除疑點(diǎn)，減少可能性。

3.列出合理的猜測(cè)

在了解現(xiàn)象知道肯定的線索的之后，就能做出合理的猜測(cè)。最后一步就是匯總前面掌握的情況，從現(xiàn)象出發(fā)，根據(jù)已知的線索，列出可能產(chǎn)生問題的原因。以主節(jié)點(diǎn)異常退出，從節(jié)點(diǎn)不能退出為例，可能的原因包括：

（1）訓(xùn)練節(jié)點(diǎn)容器沒有捕獲到退出信號(hào)；

（2）NCCL主從進(jìn)程沒有斷開，一直阻塞；

（3）NCCL主從進(jìn)程斷開，程序沒有捕獲異常。

最后逐一驗(yàn)證猜想，這個(gè)過程中可能解決問題，可能發(fā)現(xiàn)新的線索。如果不能解決問題，再來一輪，基于上一輪的掌握的新線索做出合理的猜想，驗(yàn)證所有的猜想。掌握的線索越來越多，問題的范圍越來越小，最終一定抓住這個(gè)bug。

貼一個(gè)問題的解決過程，如下。

七、收獲

經(jīng)過這一段時(shí)間的開發(fā)，接觸到很多新知識(shí)，收獲也還不錯(cuò)。

1.學(xué)習(xí)模型訓(xùn)練框架Pytorch dpp

學(xué)習(xí)了Pytorch的實(shí)現(xiàn)，了解了模型訓(xùn)練過程，了解Pytorch ddp的原理，學(xué)習(xí)了一個(gè)進(jìn)程等待的巧妙方法。

@contextmanager
def torch_distributed_zero_first(local_rank: int):
    # Decorator to make all processes in distributed training wait for each local_master to do something
    if local_rank not in [-1, 0]:
        dist.barrier(device_ids=[local_rank])
    yield
    if local_rank == 0:
        dist.barrier(device_ids=[0])


# 使用該裝飾器下載資源
with torch_distributed_zero_first(LOCAL_RANK):
    weights = attempt_download(weights)  # download if not found locally

dist.barrier 是PyTorch 的分布式通信庫，會(huì)阻塞等待，所有注冊(cè)進(jìn)程都到齊了才會(huì)通過。

從語法上來說：使用上下文管理器加yield關(guān)鍵字實(shí)現(xiàn)一個(gè)裝飾器；從功能上來說：讓所有子進(jìn)程等待，放過主進(jìn)程做一些操作，等主進(jìn)程操作完成才會(huì)放行所有進(jìn)程。用于只有主進(jìn)程才能操作的場(chǎng)景。

2.做復(fù)雜的技術(shù)方案

該技術(shù)方案是我做過最復(fù)雜的技術(shù)方案，有架構(gòu)設(shè)計(jì)、技術(shù)選型、技術(shù)優(yōu)劣對(duì)比、潛在問題、解決辦法等?？偨Y(jié)出一個(gè)寫技術(shù)方案的模板：

（1）關(guān)鍵技術(shù)分析

（2）要實(shí)現(xiàn)的功能

（3）技術(shù)難點(diǎn)

（4）實(shí)現(xiàn)方案

（5）優(yōu)劣對(duì)比

（6）最佳方案理由

（7）遺留問題解決辦法

3.技術(shù)選型

技術(shù)選型是一個(gè)比較困難的工作，我在選擇的rpc框架zerorpc和文件同步工具lsyncd一定程度上都存在問題。zerorpc不能并發(fā)、lsyncd存在雙向同步的問題。

技術(shù)選項(xiàng)要從幾個(gè)點(diǎn)出發(fā)：

（1）是否能夠滿足業(yè)務(wù)邏輯

（2）是否符合當(dāng)前技術(shù)棧

（3）技術(shù)復(fù)雜性不能過高

（4）是否有明顯的缺陷

像是zerorpc框架就有明顯的缺陷，在查閱資料的時(shí)候也見過有些文章提到zerorpc是基于協(xié)程的并發(fā)，但當(dāng)時(shí)并沒有仔細(xì)思考，直到碰見并發(fā)請(qǐng)求才發(fā)現(xiàn)問題。

4.個(gè)人感受

一個(gè)人開發(fā)一個(gè)項(xiàng)目，我最大的感受就是太爽了。這種感覺就像自己在蓋一個(gè)大別墅，圖紙?jiān)O(shè)計(jì)是我做，搬磚砌墻我能搞定，最后造出一個(gè)完全屬于自己審美風(fēng)格的別墅。