Libuv TCPStore 后端简介
作者: Xilun Wu
你将学到什么
-
什么是新的 TCPStore 后端
-
将新的 libuv 后端与旧的后端进行比较
-
如何启用旧的后端
先决条件
-
PyTorch 2.4 或更高版本
-
阅读 TCPStore API 的相关内容。
简介
最近，我们使用了一个第三方库 libuv 推出了一个新的 TCPStore 服务器后端，该库用于异步 I/O。这个新的服务器后端旨在解决大规模分布式训练任务中的可扩展性和鲁棒性挑战，例如那些超过 1024 个秩的任务。我们进行了一系列基准测试，将 libuv 后端与旧的后端进行比较，实验结果表明，在存储初始化时间方面有显著改进，并且在存储 I/O 操作方面保持了相当的性能。
基于这些发现，libuv 后端已被设置为 PyTorch 2.4 中的默认 TCPStore 服务器后端。这一更改预计将提升分布式训练任务的性能和可扩展性。
该更改在存储初始化方面引入了轻微的兼容性问题。对于希望继续使用旧版后端的用户，教程将提供如何指定使用先前 TCPStore 服务器后端的指导。
性能基准测试
为了更好地展示我们新 libuv TCPStore 后端的好处，我们设置了一个涵盖广泛作业规模的基准测试，从 1024（1K）到 98304（96K）个进程。我们首先使用以下代码片段测量了 TCPStore 的初始化时间：
importlogging
importos
fromtimeimport perf_counter
importtorch
importtorch.distributedasdist
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 服务器线程的执行会阻塞，直到所有客户端成功连接为止，因此我们将 rank 0 上测量的时间作为 TCPStore 初始化的总运行时间。实验结果如下图所示：
图1展示了一些重要证据，表明libuv后端优于传统后端：
-
使用 libuv 后端的 TCPStore 总是比传统后端初始化更快，尤其是在超大规模情况下
-
传统后端在 96K 规模（例如，超过 30 分钟）的服务器-客户端连接时会超时，而 libuv 后端在 100 秒内完成了初始化。
我们进行的第二个基准测试是测量 TCPStore store_based_barrier 操作的运行时间：
importlogging
importos
importtime
fromdatetimeimport timedelta
fromtimeimport perf_counter
importtorch
importtorch.distributedasdist
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
defstore_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."
)
我们通过将 rank 0 上测量的运行时间除以 number_runs 来计算平均值，并在下图中报告该值：
图 2 显示，libuv 后端的 I/O 性能与传统后端相当：
-
在整个排名范围内，libuv 后端的性能表现相当
-
随着排名的增加，libuv 后端的运行时间比传统后端更加稳定
影响
用户可能需要注意的一个不兼容性是，在使用 libuv 后端时，TCPStore 目前不支持通过 listen_fd 进行初始化。如果用户希望继续使用这种初始化方法，可以简单地传递 use_libuv=False 以继续使用旧的 TCPStore 后端。
importsocket
importtorch
importtorch.distributedasdist
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 方法来创建存储，只需传递 use_libuv=False 即可让用户继续使用旧的 TCPStore 后端。此配置在决定 TCPStore 服务器应选择哪个后端时具有最高优先级，将覆盖其他方法。
退出路径 2：在 ProcessGroup 初始化时向 init_method 添加 use_libuv=0
ProcessGroup 会在用户未显式传递 TCPStore 时自动创建一个。用户可以在初始化 ProcessGroup 时，在 init_method 中添加 use_libuv=0 查询选项。此方法的优先级低于退出路径 1。
importtorch
importtorch.distributedasdist
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：将环境变量 USE_LIBUV 设置为 0
当 ProcessGroup 创建 TCPStore 时，它还会检查环境变量 USE_LIBUV 以确定使用哪个 TCPStore 后端。用户可以将环境变量 "USE_LIBUV" 设置为 "0" 来指定使用旧的 TCPStore 后端。此方法的优先级低于退出路径 2，例如，如果用户将环境变量 USE_LIBUV 设置为 1，同时在 init_method 中传递了 use_libuv=0，那么旧的存储后端将被选择。
importos
importtorch
importtorch.distributedasdist
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 后端设为了默认选项。尽管新后端与通过 listen_fd 初始化的方式存在不兼容性，但它在大规模存储初始化上表现出显著的性能提升，并且在小、中、大规模存储 I/O 操作中保持了兼容的性能，这为分布式训练的控制平面带来了重要优势。本教程解释了我们的动机，展示了性能基准测试，向用户通知了可能的影响，并介绍了三种继续使用旧版后端的退出路径。长远来看，我们计划最终弃用旧版后端。