feat(docs): KEP-2442-Support JAX Training Runtime #2643

mahdikhashan · 2025-05-10T13:22:33Z

What this PR does / why we need it:

Which issue(s) this PR fixes (optional, in Fixes #<issue number>, #<issue number>, ... format, will close the issue(s) when PR gets merged):
part of #2442

Checklist:

Docs included if any changes are user facing

coveralls · 2025-05-10T13:26:58Z

Pull Request Test Coverage Report for Build 18256788359

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 55.174%

Totals
Change from base Build 18223898761:	0.0%
Covered Lines:	1093
Relevant Lines:	1981

💛 - Coveralls

mahdikhashan · 2025-06-20T17:55:32Z

@Electronic-Waste @andreyvelich please take a look when you have time.

mahdikhashan · 2025-06-20T17:57:48Z

regarding data loading I mentioned in our sync meetings, after going over jax document, as you said, it will be handled within objective function by the user, so its sharding of data in memory and no need for further mechanism to handle permanent data (PVC, PV). however user can take care of it anytime with a patch.

cc @Electronic-Waste

Electronic-Waste

@mahdikhashan Thanks for this great work. I've left my initial comments for you.

PTAL when you have time.

/cc @kubeflow/wg-training-leads @astefanutti @franciscojavierarceo

docs/proposals/2442-jax-runtime-trainer-v2/README.md

Electronic-Waste · 2025-06-26T15:07:12Z

@mahdikhashan Hey, do you get a chance to address the comments above?

mahdikhashan · 2025-06-26T15:10:39Z

@mahdikhashan Hey, do you get a chance to address the comments above?

I'll address them by tomorrow before our sync meeting, hope its fine with you.

Electronic-Waste

Thanks for the update @mahdikhashan! I left some comments for you:)

/cc @kubeflow/wg-training-leads @astefanutti @sandipanpanda @franciscojavierarceo

docs/proposals/2442-jax-runtime-trainer-v2/README.md

Electronic-Waste · 2025-07-01T16:58:13Z

docs/proposals/2442-jax-runtime-trainer-v2/README.md

+### Communication Backend
+
+#### OpenMPI
+
+**Pros:**
+
+* Compatible with existing MPI runtime in Kubeflow Trainer v2, making deployment easier.
+
+**Cons:**
+
+* Typically requires more complex environment setup compared to simpler backends like Gloo.
+
+#### Gloo
+
+**Pros:**
+
+* Lightweight and simple to use.
+
+**Cons:**
+
+* Significantly slower than OpenMPI (10–20×) for distributed JAX training on CPUs and GPUs.
+* Less optimized for multi-node scaling and lacks native support for high-speed interconnects like InfiniBand.


I remember that @sandipanpanda implemented the Jax Runtime with Gloo backend in the last year:)

Which one do you prefer? Any suggestion? @kubeflow/wg-training-leads @astefanutti @sandipanpanda

asking me, i'd say, open-mpi has upper hand, and lack of gpu support in gloo is a significant design decision maker, at least for me. i'd appreciate @sandipanpanda feedback on it.

Maybe in the future we can use the gpu ci runner: #2674. But if we decide to use Gloo, we can test it on cpu first like what we did in the last year:)

other points could be raise from @andreyvelich , since deepspeed and mlx are based on openmpi at the moment.

asking me, i'd say, open-mpi has upper hand, and lack of gpu support in gloo is a significant design decision maker, at least for me. i'd appreciate @sandipanpanda feedback on it.

The main motive behind using gloo was it was simple to set up and use, making development iterations faster, in long term open-mpi provides more value

asking me, i'd say, open-mpi has upper hand, and lack of gpu support in gloo is a significant design decision maker, at least for me. i'd appreciate @sandipanpanda feedback on it.

The main motive behind using gloo was it was simple to set up and use, making development iterations faster, in long term open-mpi provides more value

thanks for your reply

Electronic-Waste · 2025-07-01T17:01:03Z

docs/proposals/2442-jax-runtime-trainer-v2/README.md

+The number of JAX processes can be defined using the `numNodes` field within the `mlPolicy` section of the `ClusterTrainingRuntime` configuration. This allows for specifying how many JAX processes/controller run in the distributed setup.
+
+
+![user-roles](./drawing.drawio.svg)


Could you also describe the whole procedure step-by-step with words? I think, it will be more clear to readers if the graph is companied with descriptive words:)

Electronic-Waste · 2025-07-01T17:01:22Z

docs/proposals/2442-jax-runtime-trainer-v2/README.md

We also need test plans section.

you are right. i'll add it.

Signed-off-by: Mahdi Khashan <mahdikhashan1@gmail.com>

Signed-off-by: mahdikhashan <mahdikhashan1@gmail.com>

Electronic-Waste

@mahdikhashan Great work! Thanks for your instant response. Just one more suggestion. We're close to the end.

Electronic-Waste · 2025-10-05T09:09:14Z

docs/proposals/2442-jax-runtime-trainer-v2/README.md

+### Non-Goals
+
+- No TPU testing, tests will use CPU
+- No GPU testing, tests will use CPU


Now, we've had access to GPU test cluster. So it might not be accurate now. How about saying non-goals:

Set up test in TPU cluster

Add advanced JAX APIs

WDYT? @andreyvelich @mahdikhashan

you are right.

Signed-off-by: mahdikhashan <mahdikhashan1@gmail.com>

Electronic-Waste

@mahdikhashan Thanks for your great work! Let's move forward with this KEP!

/approve
/lgtm
/assign @andreyvelich @tenzen-y @astefanutti
/hold for other maintainers

andreyvelich

Thanks for this @mahdikhashan!
/lgtm
/approve
/hold cancel

google-oss-prow · 2025-10-05T13:01:27Z

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: andreyvelich, Electronic-Waste

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~~ [Electronic-Waste,andreyvelich]

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 review from Electronic-Waste and kuizhiqing May 10, 2025 13:22

google-oss-prow bot added the size/XS label May 10, 2025

google-oss-prow bot added size/L and removed size/XS labels May 10, 2025

mahdikhashan changed the title ~~[proposal] GSoC Project 8: Jax Runtime for V2~~ [proposal] GSoC Project 8: JAX Runtime for V2 May 10, 2025

mahdikhashan changed the title ~~[proposal] GSoC Project 8: JAX Runtime for V2~~ (draft)[proposal] GSoC Project 8: JAX Runtime for V2 May 28, 2025

mahdikhashan force-pushed the gsoc-2442-jax-runtime-proposal branch from c39fc67 to 271a5d1 Compare May 28, 2025 16:53

mahdikhashan marked this pull request as draft May 28, 2025 17:01

google-oss-prow bot added do-not-merge/work-in-progress size/M and removed size/L labels May 28, 2025

google-oss-prow bot added size/L and removed size/M labels Jun 20, 2025

mahdikhashan changed the title ~~(draft)[proposal] GSoC Project 8: JAX Runtime for V2~~ [proposal] GSoC Project 8: JAX Runtime for V2 Jun 20, 2025

mahdikhashan marked this pull request as ready for review June 20, 2025 17:51

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

google-oss-prow bot requested review from astefanutti and jinchihe June 20, 2025 17:51

mahdikhashan force-pushed the gsoc-2442-jax-runtime-proposal branch from b137639 to cadaf06 Compare June 20, 2025 17:53

Electronic-Waste reviewed Jun 23, 2025

View reviewed changes

google-oss-prow bot requested review from a team and franciscojavierarceo June 23, 2025 10:14

mahdikhashan force-pushed the gsoc-2442-jax-runtime-proposal branch from c92f96c to 860a8a7 Compare June 27, 2025 07:49

Electronic-Waste reviewed Jul 1, 2025

View reviewed changes

mahdikhashan added 9 commits October 4, 2025 15:54

update future work

096468a

Signed-off-by: Mahdi Khashan <mahdikhashan1@gmail.com>

update test plan

f002edc

Signed-off-by: Mahdi Khashan <mahdikhashan1@gmail.com>

update workflow description

ae58d22

Signed-off-by: Mahdi Khashan <mahdikhashan1@gmail.com>

improve design detail section

55eeba0

Signed-off-by: Mahdi Khashan <mahdikhashan1@gmail.com>

change system to Trainer Controller Manager

1b0dd1f

Signed-off-by: Mahdi Khashan <mahdikhashan1@gmail.com>

merge runtime and trainer controller manager

effee0a

Signed-off-by: Mahdi Khashan <mahdikhashan1@gmail.com>

change plugin details

57044b8

Signed-off-by: Mahdi Khashan <mahdikhashan1@gmail.com>

update goals

f9718fb

Signed-off-by: Mahdi Khashan <mahdikhashan1@gmail.com>

use kwargs

aecb6d2

Signed-off-by: Mahdi Khashan <mahdikhashan1@gmail.com>

mahdikhashan reopened this Oct 4, 2025

google-oss-prow bot added size/L and removed size/XS labels Oct 4, 2025

mahdikhashan added 2 commits October 4, 2025 15:59

fix ci

40eb886

Signed-off-by: mahdikhashan <mahdikhashan1@gmail.com>

fix ci

c014622

Signed-off-by: mahdikhashan <mahdikhashan1@gmail.com>

mahdikhashan requested review from Electronic-Waste and andreyvelich October 4, 2025 14:22

Electronic-Waste reviewed Oct 5, 2025

View reviewed changes

update non-goals

5806896

Signed-off-by: mahdikhashan <mahdikhashan1@gmail.com>

Electronic-Waste approved these changes Oct 5, 2025

View reviewed changes

google-oss-prow bot assigned andreyvelich and astefanutti Oct 5, 2025

google-oss-prow bot added the do-not-merge/hold label Oct 5, 2025

google-oss-prow bot assigned tenzen-y and Electronic-Waste Oct 5, 2025

google-oss-prow bot added lgtm approved labels Oct 5, 2025

andreyvelich reviewed Oct 5, 2025

View reviewed changes

google-oss-prow bot removed the do-not-merge/hold label Oct 5, 2025

google-oss-prow bot merged commit 92c3f9b into kubeflow:master Oct 5, 2025
33 checks passed

		The number of JAX processes can be defined using the `numNodes` field within the `mlPolicy` section of the `ClusterTrainingRuntime` configuration. This allows for specifying how many JAX processes/controller run in the distributed setup.


		![user-roles](./drawing.drawio.svg)

feat(docs): KEP-2442-Support JAX Training Runtime #2643

feat(docs): KEP-2442-Support JAX Training Runtime #2643

Conversation

mahdikhashan commented May 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

coveralls commented May 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pull Request Test Coverage Report for Build 18256788359

Details

💛 - Coveralls

Uh oh!

mahdikhashan commented Jun 20, 2025

Uh oh!

mahdikhashan commented Jun 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Electronic-Waste 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!

Uh oh!

Electronic-Waste commented Jun 26, 2025

Uh oh!

mahdikhashan commented Jun 26, 2025

Uh oh!

Electronic-Waste left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Electronic-Waste left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Electronic-Waste left a comment

Choose a reason for hiding this comment

Uh oh!

andreyvelich left a comment

Choose a reason for hiding this comment

Uh oh!

google-oss-prow bot commented Oct 5, 2025

Uh oh!

mahdikhashan commented May 10, 2025 •

edited

Loading

coveralls commented May 10, 2025 •

edited

Loading

mahdikhashan commented Jun 20, 2025 •

edited

Loading