训练脚本¶
如果您的训练脚本与 torch.distributed.launch
配合使用,则它将继续与 torchrun
配合使用,但存在以下差异
无需手动传递
RANK
、WORLD_SIZE
、MASTER_ADDR
和MASTER_PORT
。rdzv_backend
和rdzv_endpoint
可以提供。对于大多数用户,这将设置为c10d
(请参阅 rendezvous)。默认的rdzv_backend
创建一个非弹性 rendezvous,其中rdzv_endpoint
保存主机的地址。确保您的脚本中包含
load_checkpoint(path)
和save_checkpoint(path)
逻辑。当任意数量的工作进程失败时,我们将使用相同的程序参数重新启动所有工作进程,因此您将丢失到最近检查点为止的进度(请参阅 弹性启动)。use_env
标志已删除。如果您通过解析--local-rank
选项来解析本地排名,则需要从环境变量LOCAL_RANK
中获取本地排名(例如int(os.environ["LOCAL_RANK"])
)。
以下是一个说明性示例,该示例是一个在每个 epoch 进行检查点的训练脚本,因此在发生故障时,最坏情况下丢失的进度是一个完整 epoch 的训练。
def main():
args = parse_args(sys.argv[1:])
state = load_checkpoint(args.checkpoint_path)
initialize(state)
# torch.distributed.run ensures 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)
有关符合 torchelastic 的训练脚本的具体示例,请访问我们的 示例 页面。