Horovod multi-node fails to connect and hangs indefinitely

[jarednielsen]

Describe the bug
Running a horovod tensorflow job with multiple nodes gives a "Cannot connect to host algo-1" error and hangs indefinitely. The horovod job runs successfully if I specify a single node and multiple processes. I have been able to run horovod multi-node training with the same script outside of SageMaker.

To reproduce
Bit complex to add all the scaffolding, but I can put together an reproducible example if necessary. The gist of it is

from sagemaker.tensorflow import TensorFlow
from sagemaker.inputs import FileSystemInput

role = ...
image_name = "jarednielsen/albert-tf:sagemaker"
fsx_id = ...
hvd_instance_type = "ml.p3.16xlarge"
hvd_processes_per_host = 8
hvd_instance_count = 8
batch_size = 8

distributions = {
    "mpi": {
        "enabled": True,
        "processes_per_host": hvd_processes_per_host,
        "custom_mpi_options": "-verbose --NCCL_DEBUG=INFO -x OMPI_MCA_btl_vader_single_copy_mechanism=none",
    }
}

hyperparameters = {
    "model_size": "base",
    "batch_size": batch_size,
    "max_seq_length": 512,
    "gradient_accumulation_steps": 1,
    "learning_rate": 0.00176,
    "optimizer": "lamb",
    "fsx_prefix": "/opt/ml/input/data/training",
    "name": "sagemaker",
}

estimator_hvd = TensorFlow(
    entry_point="/path/to/blank/file.py",
    role=role,
    framework_version="2.1.0",
    py_version="py3",
    hyperparameters=hyperparameters,
    train_instance_count=hvd_instance_count,
    train_instance_type=hvd_instance_type,
    distributions=distributions,
    image_name=image_name,
    subnets=[subnet_id],
    security_group_ids=[security_group_id],
    enable_sagemaker_metrics=True,
)

fsx_input = FileSystemInput(
    file_system_id=fsx_id,
    file_system_type="FSxLustre",
    directory_path="/fsx",
    file_system_access_mode="rw",
)

estimator_hvd.fit(fsx_input)

Expected behavior
It to not hang :)

Screenshots or logs

$ $ python run_sagemaker.py
2020-03-20 00:49:47 Starting - Starting the training job...
2020-03-20 00:49:50 Starting - Launching requested ML instances.......................................
2020-03-20 00:56:58 Starting - Preparing the instances for training.........
2020-03-20 00:58:47 Downloading - Downloading input data
2020-03-20 00:58:47 Training - Downloading the training image...............
2020-03-20 01:01:21 Training - Training image download completed. Training in progress.2020-03-20 01:01:22,880 sagemaker-containers INFO     Starting MPI run as worker node.
2020-03-20 01:01:22,880 sagemaker-containers INFO     Waiting for MPI Master to create SSH daemon.
2020-03-20 01:01:23,049 sagemaker-containers INFO     Starting MPI run as worker node.
2020-03-20 01:01:23,049 sagemaker-containers INFO     Waiting for MPI Master to create SSH daemon.
2020-03-20 01:01:22,603 sagemaker-containers INFO     Starting MPI run as worker node.
2020-03-20 01:01:22,603 sagemaker-containers INFO     Waiting for MPI Master to create SSH daemon.
2020-03-20 01:01:23,981 sagemaker-containers INFO     Starting MPI run as worker node.
2020-03-20 01:01:23,982 sagemaker-containers INFO     Creating SSH daemon.
2020-03-20 01:01:23,986 sagemaker-containers INFO     Waiting for MPI workers to establish their SSH connections
2020-03-20 01:01:23,361 sagemaker-containers INFO     Starting MPI run as worker node.
2020-03-20 01:01:23,361 sagemaker-containers INFO     Waiting for MPI Master to create SSH daemon.
2020-03-20 01:01:25,144 sagemaker-containers INFO     Starting MPI run as worker node.
2020-03-20 01:01:25,144 sagemaker-containers INFO     Waiting for MPI Master to create SSH daemon.
2020-03-20 01:01:22,417 sagemaker-containers INFO     Starting MPI run as worker node.
2020-03-20 01:01:22,417 sagemaker-containers INFO     Waiting for MPI Master to create SSH daemon.
2020-03-20 01:01:24,060 sagemaker-containers INFO     Starting MPI run as worker node.
2020-03-20 01:01:24,061 sagemaker-containers INFO     Waiting for MPI Master to create SSH daemon.
2020-03-20 01:03:32,823 sagemaker-containers INFO     Cannot connect to host algo-1
2020-03-20 01:03:32,824 sagemaker-containers INFO     Connection failed with exception: 
 [Errno 110] Connection timed out
[repeated 64 times]

System information
A description of your system. Please provide:

• SageMaker Python SDK version: 1.50.14
• Framework name (eg. PyTorch) or algorithm (eg. KMeans): TensorFlow,
• Framework version: 2.1
• Python version: 3.7
• CPU or GPU: GPU
• Custom Docker image (Y/N): Yes

Additional context
The one lead I can think of is that I'm doing something a little different with the entrypoint script. Instead of specifying it in the Tensorflow() constructor, I specify it in the Dockerfile with

ENV SAGEMAKER_PROGRAM /opt/ml/input/data/training/myscript.py

This is a quirk specific to my situation, but works fine on single-node. Anything I should dive into to investigate the Horovod hanging issue?

I have the following in my Dockerfile, following the lead of https://github.com/aws/sagemaker-tensorflow-container/blob/master/docker/1.15.2/py3/Dockerfile.gpu

# Below here is necessary to install SSH on SageMaker machines
RUN apt-get update && apt-get install -y --no-install-recommends openssh-server && mkdir -p /var/run/sshd
RUN sed 's@session\s*required\s*pam_loginuid.so@session optional pam_loginuid.so@g' -i /etc/pam.d/sshd
RUN mkdir -p /root/.ssh/ && \
    ssh-keygen -q -t rsa -N '' -f /root/.ssh/id_rsa && \
    cp /root/.ssh/id_rsa.pub /root/.ssh/authorized_keys && \
    printf "Host * StrictHostKeyChecking no" >> /root/.ssh/config
# Allow OpenSSH to talk to containers without asking for confirmation
RUN cat /etc/ssh/ssh_config | grep -v StrictHostKeyChecking > /etc/ssh/ssh_config.new \
    && echo "    StrictHostKeyChecking no" >> /etc/ssh/ssh_config.new \
    && mv /etc/ssh/ssh_config.new /etc/ssh/ssh_config

Anything more I need to do?

[laurenyu]

the logs seem like it's hanging before it tries to get to your script, but it is possible that entry_point="/path/to/blank/file.py" is overriding something - given that it's working for single-node training, I'm less inclined to believe that's the issue.

Could you share your entire Dockerfile, or at least all of the parts that are MPI-related?

[jarednielsen]

Sure! The Dockerfile is

FROM nvcr.io/nvidia/tensorflow:20.02-tf2-py3

RUN apt-get update && \
    apt-get install -y iputils-ping net-tools openmpi-bin clustershell htop

# Below here is necessary to install SSH on SageMaker machines
RUN apt-get update && apt-get install -y --no-install-recommends openssh-server && mkdir -p /var/run/sshd
RUN sed 's@session\s*required\s*pam_loginuid.so@session optional pam_loginuid.so@g' -i /etc/pam.d/sshd
RUN mkdir -p /root/.ssh/ && ssh-keygen -q -t rsa -N '' -f /root/.ssh/id_rsa && \
    cp /root/.ssh/id_rsa.pub /root/.ssh/authorized_keys && \
    printf "Host * StrictHostKeyChecking no" >> /root/.ssh/config
# Allow OpenSSH to talk to containers without asking for confirmation
RUN cat /etc/ssh/ssh_config | grep -v StrictHostKeyChecking > /etc/ssh/ssh_config.new \
    && echo "    StrictHostKeyChecking no" >> /etc/ssh/ssh_config.new \
    && mv /etc/ssh/ssh_config.new /etc/ssh/ssh_config
RUN pip install mpi4py sagemaker-containers

ENTRYPOINT []
ENV SAGEMAKER_PROGRAM /opt/ml/input/data/training/my_training_script.py

One more piece of information - when I run single-node multi-process, I get a warning saying "MPI was unable to start process", but the job runs successfully in distributed mode anyway.

[jarednielsen]

I've replicated it with an even simpler Dockerfile based on DLC:

FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/tensorflow-training:2.1.0-gpu-py36-cu101-ubuntu18.04
ENTRYPOINT []
ENV SAGEMAKER_PROGRAM /opt/ml/input/data/training/run_simple_horovod.py

where run_simple_horovod.py has content

import horovod.tensorflow as hvd
import tensorflow as tf

print("Starting the program")
hvd.init()
hvd.allreduce(tf.constant(0))
print("Finished the program")

This runs on single-node distributed, but hangs on multi-node. This means it's not an issue with the custom container, but something else. Distributed training is crucial to our use case of training language models, and I'm able to run it on other platforms like ECS and EKS successfully. What's the best way forward to get SageMaker working?

[nadiaya]

Since the dockerfile you provided didn't seem to modify anything except specifying a SAGEMAKER_PROGRAM I tried to reproduce the problem by just running a training job using an official image and using provided script as an entrypoint:

    estimator = TensorFlow(
        entry_point=os.path.join(horovod_dir, "run_simple_horovod.py"),
        role="SageMakerRole",
        train_instance_count=8,
        train_instance_type="ml.p3.16xlarge",
        sagemaker_session=sagemaker_session,
        py_version="py3",
        script_mode=True,
        framework_version="2.1",
        distributions={"mpi": {
            "enabled": True,
            "processes_per_host": 8,
        }},
    )

And I was not able to reproduce the issue.

Do you get the same error when using official image and passing the user script instead of rebuilding the image?

[nadiaya]

Rerun once again with the custom mpi options set and training job still succeeds with no issues:

    estimator = TensorFlow(
        entry_point=os.path.join(horovod_dir, "run_simple_horovod.py"),
        role="SageMakerRole",
        train_instance_count=8,
        train_instance_type="ml.p3.16xlarge",
        sagemaker_session=sagemaker_session,
        py_version="py3",
        script_mode=True,
        framework_version="2.1",
        distributions={"mpi": {
            "enabled": True,
            "processes_per_host": 8,
            "custom_mpi_options": "-verbose --NCCL_DEBUG=INFO -x OMPI_MCA_btl_vader_single_copy_mechanism=none",
        }},
    )

[jarednielsen]

Turned out to be a security group permissions issue. Relaxing the ingress/egress requirements fixed the problem.

[rcruzgar]

Hi @jarednielsen

I am having the same issue, but I don't get to solve it.
Regarding your last comment, could you please tell me what permissions did you change from Sagemaker?

Cheers,
Rubén.