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