feat(runtimes): Support MLX Distributed Runtime with OpenMPI #2565

andreyvelich · 2025-03-24T01:06:41Z

Fixes: #2047

This is implementation of MLX distributed runtime in Kubeflow Trainer V2 🎉

I added example for MNIST distributed training with MLX and local minikube cluster.

Let's merge it after: #2559, since this PR includes commit from that branch.

/cc @kubeflow/wg-training-leads @gaocegege @Electronic-Waste @astefanutti @saileshd1402 @shravan-achar @akshaychitneni @franciscojavierarceo @awni @angeloskath @Blaizzy @jeiksegovia @mstei4176

review-notebook-app · 2025-03-24T01:06:47Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

google-oss-prow · 2025-03-24T01:06:49Z

@andreyvelich: GitHub didn't allow me to request PR reviews from the following users: shravan-achar, mstei4176, saileshd1402, awni, angeloskath, blaizzy, jeiksegovia.

Note that only kubeflow members and repo collaborators can review this PR, and authors cannot review their own PRs.

In response to this:

Fixes: #2047

This is implementation of MLX distributed runtime in Kubeflow Trainer V2 🎉

I added example for MNIST distributed training with MLX and local minikube cluster.

Let's merge it after: #2559, since this PR includes commit from that branch.

/cc @kubeflow/wg-training-leads @Electronic-Waste @astefanutti @saileshd1402 @shravan-achar @akshaychitneni @franciscojavierarceo @awni @angeloskath @Blaizzy @jeiksegovia @mstei4176

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository.

coveralls · 2025-03-24T01:11:01Z

Pull Request Test Coverage Report for Build 14096955415

Details

0 of 0 changed or added relevant lines in 0 files are covered.
No unchanged relevant lines lost coverage.
Overall coverage remained the same at 65.938%

Totals
Change from base Build 14092676780:	0.0%
Covered Lines:	1719
Relevant Lines:	2607

💛 - Coveralls

andreyvelich · 2025-03-24T19:53:54Z

manifests/base/runtimes/mlx_distributed.yaml

+  labels:
+    trainer.kubeflow.org/accelerator: cpu


Any thoughts on this label @Electronic-Waste @tenzen-y ?
If you want, I can remove it for now as well.

I think, it will be better if we remove this label for now until we come to an agreement on runtime labels: https://cloud-native.slack.com/archives/C0742LDFZ4K/p1741263570091899

andreyvelich · 2025-03-24T19:59:20Z

cmd/runtimes/mlx/Dockerfile

+ARG MPI_VERSION=5.0.7
+RUN wget https://download.open-mpi.org/release/open-mpi/v5.0/openmpi-${MPI_VERSION}.tar.gz
+RUN tar -xvf openmpi-${MPI_VERSION}.tar.gz && rm -rf openmpi-${MPI_VERSION}.tar.gz
+RUN cd openmpi-${MPI_VERSION} && ./configure --prefix=/usr/local && make -j"$(nproc)" install
+RUN rm -rf openmpi-${MPI_VERSION}


I am not sure why the image build takes more than 3 hours in GitHub actions.
Locally I was able to build it in 10 minutes.

Interestingly, for DeepSpeed image building OpenMPI from sources takes ~ 50 minutes: https://github.com/kubeflow/trainer/actions/runs/14005402431/job/39218624074#step:7:40444

Is it normal to execute this command for many times?

Not really, locally it runs only once.
Let me try to just run make install

Uhhh. It's slower than before. Almost 6h. What had happened?...

andreyvelich · 2025-03-26T17:28:41Z

cmd/runtimes/mlx/Dockerfile

@@ -0,0 +1,30 @@
+FROM mpioperator/base:v0.6.0 AS mpi
+FROM debian:trixie


We discussed with @tenzen-y that we will use debian image for MLX Runtime for now since we see problems with running OpenMPI in Fedora.
We will investigate later if we require to improve this runtime.

This interesting initiative, great job everyone!

But I wonder how you plan to run this on Linux since at the moment MLX only runs on Apple silicon.

Thanks @Blaizzy! As you can see this image builds mlx and mlx-data from source, also MLX already runs some of their CI tests on Linux machines: https://github.com/ml-explore/mlx/blob/main/.circleci/config.yml#L72

I believe, that might be useful for folks who want to experiment with distributed MLX and MPI inside Kubernetes environment.

Signed-off-by: Andrey Velichkevich <[email protected]>

andreyvelich · 2025-03-26T21:03:54Z

sdk/kubeflow/trainer/utils/utils.py

+            pip_index_url,
+            packages_str,
+            # For the OpenMPI, the packages must be installed for the mpiuser.
+            "--user" if runtime_entrypoint[0] == constants.MPI_ENTRYPOINT else "",


Found a bug that for OpenMPI, we should install packages for the user.

andreyvelich · 2025-03-26T22:09:05Z

/hold

Signed-off-by: Andrey Velichkevich <[email protected]>

andreyvelich · 2025-03-26T22:54:35Z

cmd/runtimes/deepspeed/Dockerfile

+# Give mpiuser permission to download packages and HF models.
+RUN chown -R mpiuser:mpiuser /home/mpiuser/.local
+RUN chown -R mpiuser:mpiuser /home/mpiuser/.cache


Finally, fixed the permission issue.
With that changes, the example works fine for me cc @tenzen-y

andreyvelich · 2025-03-26T22:55:03Z

/hold cancel

Signed-off-by: Andrey Velichkevich <[email protected]>

tenzen-y

Thank you
/lgtm
/approve

google-oss-prow · 2025-03-27T04:53:56Z

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: tenzen-y

The full list of commands accepted by this bot can be found here.

The pull request process is described here

Needs approval from an approver in each of these files:

~~OWNERS~~ [tenzen-y]

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

google-oss-prow bot requested a review from a team March 24, 2025 01:06

google-oss-prow bot added the do-not-merge/work-in-progress label Mar 24, 2025

google-oss-prow bot requested review from astefanutti and akshaychitneni March 24, 2025 01:06

google-oss-prow bot requested review from franciscojavierarceo and Electronic-Waste March 24, 2025 01:06

google-oss-prow bot added the size/XXL label Mar 24, 2025

andreyvelich mentioned this pull request Mar 24, 2025

fix(datasets): Fix contiguous buffer in load MNIST dataset ml-explore/mlx-data#86

Closed

andreyvelich force-pushed the mlx-runtime branch 2 times, most recently from bbf7d29 to a6d4810 Compare March 24, 2025 19:47

google-oss-prow bot added size/XL and removed size/XXL labels Mar 24, 2025

andreyvelich changed the title ~~[WIP] feat(runtimes): Support MLX Distributed Runtime with OpenMPI~~ feat(runtimes): Support MLX Distributed Runtime with OpenMPI Mar 24, 2025

google-oss-prow bot removed the do-not-merge/work-in-progress label Mar 24, 2025

andreyvelich commented Mar 24, 2025

View reviewed changes

andreyvelich mentioned this pull request Mar 24, 2025

Create MLX Runtime with Kubeflow Trainer #2047

Closed

andreyvelich force-pushed the mlx-runtime branch from a5da3c7 to 78d666b Compare March 25, 2025 11:38

andreyvelich commented Mar 26, 2025

View reviewed changes

andreyvelich added 6 commits March 26, 2025 20:39

feat(runtimes): Support MLX Distributed Runtime with OpenMPI

5cd5bb9

Signed-off-by: Andrey Velichkevich <[email protected]>

Remove sshd config

4468ac0

Signed-off-by: Andrey Velichkevich <[email protected]>

Modify make install for OpenMPI

1c00382

Signed-off-by: Andrey Velichkevich <[email protected]>

Use Debian for MLX Runtime

39db04d

Signed-off-by: Andrey Velichkevich <[email protected]>

Install Python package for mpiuser

93a3278

Signed-off-by: Andrey Velichkevich <[email protected]>

Remove generator

0ca635a

Signed-off-by: Andrey Velichkevich <[email protected]>

andreyvelich commented Mar 26, 2025

View reviewed changes

google-oss-prow bot added the do-not-merge/hold label Mar 26, 2025

Fix mpiuser permission

dd88100

Signed-off-by: Andrey Velichkevich <[email protected]>

andreyvelich force-pushed the mlx-runtime branch from 59292fb to dd88100 Compare March 26, 2025 22:53

andreyvelich commented Mar 26, 2025

View reviewed changes

google-oss-prow bot removed the do-not-merge/hold label Mar 26, 2025

andreyvelich added 2 commits March 26, 2025 23:46

Create .cache

7169faa

Signed-off-by: Andrey Velichkevich <[email protected]>

Update MLX example

30f6b31

Signed-off-by: Andrey Velichkevich <[email protected]>

tenzen-y reviewed Mar 27, 2025

View reviewed changes

google-oss-prow bot assigned tenzen-y Mar 27, 2025

google-oss-prow bot added the lgtm label Mar 27, 2025

google-oss-prow bot added the approved label Mar 27, 2025

google-oss-prow bot merged commit 41ca967 into kubeflow:master Mar 27, 2025
18 checks passed

andreyvelich deleted the mlx-runtime branch March 27, 2025 09:38

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat(runtimes): Support MLX Distributed Runtime with OpenMPI #2565

feat(runtimes): Support MLX Distributed Runtime with OpenMPI #2565

andreyvelich commented Mar 24, 2025 •

edited

Loading

review-notebook-app bot commented Mar 24, 2025

google-oss-prow bot commented Mar 24, 2025

coveralls commented Mar 24, 2025 •

edited

Loading

andreyvelich Mar 24, 2025

Electronic-Waste Mar 25, 2025

andreyvelich Mar 24, 2025

andreyvelich Mar 24, 2025

Electronic-Waste Mar 25, 2025

andreyvelich Mar 25, 2025

Electronic-Waste Mar 26, 2025

andreyvelich Mar 26, 2025

Blaizzy Mar 26, 2025

andreyvelich Mar 27, 2025

andreyvelich Mar 26, 2025

andreyvelich commented Mar 26, 2025

andreyvelich Mar 26, 2025

andreyvelich commented Mar 26, 2025

tenzen-y left a comment

google-oss-prow bot commented Mar 27, 2025

		@@ -0,0 +1,30 @@
		FROM mpioperator/base:v0.6.0 AS mpi
		FROM debian:trixie

feat(runtimes): Support MLX Distributed Runtime with OpenMPI #2565

feat(runtimes): Support MLX Distributed Runtime with OpenMPI #2565

Conversation

andreyvelich commented Mar 24, 2025 • edited Loading

review-notebook-app bot commented Mar 24, 2025

google-oss-prow bot commented Mar 24, 2025

coveralls commented Mar 24, 2025 • edited Loading

Pull Request Test Coverage Report for Build 14096955415

Details

💛 - Coveralls

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

andreyvelich commented Mar 26, 2025

Choose a reason for hiding this comment

andreyvelich commented Mar 26, 2025

tenzen-y left a comment

Choose a reason for hiding this comment

google-oss-prow bot commented Mar 27, 2025

andreyvelich commented Mar 24, 2025 •

edited

Loading

coveralls commented Mar 24, 2025 •

edited

Loading