Python PyTorch ElasticAgent用法及代碼示例

本文簡要介紹python語言中 torch.distributed.elastic.agent.server.ElasticAgent 的用法。

用法:

class torch.distributed.elastic.agent.server.ElasticAgent

代理進程負責管理一個或多個工作進程。假定工作進程是常規分布式PyTorch 腳本。當代理創建工作進程時，代理會為工作進程提供必要的信息，以正確初始化 Torch 進程組。

agent-to-worker 的確切部署拓撲和比例取決於代理的具體實現和用戶的工作安置偏好。例如，要在具有 8 個訓練器(每個 GPU 一個)的 GPU 上運行分布式訓練作業，可以：

1. 使用 8 個單 GPU 實例，每個實例放置一個代理，每個代理管理 1 個工作人員。

2. 使用 4 個雙 GPU 實例，每個實例放置一個代理，每個代理管理 2 個工作人員。

3. 使用 2 個四核 GPU 實例，每個實例放置一個代理，每個代理管理 4 個工作人員。

4. 使用 1 x 8 GPU 實例，每個實例放置一個代理，每個代理管理 8 個工作人員。

用法

group_result = agent.run()
 if group_result.is_failed():
   # workers failed
   failure = group_result.failures[0]
   log.exception(f"worker 0 failed with exit code : {failure.exit_code}")
 else:
   return group_result.return_values[0] # return rank 0's results