Libuv TCPStore 后端介绍¶
创建时间： Jul 22， 2024 |上次更新时间： 2024-7-24 |上次验证： Nov 05， 2024
作者： Xilun Wu
注意
在 github 中查看和编辑本教程。
什么是新的 TCPStore 后端
将新的 libuv 后端与旧后端进行比较
如何启用使用旧版后端
PyTorch 2.4 或更高版本
阅读有关 TCPStore API 的信息。
介绍¶
最近，我们使用 libuv 推出了一个新的 TCPStore 服务器后端，libuv 是一个用于异步 I/O 的第三方库。这个新的服务器后端旨在 解决大规模分布式训练作业（例如排名超过 1024 的作业）中的可扩展性和稳健性挑战。我们运行了一系列 基准测试将 libuv 后端与旧后端进行比较，实验结果表明 store 初始化有显著改进 时间，并在存储 I/O 操作中保持了相当的性能。
由于这些发现，libuv 后端已在 PyTorch 2.4 中设置为默认的 TCPStore 服务器后端。预计此更改将增强 分布式训练作业的性能和可扩展性。
此更改引入了与 store 初始化的轻微不兼容。对于希望继续使用旧版后端的用户，本教程将 提供有关如何指定使用以前的 TCPStore 服务器后端的指导。
性能基准¶
为了更好地展示我们新的 libuv TCPStore 后端的优势，我们在从 1024 （1K） 到 98304 （96K） 等级的广泛作业大小上建立了基准。 我们首先使用以下代码片段测量了 TCPStore 初始化时间：
import logging
import os
from time import perf_counter
import torch
import torch.distributed as dist
logger: logging.Logger = logging.getLogger(__name__)
# Env var are preset when launching the benchmark
env_rank = os.environ.get("RANK", 0)
env_world_size = os.environ.get("WORLD_SIZE", 1)
env_master_addr = os.environ.get("MASTER_ADDR", "localhost")
env_master_port = os.environ.get("MASTER_PORT", "23456")
start = perf_counter()
tcp_store = dist.TCPStore(
env_master_addr,
int(env_master_port),
world_size=int(env_world_size),
is_master=(int(env_rank) == 0),
)
end = perf_counter()
time_elapsed = end - start
logger.info(
f"Complete TCPStore init with rank={env_rank}, world_size={env_world_size} in {time_elapsed} seconds."
)
由于在所有客户端都成功连接之前，TCPStore 服务器线程的执行将被阻止，因此我们将排名 0 上测得的时间作为总时间 TCPStore 初始化运行时。实验编号如下图所示：
图 1.显示了一些重要的证据表明 libuv 后端优于传统后端：
带有 libuv 后端的 TCPStore 总是比传统后端具有更快的初始化速度，尤其是在超大规模的情况下
传统后端将在 96K 规模的服务器-客户端连接时超时（例如，超过 30 分钟），而 libuv 后端在 100 秒内完成初始化。
我们做的第二个基准测试是测量 TCPStore 操作的运行时间：store_based_barrier
import logging
import os
import time
from datetime import timedelta
from time import perf_counter
import torch
import torch.distributed as dist
DistStoreError = torch._C._DistStoreError
logger: logging.Logger = logging.getLogger(__name__)
# since dist._store_based_barrier is a private function and cannot be directly called, we need to write a function which does the same
def store_based_barrier(
rank,
store,
group_name,
rendezvous_count,
timeout=dist.constants.default_pg_timeout,
logging_interval=timedelta(seconds=10),
):
store_key = f"store_based_barrier_key:{group_name}"
store.add(store_key, 1)
world_size = rendezvous_count
worker_count = store.add(store_key, 0)
last_worker_key = f"{store_key}:last_worker"
if worker_count == world_size:
store.set(last_worker_key, "1")
start = time.time()
while True:
try:
# This will throw an exception after the logging_interval in which we print out
# the status of the group or time out officially, throwing runtime error
store.wait([last_worker_key], logging_interval)
break
except RuntimeError as e:
worker_count = store.add(store_key, 0)
# Print status periodically to keep track.
logger.info(
"Waiting in store based barrier to initialize process group for "
"rank: %s, key: %s (world_size=%s, num_workers_joined=%s, timeout=%s)"
"error: %s",
rank,
store_key,
world_size,
worker_count,
timeout,
e,
)
if timedelta(seconds=(time.time() - start)) > timeout:
raise DistStoreError(
"Timed out initializing process group in store based barrier on "
"rank {}, for key: {} (world_size={}, num_workers_joined={}, timeout={})".format(
rank, store_key, world_size, worker_count, timeout
)
)
logger.info(
"Rank %s: Completed store-based barrier for key:%s with %s nodes.",
rank,
store_key,
world_size,
)
# Env var are preset when launching the benchmark
env_rank = os.environ.get("RANK", 0)
env_world_size = os.environ.get("WORLD_SIZE", 1)
env_master_addr = os.environ.get("MASTER_ADDR", "localhost")
env_master_port = os.environ.get("MASTER_PORT", "23456")
tcp_store = dist.TCPStore(
env_master_addr,
int(env_master_port),
world_size=int(env_world_size),
is_master=(int(env_rank) == 0),
)
# sync workers
store_based_barrier(int(env_rank), tcp_store, "tcpstore_test", int(env_world_size))
number_runs = 10
start = perf_counter()
for _ in range(number_runs):
store_based_barrier(
int(env_rank), tcp_store, "tcpstore_test", int(env_world_size)
)
end = perf_counter()
time_elapsed = end - start
logger.info(
f"Complete {number_runs} TCPStore barrier runs with rank={env_rank}, world_size={env_world_size} in {time_elapsed} seconds."
)
我们通过将排名 0 上测得的运行时间除以来计算平均值，并在下图中报告：number_runs
图 2.显示 libuv 后端的 I/O 性能与 legacy backend 相当：
libuv 后端在秩数方面在整个范围内具有可比的性能
随着等级数量的增加，libuv 后端运行时比传统后端更稳定
冲击¶
用户可能需要注意的一个不兼容问题是，TCPStore 目前在使用 libuv 后端时不支持使用 a 进行初始化。
如果用户想继续使用这个初始化方法，用户只需传递 to stay 使用旧的 TCPStore 后端即可。listen_fduse_libuv=False
import socket
import torch
import torch.distributed as dist
listen_sock: socket.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
listen_sock.bind(("localhost", 0))
addr, port, *_ = listen_sock.getsockname()
listen_fd = listen_sock.detach()
tcpstore = dist.TCPStore(addr, port, 1, True, master_listen_fd=listen_fd) # expect NotImplementedError
tcpstore = dist.TCPStore(addr, port, 1, True, master_listen_fd=listen_fd, use_libuv=False) # OK. Use legacy backend
退出路由 1：传递到 TCPStore 初始化use_libuv=False¶
如上面的代码片段所示，如果用户调用 TCPStore init 方法来创建一个 store，只需传递即可允许用户继续使用旧的
TCPStore 后端。此覆盖的优先级高于确定 TCPStore 服务器应选择哪个后端的其他方法。use_libuv=False
退出路由 2：在 ProcessGroup 初始化时添加到use_libuv=0init_method¶
ProcessGroup如果用户未显式将 TCPStore 传递给其初始化，则创建 TCPStore。用户可以在
初始化 .此方法的优先级低于 Exit Route 1。use_libuv=0init_methodProcessGroup
import torch
import torch.distributed as dist
addr = "localhost"
port = 23456
dist.init_process_group(
backend="cpu:gloo,cuda:nccl",
rank=0,
world_size=1,
init_method=f"tcp://{addr}:{port}?use_libuv=0",
)
dist.destroy_process_group()
退出路由 3：将 Environment Variable 设置为USE_LIBUV0¶
当 ProcessGroup 创建 TCPStore 时，它还会检查环境 vairable 以确定要使用的 TCPStore 后端。用户可以设置环境
变量来指定使用旧的 TCPStore 后端。此方法的优先级低于 Exit Route 2（退出路由 2），例如，如果用户将 environment （环境）
变量传递给 ，并且还传入 ，则将选择旧的 store 后端。USE_LIBUV"USE_LIBUV""0"USE_LIBUV1use_libuv=0init_method
import os
import torch
import torch.distributed as dist
addr = "localhost"
port = 23456
os.environ["USE_LIBUV"] = "0"
dist.init_process_group(
backend="cpu:gloo,cuda:nccl",
rank=0,
world_size=1,
init_method=f"tcp://{addr}:{port}",
)
dist.destroy_process_group()
结论¶
在 PyTorch 2.4 中，我们将新的 libuv TCPStore 后端设为默认后端。尽管新后端与从 初始化 不兼容，但它
在大规模存储初始化方面显示出显著的性能改进，在小型/中型/大规模的存储 I/O 上表现出兼容性能，这
为分布式训练的控制平面带来了重大好处。本教程解释了我们的动机，介绍了性能基准测试，并通知了用户
的潜在影响，并引入了三种退出路由以继续使用旧版后端。从长远来看，我们的目标是最终弃用遗留后端。listen_fd