Revamp the rabit implementation. #10112

trivialfis · 2024-03-11T19:21:30Z

motivation

With the increasing complexity of the networking module due to support for vertical and horizontal federated learning and GPU-based training for scaling, the existing rabit module is no longer sufficient. We have multiple dimensions of features to support:

data split.
federated learning.
GPU acceleration.

Among the above features, GPU acceleration and federated learning require loading optional external libraries. Lastly, we are trying to support resilience.

features

This PR replaces the original RABIT implementation with a new one, which has already been partially merged into XGBoost. The new one features:

Federated learning for both CPU and GPU.
NCCL.
More data types.
A unified interface for all the underlying implementations.
Improved timeout handling for both tracker and workers.
Exhausted tests with metrics (fixed a couple of bugs along the way).
A reusable tracker for Python and JVM packages.

todos:

JVM
Standardize the naming of worker parameters and tracker methods.
Worker sortby.

working in progress

Retry is still in progress. This is to provide essential support for handling exception (e.g., a network error or an OOM). Segfault handling has to be done with additional cooperation with the distributed framework and is out of scope for this work.

note for review

Breaking changes are made to the tracker in Python and JVM interfaces.
communicator parameters are reworked. For a list of parameters, see the document in the C header.
Federated tracker requires an additional n_workers parameter.

trivialfis · 2024-03-12T11:11:33Z

@rongou Is using rank{} as process name instead of host name a deliberate choice for federated learning?

rongou · 2024-03-12T17:33:02Z

@rongou Is using rank{} as process name instead of host name a deliberate choice for federated learning?

Yes, in a federated setting, a participant may not want to expose the host name to the rest of the group.

trivialfis · 2024-04-08T09:15:42Z

Tracking apache/arrow#41058 Need to remove the war in CI once a new snappy is published.

trivialfis · 2024-04-28T15:34:58Z

@wbo4958 Please help look into the changes to JVM packages when you are available.

wbo4958 · 2024-04-29T02:33:15Z

jvm-packages/xgboost4j/src/main/java/ml/dmlc/xgboost4j/java/RabitTracker.java

+    XGBoostJNI.checkCall(XGBoostJNI.TrackerRun(this.handle));
+    this.tracker_daemon = new Thread(() -> {
+      try {
+        XGBoostJNI.checkCall(XGBoostJNI.TrackerWaitFor(this.handle, 0));


May I ask why tracker_daemon is needed here?

I kept it as a handle for the implementation of uncaughtException, do you have suggestions for handling it differently?

trivialfis · 2024-04-29T03:42:28Z

@rongou The PR is ready for an initial review. I can't extract more small PRs since the global communicator is swapped.

Need to investigate the JVM dependency and the flaky error.

get host name. send port. all gather. assert. op. begin work on bootstrap. utilities. work on bootstrap. move listener. catch. comm. async send. block. block. Start work on async. batch poll. tests. move. Start working on tracker. better tests. work on tracker. bind. work on accepting workers. complete allgather. Move. Send with JSON. work on shutdown. msg. compare task. rename. Move. cleanup. Start work on broadcast. Work comm. Cleanup bootstrap. hide. Move to bootstrap. non blocking. cleanup. any op. shift. cleanup. Cleanup. per-thread. checks. start working on nccl. backend. Get the prototype compile. log. test print. timeout on connection. get nccl allreduce basic. look into federated. proto. scatter reduce. allreduce prototype. Work on tests. cleanup. Initialization. Init. Work on Python. get args. Start working on allgatherv. convert some allreduce. remove some old use. remove cpu impl. work on gpu. play with dlopen. convert. convert. convert. placeholder. backend. work on federated. remove. Move. Federated tracker. Move. move into comm. GPU variant. not just nccl. fix. fix. Convert. bitwise. stream. copying allgather. replace. Remove. remove. Remove device. Remove rabit. Remove rabit. cmake. tests. use gmock. Move. Split. init. Extract. compiler. test timeout. exc. comments. Tests for federated. Remove. remove. Split up. refactor tests. format. extract magic number. Extract more commands. refactor. Remove. Reduce dependency on c api. remove old code. throw. coll error. indirect. look into dask module. parameters. command. probing. listen for error. debug. host. cleanup. dask. loop. working basic. header. guard. test. type. socket. cleanup & notes. use a state machine. work on tests. header. test channel. cleanup. cleanup broadcast. unneeded changes. allgather string. Fixes. cleanup rebase. fixes after rebase. split up nccl comm. Move data copying. allgatherv test. Extract. tests. test allreduce. remove the use of ctx. tests. rebase. work on fed. work on allgatherv. name. lint. Split. split. remove gmock. move. CPU. CUDA. compile. Cleanup. header. Work on tests. checks. fixes. tests. work on CUDA test. comm. Share the implementation. tests. cleanup. cleanup. cleanup cleanup. set device. cleanup. cleanup. more. cleanup. Get it work. wait. revert dask changes. time. remove reference to encoder. extract. extract. split up the training function. Fix. deterministic. Fix. debug. Fixes. remove. cleanup. fix. Move worker env. cleanup. cleanup. wait. cleanup. extract error handling. get abort to work as well. Move. policy. cleanups. cleanup. Split up. doc. Cleanup ctor. tests. tests. tests. configuration. tests. task id. start working on metric tests. Remove. type. agg. fix seq. tests. start working on cuda test. type. fixes. tests. Use device ord. Remove auc. remove elementwise. remove multi-class cleanup aft. cleanup ranking. remove old tests. headers. Move. move. single gpu tests. Cleanup C API. unknown. C API. Small cleanup. cleanup. Fix. cleanup. work on async queue. work on sync. Use blocking op. result. Fuzzing. result. Remove coll error. Move. cleanup. cleanup. cleanup. cleanup. cleanup. cleanup. cleanup. Fix removal. test lint. remove. cleanup. cleanup. cleanup. invoke result. note. Fix rebase. Fix rebase. fix & cleanup. Fix. remove coll error for now. cleanup. replace. replace. replace. test basic. cleanup, fix. deduced size. cleanup. Convert to new routines. Fixes. Fix. Add test. use vector. cleanup. safe coll. Fix. Fix. Fix. Fix. cleanup. cleanup. Don't throw. Fixes. v6 timeout. Cleanups. Remove error handling for now. Timeout. syc. mac. cli. remove. types. federated. build. build. sortby. build. windows, macos. macos. federated. macos. lint. skip finalize. Forbid empty data. Shutdown before dtor. small allreduce. rounddown. empty input. remove warning. remove get host IP. take down the jvm package for now. stop early. annotation. lint. debug check. windows. np.bool. Work on shutdown. test blocking. remove dask error. Detach. comments. blocking. display timeout. debug github error. checks. Switch the order. revert debug log. fix tests. delete. reverse. don't block. improved error. clear; shutdown the tracker. release lock early. comments. macos. windows. Move. const. Unify ctor. remove exceptions. lint, comment. freeze pyarrow. windows. r package. Fix CI. Start looking into jvm chpk jni. c test. remove extra argument. tracker. interrupt. cleanup. Fix spark profiling. Compile. Start convert the scala package. Log init. cleanup test. communicator. alive. tests log. Revert "log." This reverts commit 3bc6d82. Shutdown when exit. remove tracker return code. windows build. shutdown only if not closed. lint. protect the listener. concat. Debug log. detect EOF. Revert "Debug log." This reverts commit a3e0bd9. Cleanup. Fixes. lint. don't omit frame pointer. Refactor tests. Fix minimum build. Fix distributed tests on single GPU. cleanup & win build. MacOS compilation. typo. macos, jvm. Windows socket. Ignore POLLHUP Handle shutdown. unix socket fix MSVC Sock error Restore the shutdown signal. states windows sock enable win tests lint. update. Skip tests. skip only for gpu. cleanup. cleanup. remove error code for now. Documents. lower case. rename. Fix. Fix.

This reverts commit 0525394.

test.

This reverts commit ee8e203.

This reverts commit 1baa02a.

This reverts commit d7540d8.

This reverts commit eb7b88e.

trivialfis · 2024-05-09T00:51:55Z

@rongou Hi, the PR is ready for review. It's quite large, if you need to have an online review please let me know.

trivialfis force-pushed the rabit-worker-get-all branch from 17d179b to 834e598 Compare March 15, 2024 10:53

trivialfis force-pushed the rabit-worker-get-all branch 2 times, most recently from 3f702cf to 0380974 Compare March 27, 2024 11:31

trivialfis mentioned this pull request Mar 27, 2024

Update collective implementation. #10152

Merged

trivialfis force-pushed the rabit-worker-get-all branch from 17fa5a8 to 85ac8f8 Compare April 1, 2024 02:51

trivialfis force-pushed the rabit-worker-get-all branch from 3001987 to 1421115 Compare April 11, 2024 12:28

trivialfis mentioned this pull request Apr 17, 2024

[coll] Improve event loop. #10199

Merged

trivialfis force-pushed the rabit-worker-get-all branch from 6ac4fe2 to 633da12 Compare April 18, 2024 02:03

trivialfis mentioned this pull request Apr 18, 2024

[coll] Add global functions. #10203

Merged

trivialfis force-pushed the rabit-worker-get-all branch from 633da12 to 3bf7a7f Compare April 18, 2024 19:29

trivialfis mentioned this pull request Apr 18, 2024

[coll] Implement shutdown for tracker and comm. #10208

Merged

trivialfis force-pushed the rabit-worker-get-all branch 5 times, most recently from 852a9f8 to 45c8fa8 Compare April 27, 2024 23:52

wbo4958 reviewed Apr 29, 2024

View reviewed changes

trivialfis force-pushed the rabit-worker-get-all branch from 931010e to ee8e203 Compare April 29, 2024 02:54

trivialfis marked this pull request as ready for review April 29, 2024 03:25

trivialfis changed the title ~~[WIP] Revamp the rabit implementation.~~ Revamp the rabit implementation. Apr 29, 2024

trivialfis mentioned this pull request May 6, 2024

[jvm-packagaes] Freeze spark to 3.4.1 for now. #10253

Merged

trivialfis added 4 commits May 9, 2024 04:54

cleanup.

abf2419

cleanup jvm packages.

9fa105d

Fix demo chunks.

3929daa

trivialfis added 18 commits May 9, 2024 04:54

consistent jvm.

ae0f49d

Revert "consistent jvm."

08ccefc

This reverts commit 0525394.

non-jvm changes.

22194a7

bisect.

bb8e40e

test.

revert test changes.

e5c35ec

revert gh changes.

81d8692

jvm changes.

99d55b3

unused import.

baf99f3

Revert "unused import."

353f2de

This reverts commit ee8e203.

Revert "Revert "unused import.""

4690758

This reverts commit 1baa02a.

update jackson.

72ef55b

try latest spark.

594bbc9

scope.

2232a6e

Revert "try latest spark."

33876f6

This reverts commit d7540d8.

Fix.

c3547ab

GPU package.

269090c

more.

eb7b88e

Revert "more."

b4e97f9

This reverts commit eb7b88e.

trivialfis force-pushed the rabit-worker-get-all branch from aef7630 to b4e97f9 Compare May 8, 2024 20:54

trivialfis added 2 commits May 9, 2024 05:23

Revert jackson version

b56ba97

Merge branch 'master' into rabit-worker-get-all

ae44c1a

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Revamp the rabit implementation. #10112

Revamp the rabit implementation. #10112

trivialfis commented Mar 11, 2024 •

edited

trivialfis commented Mar 12, 2024

rongou commented Mar 12, 2024

trivialfis commented Apr 8, 2024

trivialfis commented Apr 28, 2024

wbo4958 Apr 29, 2024

trivialfis Apr 29, 2024

trivialfis commented Apr 29, 2024 •

edited

trivialfis commented May 9, 2024

Revamp the rabit implementation. #10112

Are you sure you want to change the base?

Revamp the rabit implementation. #10112

Conversation

trivialfis commented Mar 11, 2024 • edited

motivation

features

working in progress

note for review

trivialfis commented Mar 12, 2024

rongou commented Mar 12, 2024

trivialfis commented Apr 8, 2024

trivialfis commented Apr 28, 2024

wbo4958 Apr 29, 2024

Choose a reason for hiding this comment

trivialfis Apr 29, 2024

Choose a reason for hiding this comment

trivialfis commented Apr 29, 2024 • edited

trivialfis commented May 9, 2024

trivialfis commented Mar 11, 2024 •

edited

trivialfis commented Apr 29, 2024 •

edited