Fail to build connection at crypten.init() on two machines or a single machine's two sessions

[dorafan]

I just try to let my two machines run a simple DNN with crypten. However, I always fail to establish communication at crypten.init(0, both on two different machines, or on single machine's two sessions using 127.0.0.1

I'm using conda:
crypten 0.4.1
python 3.12.7
pytorch 2.5.1

I then try to check my settings using a simple testing script. Here is the code:

import os
import argparse
import crypten
import crypten.communicator as comm
from crypten.config import cfg

import torch.distributed as dist

def run(rank, world_size, master_addr, master_port):
    # Set environment variables for CrypTen
    os.environ["RANK"] = str(rank)
    os.environ["WORLD_SIZE"] = str(world_size)
    os.environ["MASTER_ADDR"] = master_addr
    os.environ["MASTER_PORT"] = master_port

    print(f"Process {rank}: Initializing CrypTen with MASTER_ADDR={master_addr}, MASTER_PORT={master_port}")

    # dist.init_process_group("gloo", init_method='env://')

    # Initialize CrypTen
    crypten.init()
    print(f"Process {rank}: CrypTen initialized.")

    # Check communicator status
    print(f"Process {rank}: Communicator rank={comm.get().get_rank()}, world_size={comm.get().get_world_size()}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="CrypTen Initialization Debug Test")
    parser.add_argument("--rank", type=int, required=True, help="Rank of the current process")
    parser.add_argument("--world_size", type=int, default=2, help="Total number of processes")
    parser.add_argument("--master_addr", type=str, default="127.0.0.1", help="Master address")
    parser.add_argument("--master_port", type=str, default="29500", help="Master port")
    parser.add_argument("--local_rank", type=int)
    args = parser.parse_args()

    run(args.rank, args.world_size, args.master_addr, args.master_port)

after run python script.py --rank=0 --world_size=2 --master_addr=127.0.0.1 --master_port=29500
and python script.py --rank=1 --world_size=2 --master_addr=127.0.0.1 --master_port=29500
both session hangs at crypten.init()

I'm using Mac, and already verified that torch.distributed is available.

[dorafan]

I also have these two scripts:

# simple_distributed_test.py
import os
import argparse
import crypten
import crypten.communicator as comm
from crypten.config import cfg

def run(rank, world_size, master_addr, master_port):
    # Set environment variables for CrypTen
    os.environ["RANK"] = str(rank)
    os.environ["WORLD_SIZE"] = str(world_size)
    os.environ["MASTER_ADDR"] = master_addr
    os.environ["MASTER_PORT"] = master_port

    # Initialize CrypTen
    crypten.init()
    print(f"Process {rank} initialized.")

    if rank == 0:
        # Sender
        a = crypten.cryptensor([1.0, 2.0, 3.0], requires_grad=True)
        comm.get().send(a, dst=1)
        print(f"Process {rank} sent tensor: {a}")
    elif rank == 1:
        # Receiver
        b = comm.get().recv(src=0)
        print(f"Process {rank} received tensor: {b}")
        decrypted_b = b.get_plain_text()
        print(f"Process {rank} decrypted tensor: {decrypted_b}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Simple CrypTen Distributed Test")
    parser.add_argument("--rank", type=int, required=True, help="Rank of the current process")
    parser.add_argument("--world_size", type=int, default=2, help="Total number of processes")
    parser.add_argument("--master_addr", type=str, default="127.0.0.1", help="Master address")
    parser.add_argument("--master_port", type=str, default="29500", help="Master port")
    args = parser.parse_args()

    run(args.rank, args.world_size, args.master_addr, args.master_port)

which hangs at crypten.init()

and this pytorch version:

# simple_pytorch_distributed_test.py
import os
import argparse
import torch
import torch.distributed as dist

def run(rank, world_size, master_addr, master_port):
    os.environ["RANK"] = str(rank)
    os.environ["WORLD_SIZE"] = str(world_size)
    os.environ["MASTER_ADDR"] = master_addr
    os.environ["MASTER_PORT"] = master_port

    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    print(f"Process {rank} initialized.")

    if rank == 0:
        tensor = torch.tensor([1.0, 2.0, 3.0])
        dist.send(tensor, dst=1)
        print(f"Process {rank} sent tensor: {tensor}")
    elif rank == 1:
        tensor = torch.zeros(3)
        dist.recv(tensor, src=0)
        print(f"Process {rank} received tensor: {tensor}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Simple PyTorch Distributed Test")
    parser.add_argument("--rank", type=int, required=True, help="Rank of the process")
    parser.add_argument("--world_size", type=int, default=2, help="Number of processes")
    parser.add_argument("--master_addr", type=str, default="127.0.0.1", help="Master address")
    parser.add_argument("--master_port", type=str, default="29500", help="Master port")
    args = parser.parse_args()

    run(args.rank, args.world_size, args.master_addr, args.master_port)

which works well on the single machine, two sessions.

[dorafan]

UPDATE:

after I modified the simple_distributed_test.py as the following:

# simple_distributed_test.py
import os
import argparse
import crypten
import crypten.communicator as comm
from crypten.config import cfg

import torch.distributed as dist

def run(rank, world_size, master_addr, master_port):
    # Set environment variables for CrypTen
    env = os.environ.copy()
    env["RENDEZVOUS"] = "env://"
    env["WORLD_SIZE"] = str(world_size)
    env["RANK"] = str(rank)
    env["MASTER_ADDR"] = master_addr
    env["MASTER_PORT"] = str(master_port)

    print(f"Process {rank}: Initializing CrypTen with MASTER_ADDR={master_addr}, MASTER_PORT={master_port}")

    # Initialize CrypTen
    crypten.init()
    print(f"Process {rank} initialized.")

    # Introduce a barrier to synchronize both ranks
    comm.get().barrier()
    print(f"Process {rank}: Passed barrier.")

    if rank == 0:
        # Sender
        a = crypten.cryptensor([1.0, 2.0, 3.0], requires_grad=True)
        comm.get().send(a, dst=1)
        print(f"Process {rank} sent tensor: {a}")
    elif rank == 1:
        # Receiver
        b = comm.get().recv(src=0)
        print(f"Process {rank} received tensor: {b}")
        decrypted_b = b.get_plain_text()
        print(f"Process {rank} decrypted tensor: {decrypted_b}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Simple CrypTen Distributed Test")
    parser.add_argument("--rank", type=int, required=True, help="Rank of the current process")
    parser.add_argument("--world_size", type=int, default=2, help="Total number of processes")
    parser.add_argument("--master_addr", type=str, default="127.0.0.1", help="Master address")
    parser.add_argument("--master_port", type=str, default="29500", help="Master port")
    parser.add_argument("--local_rank", type=int)
    args = parser.parse_args()

    run(args.rank, args.world_size, args.master_addr, args.master_port)

It could pass the init, but the rank 0 will directly end and won't wait for rank 1

[knottb]

CrypTen uses torch.distributed to handle the communication. This connection is set up here:

https://github.com/facebookresearch/CrypTen/blob/main/crypten/communicator/distributed_communicator.py#L33-L66

You can test whether this should work by debugging using torch.distributed with your arguments. If it doesn't work with torch.distributed, you will have to change your inputs.