[distributed] Slow xgboost distributed communicator setup when scaling past 150 nodes

[justinvyu]

I'm running distributed xgboost on a large CPU cluster with 200 m5.4xlarge (16 CPU, 64 GB memory) instances. I'm running with Ray Train's XGBoostTrainer, which just launches processes on all 200 machines and sets the necessary distributed parameters and wraps the user code with xgboost.collective.CommunicatorContext(**distributed_kwargs).

When using 200 nodes, a worker will eventually crash with a socket connection timed out error, and it will never get passed this CommunicatorContext context manager call:

  File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/train/xgboost/config.py", line 44, in collective_communication_context
    with CommunicatorContext(**_get_xgboost_args()):
  File "/home/ray/anaconda3/lib/python3.11/site-packages/xgboost/collective.py", line 279, in __enter__
    init(**self.args)
  File "/home/ray/anaconda3/lib/python3.11/site-packages/xgboost/collective.py", line 51, in init
    _check_call(_LIB.XGCommunicatorInit(make_jcargs(**args)))
  File "/home/ray/anaconda3/lib/python3.11/site-packages/xgboost/core.py", line 284, in _check_call
    raise XGBoostError(py_str(_LIB.XGBGetLastError()))
xgboost.core.XGBoostError: [02:15:49] /workspace/src/collective/result.cc:78: 
- [comm.cc:219|02:15:49]: Failed to bootstrap the communication group.
- [comm.cc:330|02:15:49]: Failed to connect to other workers.
- [socket.cc:185|02:15:49]: Failed to connect to 172.24.35.81:34313
- [socket.h:357|02:15:49]: Socket error. Connection timed out

I can avoid the error by setting the CommunicatorContext({"dmlc_retry": 100}) to a large number. However, I often see that a few workers (maybe 4 of the 200) never enter the context (the call to init never finishes).

Profile of communicator init() time

I timed how long it takes to enter the with CommunicatorContext(...): with a few different settings of num_workers. Each of the graphs shows the world rank of the worker on the x axis, and the y axis shows the time it took to enter the distributed context.

Observation: the straggler among the first 20 workers seems to be the bottleneck for the rest of workers past rank 20.

100 workers

120 workers

120 workers takes longer.

200 workers

Sometimes all workers finish initializing properly, but after a delay (not consistent how long this delay is). Here's 2 examples:

Sometimes all but a few workers finish. The stragglers are hanging there forever without any more warning messages. See the gaps in data at around rank 25.

Questions

1. Is this scale (200+ nodes) something that the xgboost team has tested before?
2. How to initialize the group more quickly?

[justinvyu]

Here's a repro script that I run on a Ray cluster with 200 m5.4xlarge instances:

import xgboost

import ray.train
from ray.train.xgboost.v2 import XGBoostTrainer


def train_fn_per_worker(config: dict):
    print("STARTED")
    import time

    for i in range(600):
        print(i)
        time.sleep(1)



if __name__ == "__main__":
    trainer = XGBoostTrainer(
        train_fn_per_worker,
        scaling_config=ray.train.ScalingConfig(
            num_workers=200, resources_per_worker={"CPU": 10}, trainer_resources={"CPU": 0}
        ),
    )
    trainer.fit()

[trivialfis]

Thank you for raising the issue. We will look into improving the comm module in the future. At the moment, it requires all to all communication network, which is quite inefficient.

[trivialfis]

Ref: #10840

[justinvyu]

@trivialfis Thanks for linking the PR! Great to know that it's on your radar.

I'm wondering if you have any benchmarks running distributed (including the cluster setup, etc.)? What's the largest scale that you've tested with?

I'm looking to set up some baselines running with Ray, and I wanted to have some reference numbers.

[trivialfis]

Most of my benchmark concerns about GPUs instead of CPU. Largest cluster I have tested is somewhere around 64 GPUs if memory serves.

When the cluster is large, the biggest bottleneck for me is imbalance between worker. In a communication group, if workers wait for each other, the wait time can accumulate instead of simply bottlenecked by the slowest one. As a result, the initialization time can be a problem when members within group has to wait for each other.

In your case, it might be over congestion. After we can overcome it, the waiting issue still needs to be resolved.

I don't have any large CPU baseline for initialization yet. As mentioned, it's difficult to profile when the wait time accumulates through the ring. When the data is well balanced and all workers can arrive at the Comm Context at the same time, the wait time is mostly negligible.

That said, we are working on it along with more reliable resilience support.