背景
機器學習工作負載與傳統的工作負載相比，一個比較顯著的特點是對 GPU 的需求旺盛。在之前的文章中介紹過（https://mp.weixin.qq.com/s/Nasm-cXLtJObjLwLQHALmw 和 https://mp.weixin.qq.com/s/X4VDynLfKdVp-tyciQccyQ），目前 GPU 的視訊記憶體已經不足以跟上模型引數規模的發展。隨著 Transformer 等新的模型結構的出現，這一問題越來越顯著。演算法工程師們訓練模型所需要的資源越來越多，分散式訓練也隨之成為了工業界進行模型訓練的標準方式。
彈性訓練能夠在訓練過程中動態地調整參與訓練的例項數量，極大程度提高叢集資源的利用率。同時，配合雲上的競價例項等資源型別，能夠以更低的成本進行模型調優，進一步降本增效。在 PyTorch 最新發布的 1.9.0 版本中，其原本分散式訓練的方式torch.distributed.launch 即將被廢棄，轉而推薦使用者使用彈性的分散式訓練介面 torch.distributed.run。
藉此機會，我們對這一新特性進行簡單地介紹，並且與 Horovod Elastic 進行簡單地對比和分析。最後總結一下使用彈性訓練時，需要注意的問題。
PyTorch 1.9.0 之前的設計
PyTorch 是目前最流行的深度學習框架之一，它最讓人稱道的是易用性。無論是單機訓練還是分散式訓練，PyTorch 都提供了簡潔的 API。
PyTorch 1.9.0 版本之前，分散式訓練的方式通常是通過如下的方式進行。
python -m torch.distributed.launch
--nnodes=NODE_SIZE
--nproc_per_node=TRAINERS_PER_NODE
--node_rank=NODE_RANK
--master_port=HOST_PORT
--master_addr=HOST_NODE_ADDR
YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)
其中 nnodes 是參與訓練的節點個數，nproc_per_node 是每個節點上執行的程式數量。node_rank 是當前節點的識別符號，master_addr 和 master_port 是 master 監聽的地址和埠。torch.distributed.launch 會設定一些環境變數，其中包括 WORLD_SIZE 和 MASTER_PORT、MASTER_ADDR 等。
隨後在當前機器上會建立對應程式進行訓練，當前機器會有 TRAINERS_PER_NODE 個程式，這些程式組成了一個 local worker group。一共有 NODE_SIZE 個機器參與訓練，一共有 NODE_SIZE * TRAINERS_PER_NODE 個程式。如果想要發起一個分散式訓練任務，需要在所有的機器上執行相應的命令。
PyTorch 1.9.0 中的新設計
在 PyTorch 1.9 中，torch.distributed.launch 即將被廢棄，取而代之的是基於 pytorch/elastic 的 torch.distributed.run。這一新的方式與之前相比有一些使用上的改動，如下所示。
python -m torch.distributed.run
--nnodes=MIN_SIZE:MAX_SIZE
--nproc_per_node=TRAINERS_PER_NODE
--rdzv_id=JOB_ID
--rdzv_backend=c10d
--rdzv_endpoint=HOST_NODE_ADDR
YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)
它提供了一些新的能力：首先是更好的容錯，當 worker 失敗後會自動重啟繼續訓練；其次是 RANK 和 WORLD_SIZE 這些欄位不再需要手動設定。最後也是最重要的，支援彈性訓練，動態地增加或減少參與訓練的 worker 數量。在上面的例子中，nnodes 的設定不再是一個固定的值，而是一個區間。訓練任務可以容忍在這一區間範圍內的 worker 數量變化。
如果要支援彈效能力，訓練程式碼也需要進行一些修改。
def main():
args = parse_args(sys.argv[1:])
state = load_checkpoint(args.checkpoint_path)
initialize(state)
# torch.distributed.run ensure that this will work
# by exporting all the env vars needed to initialize the process group
torch.distributed.init_process_group(backend=args.backend)
for i in range(state.epoch, state.total_num_epochs)
for batch in iter(state.dataset)
train(batch, state.model)
state.epoch += 1
save_checkpoint(state)
其中比較明顯的變化是，使用者需要手動地處理 checkpoint。這是因為當 worker 出現失效時，所有的 worker 都會重啟，所以需要 checkpoint 機制來保證重啟後訓練能夠繼續下去。這一新的分散式訓練方式引入不少新的概念，包括 agent、rendezvous 等。接下來我們自使用者能接觸到的 torch.distributed.run 開始，介紹這些新的設計。
def run(args):
if args.standalone:
args.rdzv_backend = "c10d"
args.rdzv_endpoint = "localhost:29400"
args.rdzv_id = str(uuid.uuid4())
log.info(
f"\n**************************************\n"
f"Rendezvous info:\n"
f"--rdzv_backend={args.rdzv_backend} "
f"--rdzv_endpoint={args.rdzv_endpoint} "
f"--rdzv_id={args.rdzv_id}\n"
f"**************************************\n"
)
config, cmd, cmd_args = config_from_args(args)
elastic_launch(
config=config,
entrypoint=cmd,
)(*cmd_args)
其中主要區分了兩個模式，Standalone 模式和分散式模式。Standalone 模式是分散式模式的一種特例，它主要針對單機多 Worker 的方式提供了一些便利的設定，不再需要設定一些多餘的引數如 rdzv_backend 和 rdzv_endpoint 等。
兩者最後都會通過 elastic_launch 發起真正的訓練程式。elastic_launch 會通過 elastic agent 來管理 worker 的生命週期，它的返回是每個 worker 的輸出。
class elastic_launch:
...
def __call__(self, *args):
return launch_agent(self._config, self._entrypoint, list(args))
def launch_agent(
config: LaunchConfig,
entrypoint: Union[Callable, str, None],
args: List[Any],
) -> Dict[int, Any]:
...
agent = LocalElasticAgent(
spec=spec, start_method=config.start_method, log_dir=config.log_dir
)
...
result = agent.run()
...
return result.return_values
Elastic Agent 的設計：如何管理多個 worker 程式
elastic agent 是一個獨立的程式，負責管理其下的 workers。它起到了類似程式管理系統 supervisor 的作用，會在啟動的時候確保每個 worker 的設定正確。由於有關 WORLD_SIZE 和 RANK 的資訊不再需要使用者提供，elastic agent 會負責處理。
除此之外，worker 的失效也是由 elastic agent 負責捕獲處理。可以說 elastic agent 是彈性訓練中最核心的抽象概念。
上圖展示的是elastic agent 的工作原理。
不同的 elastic agent 之間通過 rendezvous 進行 worker 之間的相互發現和對成員變動的同步。與此同時，通過對 worker 程式的監控，來捕獲訓練過程中的失效。其中核心的邏輯都包裝在 LocalElasticAgent.run() 這一函式呼叫中。
def run(self, role: str = DEFAULT_ROLE) -> RunResult:
...
result = self._invoke_run(role)
return result
def _invoke_run(self, role: str = DEFAULT_ROLE) -> RunResult:
...
self._initialize_workers(self._worker_group)
while True:
...
run_result = self._monitor_workers(self._worker_group)
state = run_result.state
...
if state == WorkerState.SUCCEEDED:
...
return run_result
elif state in {WorkerState.UNHEALTHY, WorkerState.FAILED}:
if self._remaining_restarts > 0:
...
self._restart_workers(self._worker_group)
else:
...
return run_result
elif state == WorkerState.HEALTHY:
...
if num_nodes_waiting > 0:
...
self._restart_workers(self._worker_group)
else:
raise Exception(f"[{role}] Worker group in {state.name} state")
可以看到，核心的邏輯在 _invoke_run 中。其中 _initialize_workers 執行了大部分初始化的工作，其中包括為每個 worker 分配 RANK 等。在預設的實現中 elastic agent 和 workers 程式在同一機器上，因此 self._monitor_workers(self._worker_group) 通過 multiprocessing 對 workers 的執行狀態進行了監控。並且根據不同的狀態，進行不同的處理。
elastic agent 的可擴充套件性非常好，在 1.9.0 版本中，一共有三個 Agent，分別是 ElasticAgent、SimpleElasticAgent 和 LocalElasticAgent。
其中 ElasticAgent 是一個 Abstract Class，SimpleElasticAgent 對其中的某些函式進行了實現，而 LocalElasticAgent 則實現了管理單機上所有 worker 程式的 elastic agent。
SimpleElasticAgent 這一個抽象主要是為了方便擴充套件新的 agent 實現，比如如果你想通過一個 agent 管理多機上所有的 worker，而不只是本機上的 worker，則可以通過擴充套件 SimpleElasticAgent 來實現。
rendezvous 的設計：如何在不同的節點間確定 RANK
接下來，我們再看另外一個核心的抽象 rendezvous。為了實現彈性訓練，worker 之間要能夠動態地進行 membership 的變更。rendezvous 就是實現這一特性的用於同步的元件。
rendezvous 最核心的方法是：
@abstractmethod
def next_rendezvous(
self,
) -> Tuple[Store, int, int]:
"""Main entry-point into the rendezvous barrier.
Blocks until the rendezvous is complete and the current process is
included in the formed worker group, or a timeout occurs, or the
rendezvous was marked closed.
Returns:
A tuple of :py:class:`torch.distributed.Store`, ``rank``, and
``world size``.
Raises:
RendezvousClosedError:
The rendezvous is closed.
RendezvousConnectionError:
The connection to the rendezvous backend has failed.
RendezvousStateError:
The rendezvous state is corrupt.
RendezvousTimeoutError:
The rendezvous did not complete on time.
"""
如註釋所示，這一函式呼叫會被阻塞，直到 worker 的數量達到了要求。在 worker 被初始化，或者重啟的時候，這一函式都會被呼叫。當函式返回時，不同的 worker 會以返回中的 rank 作為唯一的標示。rendezvous 一共有四個實現，分別是 etcd、etcd-v2、c10d 和 static。
class EtcdRendezvousHandler(RendezvousHandler):
def next_rendezvous(self):
rdzv_version, rank, world_size = self._rdzv_impl.rendezvous_barrier()
log.info("Creating EtcdStore as the c10d::Store implementation")
store = self._rdzv_impl.setup_kv_store(rdzv_version)
return store, rank, world_size
其中 etcd 相關的是之前推薦使用的實現，在 c10d 出現後就不再推薦了。etcd 的實現中，不同 worker 之間的狀態通過 etcd 的 KV 介面儲存。
確定參與訓練的例項和對應的 RANK 的過程如下圖所示。
首先會在 /rdzv/active_version 下嘗試寫一個值 status: setup。在整個過程中，/rdzv/active_version 會作為儲存 rendezvous 過程中間狀態的 KV store，以及 rendezvous 過程中的排他鎖來使用。
如果寫失敗了，說明目前已經有對應的 rendezvous 過程正在進行中。
在成功後，會更新 /rdzv/version_counter 為原值加一。然後會建立一個目錄 /rdzv/v_${version_counter}。這些操作做完後，會將 /rdzv/active_version 的狀態寫為 joinable，這時就進入了 join 階段。
在 join 階段，不同的 agent 在鎖的保護下，會依次更新 /rdzv/active_version 下的 paticipants，分配到遞增的 rank，這裡的 rank 並不是每個 worker 程式分配到的 global rank，而是 agent 自己的 rank。worker 程式的 rank 會根據 agent rank 經過一定的計算得到。這也是一個非常容易混淆的設計，竊以為有優化的空間。
def init_phase(self):
try:
active_version = self.try_create_rendezvous()
state = json.loads(active_version.value)
log.info("New rendezvous state created: " + str(state))
except etcd.EtcdAlreadyExist:
# 已經有了一個新的 rendezvous 過程
active_version, state = self.get_rdzv_state()
# Note: it is possible for above query to fail (etcd.EtcdKeyNotFound),
# but this is ok for us - just means we'll restart from beginning.
log.info("Observed existing rendezvous state: " + str(state))
if state["status"] == "closed":
raise RendezvousClosedError()
if state["status"] == "joinable":
return self.join_phase(state["version"])
if state["status"] == "final":
self.handle_existing_rendezvous(state["version"])
raise EtcdRendezvousRetryImmediately()
self.try_wait_for_state_change(etcd_index=active_version.etcd_index + 1)
raise EtcdRendezvousRetryableFailure()
在參與訓練的節點達到 nnodes 的命令列引數中傳入的最小值時，會等待一定時間，在等待時間結束或者參與訓練的節點達到了 nnodes 設定的最大值時，會進入 frozen 階段。
在 fronzen 階段中，每個參與訓練的節點都需要通過在 /rdzv/v_${version_counter}/rank_${agent_rank} 下寫值的方式進行確認。在所有節點都確認完畢後，會進入最後的 final 階段。
在最後的 final 階段中，後續進入的 agent 都會 pending，已經達成 rendezvous 的節點上的 agent 會為其管理的 worker 程式分配 RANK。RANK 0 的例項會作為 master 的角色存在。隨後就會直接建立對應的 worker 程式。在預設的 LocalElasticAgent 中，會利用 python.multiprocessing 在本地建立多個程式。
@prof
def _start_workers(self, worker_group: WorkerGroup) -> Dict[int, Any]:
spec = worker_group.spec
store = worker_group.store
...
for worker in worker_group.workers:
local_rank = worker.local_rank
worker_env = {
"LOCAL_RANK": str(local_rank),
"RANK": str(worker.global_rank),
...
}
...
args[local_rank] = tuple(worker_args)
...
self._pcontext = start_processes(
name=spec.role,
entrypoint=spec.entrypoint,
args=args,
envs=envs,
log_dir=attempt_log_dir,
start_method=self._start_method,
redirects=spec.redirects,
tee=spec.tee,
)
return self._pcontext.pids()
c10d 新的設計
前文介紹了基於 etcd 的 rendezvous 實現，它可以保證多個例項之間對於參與訓練的節點共識的強一致，但是這也為 PyTorch 執行訓練任務引入了額外的依賴。因此 PyTorch 也提供了一個內建的實現 c10d。相比於基於 etcd 的實現，c10d 基於 TCP 來進行同步。
def create_backend(params: RendezvousParameters) -> Tuple[C10dRendezvousBackend, Store]:
...
if store_type == "file":
store = _create_file_store(params)
elif store_type == "tcp":
store = _create_tcp_store(params)
...
backend = C10dRendezvousBackend(store, params.run_id)
def _create_tcp_store(params: RendezvousParameters) -> TCPStore:
host, port = parse_rendezvous_endpoint(params.endpoint, default_port=29400)
...
for is_server in [is_host, False]:
...
store = TCPStore(
host, port, is_master=is_server, timeout=timedelta(seconds=read_timeout)
)
...
break
return store
c10d 是一個 client-server 的架構，其中的一個 agent 上會執行 c10d 的 TCPServer，它監聽給定的埠，提供了 compareAndSet、add 等原語。它也可以被理解為一個簡化的，提供 KV 介面的記憶體資料庫，類似於 Redis。有關 rendezvous 的同步，都是由各個 agent 通過一箇中心化的 agent 上的 c10d TCPServer 完成的。可以預見這樣的實現在可用性上相比於 etcd 是有一定差距的，但是勝在易用性。使用者如果使用 c10d，那麼不再需要運維一個 etcd 叢集。
PyTorch Elastic on Kubernetes
為了能夠享受到彈性訓練帶來的便利，PyTorch 同時提供了在 Kubernetes 上的支援。相比於 1.9.0 之前的版本，新版本的分散式訓練新增了一些新的引數。因此 PyTorch 社群在 Kubeflow PyTorch operator 的基礎上，對 CRD 進行了一些修改。一個典型的彈性訓練示例如下所示：
apiVersion: elastic.pytorch.org/v1alpha1
kind: ElasticJob
metadata:
name: imagenet
namespace: elastic-job
spec:
# Use "etcd-service:2379" if you already apply etcd.yaml
rdzvEndpoint: "<your_etcd_endpoint>:<your_etcd_port>"
minReplicas: 1
maxReplicas: 2
replicaSpecs:
Worker:
replicas: 2
restartPolicy: ExitCode
template:
apiVersion: v1
kind: Pod
spec:
containers:
- name: elasticjob-worker
image: torchelastic/examples:0.2.0
imagePullPolicy: Always
args:
- "--nproc_per_node=1"
- "/workspace/examples/imagenet/main.py"
- "--arch=resnet18"
- "--epochs=20"
- "--batch-size=32"
# number of data loader workers (NOT trainers)
# zero means load the data on the same process as the trainer
# this is set so that the container does not OOM since
# pytorch data loaders use shm
- "--workers=0"
- "/workspace/data/tiny-imagenet-200"
resources:
limits:
nvidia.com/gpu: 1
由於在最開始，基於 c10d 的 rendezvous 還沒有被支援，所以 CRD 中需要定義 rdzvEndpoint，指向一個已經部署好的 etcd 叢集。同時，使用者需要指定 minReplicas 和 maxReplicas。其他就與 Kubeflow PyTorchJob 並無二致。
PyTorch Elastic 與 Horovod Elastic
目前，兩者的設計從原理上來說並無二致。相比於 Horovod Elastic，PyTorch Elastic 提供了更靈活的擴充套件性，它提供了 agent、rendezvous 等介面，使用者可以根據需要進行擴充套件。但是從另外一個角度講，Horovod 的易用性做的更好。
PyTorch 並沒有提供儲存狀態的內建支援，為了能夠在 worker 程式失敗，重建訓練任務的時候，需要使用者自己實現儲存會載入 checkpoint 的邏輯；而 Horovod 則提供了內建的實現。
Horovod 和 PyTorch 在同步機制上也具有比較大的差異。Horovod Elastic 需要使用者提供一個指令碼 discovery_hosts.sh，幫助其在執行時獲得正在參與訓練的節點。
$ horovodrun -np 8 --host-discovery-script discover_hosts.sh python train.py
...
$ ./discover_hosts.sh
host-1:29500
host-2:29500
host-3:29500
這相當於將節點發現的邏輯交給使用者來實現。反觀 PyTorch，它利用 etcd、自身實現的 c10d 等元件解決節點間的相互發現問題，顯得更為精巧。
總結
在文章最後，我們總結一下目前實現彈性訓練需要注意的問題。
首先，也是最重要的，彈性訓練需要一種機制來解決節點/訓練程式間相互發現的問題。訓練過程中節點會動態地加入或者退出，如何讓其他的節點感知到這一變化，是這一機制主要面對的問題。目前的設計中，Horovod 將這一問題交給使用者來解決，Horovod 定期執行使用者定義的邏輯來發現目前的節點。PyTorch 通過第三方的分散式一致性中介軟體 etcd 等來實現高可用的節點發現。除此之外，也有一些探索性的工作，利用基於 Gossip 的協議來進行同步，在兼顧高可用的同時也沒有引入過多的元件。
其次，要實現彈性訓練還需要捕獲訓練失效。Horovod 和 PyTorch 都通過一個後臺程式（Horovod 中是 Driver，PyTorch 中是每個節點的 Local Elastic Agent）來實現這一邏輯。當程式 crash，或在梯度通訊中遇到問題時，後臺程式會捕獲到失效並且重新進行節點發現，然後重啟訓練。
最後，訓練時的資料切分的邏輯和學習率/ batch size 的設定也要對應進行修改。由於參與訓練的程式會動態的增減，因此可能需要根據新的訓練程式的規模來重新設定學習率和資料分配的邏輯，避免影響模型收斂。
在本文中，我們首先介紹了 PyTorch 1.9.0 版本中彈性訓練的設計與實現。然後分析總結了實現彈性訓練的方式和不同框架之間的設計差異。從我們的角度來看，彈性訓練能夠很好地貼合雲原生的趨勢，以極致的彈性來降低成本提高資源利用率，是未來的趨勢。因此目前我們也在積極參與 TensorFlow、PyTorch 和 Kubeflow 等社群的彈性訓練的社群貢獻工作，後續會發布更多的相關文章，感謝關注。
【騰訊雲原生】雲說新品、雲研新術、雲遊新活、雲賞資訊，掃碼關注同名公眾號，及時獲取更多幹貨！！
