Extending mpi4jax with XPU support #226

jczaja · 2024-02-01T12:59:15Z

Hi,

This PR is a draft of changes enabling Intel XPU (via SYCL) . Changes are modeled after GPU (Cuda) integration so hopefully it should be easy to understand.

Unit tests pass rate:

Currently on our setup with XPU We got:

mpirun -np 2 pytest .
============ 1 failed, 122 passed, 3 skipped, 2 warnings in 59.98s =============

which match result when We run on purely CPU (test_common.py I got failing for some reason on both XPU and CPU execution)

Usage:

New env var was added MPI4JAX_USE_SYCL_MPI by analogy to MPI4JAX_USE_CUDA_MPI. MPI4JAX_USE_SYCL_MPI=1 makes mpi4jax to assume that MPI implementation used can work with XPU buffers (GPU/XPU aware) and MPI4JAX_USE_SYCL_MPI=0 makes mpi4jax to assume that MPI implementation cannot work XPU buffers and there is a copy of XPU buffer into allocated CPU buffer (analogically like it is done for CUDA).

TODO:

Extend github action to have test for building XPU support.
Test against different implementation of MPI (OpenMPI, IntelMPI)
Write a doc
Polish code

Concerns:

Collective ops (python code) of CUDA are almost the same code as SYCL (just different import of cython module) , so It would be good to make a common code for that for the ease of maintaining.
Cython part of XPU is using C++ (see setup.py) so perhaps It wold be good to discuss what C++ features (which standard) can be used within code e.g. Currently very minimal C++ features are used .

Any suggestions are welcomed!

@wozna, @sfraczek, @bartekkuncer, @shssf

Added inital sycl infrastructure

Jczaja/sycl except

- Fix to sendrecv problem

Remove redundant export

jczaja · 2024-02-23T16:10:37Z

@dionhaefner , @PhilipVinc
This PR is ready to be reviewed now. One concern is that codecov CI is failing and I'm not sure if we can do something about as unit tests are not running on XPU neither on GPU in this github CI. So we added more code that is not covered in github CI.

Anyway , We are looking forward to your feedback.

dionhaefner

Thanks! This looks high-quality, let's fix some minor issues.

In case this breaks in the future: Who is going to maintain this code?

docs/installation.rst

mpi4jax/_src/xla_bridge/mpi_xla_bridge_xpu.pyx

setup.py

jczaja · 2024-02-29T10:37:55Z

Thanks! This looks high-quality, let's fix some minor issues.

In case this breaks in the future: Who is going to maintain this code?

@dionhaefner Thanks for reviewing this PR. This functionality will be maintained by Intel JAX GPU teams . I will send some more info via email

We started to work on suggestions you made and will update PR as soon as we have something.

Jczaja/xpu support setup

jczaja · 2024-03-05T13:58:50Z

@dionhaefner , @PhilipVinc We adressed (hopefully) all the questions and comments and updated PR with relevant changes. So we are ready for second round of review. Please review!

mpi4jax/_src/collective_ops/allgather.py

mpi4jax/_src/xla_bridge/device_descriptors.pxd

mpi4jax/_src/xla_bridge/device_descriptors.pyx

mpi4jax/_src/xla_bridge/mpi_xla_bridge_gpu.pyx

mpi4jax/_src/xla_bridge/mpi_xla_bridge_xpu.pyx

mpi4jax/experimental/notoken/collective_ops/allgather.py

setup.py

docs/installation.rst

dionhaefner · 2024-03-05T23:00:17Z

Thanks, I like the reduced code duplication. Let's get some final changes done then merge this.

dionhaefner · 2024-04-08T10:47:32Z

@jczaja Could you please address the last 3 comments so we can get this merged?

jczaja · 2024-04-11T08:28:04Z

@dionhaefner I apologize for lack of activity from our side. The reason were explained in an email I sent some time ago. Let my introduce my colleague @Zantares who will take over maintaining this work.

Zantares · 2024-04-12T07:32:20Z

@jczaja Could you please address the last 3 comments so we can get this merged?

Hi @dionhaefner , we have taken related work and will give the feedback in 1 week, thanks for your patience!

Dboyqiao · 2024-04-15T08:21:12Z

@jczaja Could you please address the last 3 comments so we can get this merged?

@dionhaefner we have addressed the left 3 comments. could you help to review？ Thank you!

dionhaefner

LGTM, thanks all!

@PhilipVinc You want to give it a final pass?

jczaja added 30 commits November 28, 2023 15:36

- Sylwek changes absorbed

8a2fb83

- processed first commit of zh.

e219040

- Added XPU module

5c63562

- more dummy methods to pass registration

61527c7

- more compilation fixes

7265dc6

- yet another fix

769127e

- another fix

d1f9e86

- fix

647f1f4

- renamed sycl to xpu

27afea4

- Initial changes to setup.py to add support to XPU module

6ddac92

- fix

bc1dd00

- mpi dependency fixed

35e38cc

- Fix to align with XLA ext changes for custom call

f90233f

- WIP

c9a745a

- does not work yet

0b1213c

- added sycl runtime file

2d05e02

- does compile

1c21cd8

- compilable version

08a8095

- Building fix for XPU

8dab550

- added logging to setup.py

45d203f

- removed dpctl dependency

419aa8a

- diagnostic

11bd7a4

- Another try

44074a8

- Fix

e5c41cf

-cleanup

d580fa1

- cleanup

7172ad2

Merge pull request #1 from jczaja/jczaja/sycl-added

395fdb8

Added inital sycl infrastructure

- Added exception handling

48243fc

- added draft of allreduce_mpi working for GPU not aware MPI

e2e82b4

Merge pull request #3 from jczaja/jczaja/sycl-except

d406e25

Jczaja/sycl except

jczaja and others added 4 commits February 22, 2024 10:07

Merge pull request #21 from jczaja/jczaja/xpu-sendrecv-fix

fb3a074

- Fix to sendrecv problem

$@sfraczek$

move build_sendrecv_descriptor to last arg

b6029bc

Remove redundant export

e1b2ac6

Merge pull request #22 from Silv3S/patch-1

8309ae2

Remove redundant export

jczaja marked this pull request as ready for review February 23, 2024 16:01

dionhaefner requested changes Feb 28, 2024

View reviewed changes

jczaja added 6 commits February 29, 2024 13:51

- First set pf changes after firest round of review

93be5d9

- Added comment explaining ignoring OpenMPI C++ bindings

6300f15

- added default path

ecc655b

- Drafto of refactoring

5d00c03

- fixes

ce74f83

- Fixes

1941b6e

$@sfraczek$ sfraczek mentioned this pull request Mar 4, 2024

Jczaja/xpu support setup jczaja/mpi4jax#23

Merged

sfraczek and others added 4 commits March 5, 2024 03:51

$@sfraczek$

add notoken

3ee7ca0

$@sfraczek$

simplify

0266992

Merge pull request #23 from jczaja/jczaja/xpu-support-setup

08647d2

Jczaja/xpu support setup

$@sfraczek$

simplify w/token

8353c2a

dionhaefner requested changes Mar 5, 2024

View reviewed changes

cleanup

072aa23

refactor function get_sycl_path

cb25ca5

dionhaefner approved these changes Apr 15, 2024

View reviewed changes

dionhaefner merged commit ffd431b into mpi4jax:master May 4, 2024

Extending mpi4jax with XPU support #226

Extending mpi4jax with XPU support #226

Uh oh!

Conversation

jczaja commented Feb 1, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Unit tests pass rate:

Usage:

TODO:

Concerns:

Uh oh!

jczaja commented Feb 23, 2024

Uh oh!

dionhaefner left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jczaja commented Feb 29, 2024

Uh oh!

jczaja commented Mar 5, 2024

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dionhaefner commented Mar 5, 2024

Uh oh!

dionhaefner commented Apr 8, 2024

Uh oh!

jczaja commented Apr 11, 2024

Uh oh!

Zantares commented Apr 12, 2024

Uh oh!

Dboyqiao commented Apr 15, 2024

Uh oh!

dionhaefner left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jczaja commented Feb 1, 2024 •

edited

Loading