Ctrlk
For the complete documentation index, see llms.txt. This page is also available as Markdown.

PyTorch with WireGuard

How Kinesis supports PyTorch DDP?

By Toshihito Kikuchi

PyTorch DDP (= Distributed Data Parallel) is becoming the industry standard to parallelize your model training process across multiple GPUs and machines. Since one of our missions is to provide scalable GPU computes, it’s essential to support PyTorch DDP on Kinesis Network.

Does Kinesis Network support PyTorch DDP? Yes, but it isn't quite "plug-and-play" yet — it requires a bit of manual tuning to get everything running smoothly. We will fully integrate it soon, but for now, you need some special configuration to run it. In this article, I’d like to explain why it’s a little bit tricky and what we're implementing behind the scene.

Challenge in PyTorch networking

Okay, I know you have a model to train. You wrote a python script with DistributedDataParallel, something like this.

import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup():
    # Initialize the process group
    # NCCL is the standard backend for NVIDIA GPUs
    dist.init_process_group(backend="nccl")
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))

def cleanup():
    dist.destroy_process_group()

def run_training():
    setup()

    local_rank = int(os.environ["LOCAL_RANK"])
    steps = int(os.environ["STEPS"])
    rank = int(os.environ["RANK"])
    device = torch.device(f"cuda:{local_rank}")

    # 1. Define a tiny model
    model = nn.Linear(10, 10).to(device)
    model = DDP(model, device_ids=[local_rank])

    # 2. Setup Loss and Optimizer
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001)

    # 3. Simple training loop (Synthetic data)
    print(f"[Rank {rank}] Starting training...")

    for step in range(steps):
        # Create random data on the fly
        inputs = torch.randn(20, 10).to(device)
        labels = torch.randn(20, 10).to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_fn(outputs, labels)
        loss.backward()
        optimizer.step()

        if step % 10 == 0 and rank == 0:
            print(f"Step {step} | Loss: {loss.item():.4f}")

    print(f"[Rank {rank}] Training complete.")
    cleanup()

if __name__ == "__main__":
    run_training()

We usually use torchrun to kick a distributed training process. Below is the output to run a single-node, single-gpu training inside a container with detailed logs.

It looks pretty straightforward. Can we just specify --nnodes and --node_rank accordingly to run this on multiple gpu nodes? Well, unfortunately no, that’s why I’m writing this article.

The challenge is networking. In the example above, I specified --rdzv_endpoint=127.0.0.1:29500. Does this mean each node communicates with this endpoint during a training process, like the Hub-and-Spoke topology? The answer is no. As the name implies, it’s a rendezvous point. Every node meets there, to retrieve the actual endpoint to communicate with. In the output above, you see master_addr=2f3d44d8b0ea and master_port=43093 , that’s the actual endpoint each node communicates with. The master port, 43093 in this case, is dynamically assigned, which means we cannot predict which port needs to be opened. This is a problem in a Docker environment because we need to expose ports on creation but we don't know it on creation. Should we use the host network --network=host or publish a wide range of ports like -p 1000:65535? Apparently that design is far from ideal. We want our containers to be contained as much as possible.

In the earlier article, I wrote we leverage WireGuard to bring home computers. We can get help from WireGuard to solve this PyTorch situation too. Once we set up a WireGuard network on every Docker containers, they communicate with one another via a single UDP port. This journey starts here.

WireGuard Setup

Setting up a WireGuard network is pretty easy. I skip the detailed setup steps in this article. Basically you need two things: 1) create a conf file like /etc/wireguard/wg0.conf and run wg-quick up wg0. Besides, when you create a Docker container, you need to specify --cap-add=NET_ADMIN and publish a UDP port. Lastly, you need to install the packages wireguard-tools iptables iproute2. The wireguard driver exists in the host as a part of Linux kernel, but you still need client tools to use it.

Once setup is done, you will see a virtual interface like wg0.

Debugging with GDB

Let's run torchrun with the WireGuard network address 10.100.10.1 as the rendezvous point.

It failed with timeout! It seems that the process cannot connect to the rendezvous point 10.100.10.1:29500.

The log says throwTimeoutError was raised at this line. Is it time to file an issue there? No, it’s not what this blog does. It’s time to attach debugger!

As you may know, torchrun is just a python script to kick torch.distributed.run.

To debug it, you launch python with gdb, run the script, and break it when it’s stuck before the timeout exception is thrown.

It’s running SocketConnectOp::tryConnect , which calls the standard connect function via tryConnectCore , that is expected to fail. Let’s double check. You can just set a breakpoint there.

It hit twice. The first one is from getnameinfo, which looks unrelated and we skip. The second one was called from SocketConnectOp::tryConnect and failed with EINPROGRESS (= 115). This is the one we're interested in. And we’re interested in the parameters of connect. Here’s assembly of where we are, immediately after the call to connect.

What we want is the 2nd parameter const struct sockaddr *addr, which is passed via $rsi in System V ABI. Since we already lost $rdi, we need restore it from the stack.

How to read this? The address family is 0x0a 0x00 , which means AF_INET6, and we can see the address is 0xff 0xff 0x0a 0x64 0x0a 0x01, which is an IPv4-mapped IPv6 address ::ffff:10.100.10.1. The port is 0x73 0x3c , which is 0x733c=29500. This means the script simply tries to connect to the rendezvous point we specified, 10.100.10.1:29500, but it failed. This means somebody should be listening on the endpoint. Let's find out.

Okay, nobody is listening on the endpoint, that's why connect failed.

The next thing to do is to see the positive behavior. We know this works with 127.0.0.1. Let’s see if the endpoint is listened on in that case.

And see this! The endpoint is there.

So the problem is the script doesn’t start listening on the endpoint if the address is 10.100.10.1, while it does on 127.0.0.1.

Do you know which function to set a breakpoint on? Probably listen or bind? Well, in this case, I did some homework for you already and it turned out bind was the one. So let’s do it.

Okay, we got it. In this positive scenario, we start listening on the endpoint through TCPServer::start.

Now, what happens if the address is 10.100.10.1? If you look at the debugger output earlier carefully, we called connect inside TCPStore::TCPStore through TCPClient::connect. So we know TCPStore::TCPStore is surely called. Let’s see if we call TCPServer::start or not.

See? We reached TCPClient::connect without hitting TCPServer::start. This is the problem. Let’s go check the code and see how we call TCPServer::start. Code is here. It’s behind the check if (opts.isServer) {. Does this mean isServer was false in our case? Let's confirm it on debugger.

As I commented inline, cmpb $0x0,0x2(%rsi) is checking the flag.

It’s zero! This is why we skip listening on the rendezvous point!

And it’s from the parameter opts. Where does it come from? Who instantiates TCPStore ?

The symbol of frame #1 is insanely long. See pybind11 namespace. It’s python binding, meaning it’s python script instantiating TCPStore in C++. opts.isServer is also from python.

Debugging with PDB

How to debug Python? Do we need VSCode or some fancy IDE to debug it? That’s not what this blog does. Do we debug python binding? Well, that's too ambitious.

One of the great functionality Python provides, compared to other script languages, is Python has a built-in console debugger, pdb. You don’t need additional components to live debug python code. Very handy.

To use pdb, you modify your script to break at the beginning. Since we’re running our script with torchrun, we make a little modification in torchrun itself, just adding one line breakpoint() before main().

Where does it call C++? A formal approach would be to start from the log Starting elastic_operator with launch configs:, which is from this line. In our case, however, let’s simply search the repo for the keyword TCPStore . This function _create_tcp_store looks suspicious. Since python is a script language, you can set a breakpoint accurately with a source line.

Okay, we got it. _create_tcp_store is instantiating TCPStore with is_master=is_server, which is False.

Where does this is_host come from? Since cfg_is_host is None , it must be from _matches_machine_hostname. Let’s run this function line by line.

We clearly see the problem. First, we get the hostname with gethostname(), which is 2f3d44d8b0ea (It matches the container’s ID). And we get the IP address associated with the hostname via getaddrinfo, which returns only 172.17.0.2, the one mapped to the default bridge interface (eth0). Since it doesn’t match 10.100.10.1, it thinks “I’m not the master. Somebody else should start listening on the rendezvous endpoint.”

Ideally _matches_machine_hostname should iterate all IP addressed assigned to check any of them matches the host. We may consider sending a PR to them, but for now, is there a good way to work around this behavior?

There is. See the beginning of _create_tcp_store. It’s overwritable through cfg_is_host, coming from params.

Where is this RendezvousParameters created? It’s way above _matches_machine_hostname. It’s in the function run. torchrun takes a rarely-used parameter --rdzv_conf where we can specify key-value pairs. What we want to specify is is_host. This looks promising. Let’s try it.

It worked like a charm!

Now, it’s time to run a training on two nodes?

Training on two nodes

This is the first node; master node, rank 0.

This is the second node; rank 1. Ouch! It failed!

The reason is obvious. The message TCP client failed to connect/validate to host 2f3d44d8b0ea:44403 implies It tried to communicate via hostname, which is not reachable.

Looking at the first node’s log carefully, you see master_addr=2f3d44d8b0ea. This means the first node advertised itself with the hostname. In this case, we want to use the IP address.

Let’s find out where torchrun prints this log. Well, it’s in the log: torch/distributed/elastic/agent/server/api.py:539, here​. The line master_addr = spec.master_addr or rdzv_info.bootstrap_store_info.master_addr does that. It’s easy to confirm it with pdb.

It’s coming from rdzv_info.bootstrap_store_info.master_addr. Who set the hostname in it? We need to step in the line spec.rdzv_handler.next_rendezvous().

Alright, we have the hostname in self._this_node.addr, which is passed to RendezvousStoreInfo.build to build _bootstrap_store_info. Who set self._this_node? It’s in DynamicRendezvousHandler.__init__.

And node is coming from cls._node_desc_generator.generate(local_addr) in from_backend. Let’s see this generate function.

local_addr is None, and see the function generate. It does local_addr or socket.getfqdn()! This is where PyTorch prefers the hostname though the varialbe name is local_addr.

Now the question is how to overwrite it. In other words, how to set the IP address in this local_addr? We backtrack a little bit more. In create_handler, we pass params.local_addr to DynamicRendezvousHandler.from_backend. And params is RendezvousParameters. We already know this class, right? It’s where we store is_host earlier.

So we specify --rdzv_conf=is_host=True,local_addr=10.100.10.1?

It was close, but failed with TypeError: torch.distributed.elastic.rendezvous.api.RendezvousParameters() got multiple values for keyword argument 'local_addr' because PyTorch instantiates RendezvousParameters as below. It explicitly specifies local_addr already, so we cannot include it in config.rdzv_configs.

Where does config.local_addr come from? It’s from args.local_addr in config_from_args. torchrun directly takes --local_addr parameter.

Let’s try one more time! This is from the master node.

And this is from the second node. It all worked!

Conclusion

We just confirmed PyTorch DDP worked with a WireGuard network established between Docker containers. There are special configurations needed:

  • Specify --rdzv_conf=is_host=1 for the master node because PyTorch doesn't see secondary IP addresses to check if the rendezvous endpoint is itself or not

  • Specify --local_addr for every node to communicate via IP addresses instead of hostname

Actually specifying these parameters along with other standard DDP parameters such as rdzv_endpoint or nnodes is error-prone. I created a sample Docker image to semi-automated these configurations.

The question remains: Is this a bug in PyTorch we should report?

My answer is yes. --rdzv_conf=is_host=1 is a great workaround, but PyTorch should check all IP addresses assigned. At the same time, I cannot find a clean solution for this yet. It's strangely difficult to enumerate all IP addressese on Linux. One solution AI suggested is to use fcntl.ioctl , but I believe it will be rejected because it looks too "C-style" or less compatible. I'll think through it further.

Anyway, Kinesis now supports PyTorch DDP. If you have a model to train, go to https://portal.kinesis.network and run it!

PreviousSupportNextWhy Compute Should Work Like Power on Tap

Last updated

Was this helpful?