[GSOC] Tracking Issue: Integrate JAX in Kubeflow Training Operator

[sandipanpanda]

Milestone: Integrate JAX in Kubeflow Training Operator

Description

This milestone tracks the progress of integrating JAX into the Kubeflow Training Operator to enable distributed training and fine-tuning jobs on Kubernetes. This involves leveraging the JAX jax.distributed.initialize API and utilizing Kubernetes JobSet API for managing job lifecycle.

Checklist

• Review JAX documentation and distributed training requirements.
• Review Kubeflow Training Operator and JobSet API documentation.
• Draft the initial design document for JAX integration. JAX Integration Enhancement Proposal #2125
• Create a new Custom Resource Definition (CRD) for JAX jobs (e.g., JaxJob).
• Update the Kubeflow Training Operator to manage JaxJob resources.
• Implement webhook validations for the JAXJob
• Implement a mechanism to initialize and manage JAX distributed training processes using jax.distributed.initialize.
• Extend the Training Operator Python SDK to simplify the creation and management of JaxJob resources.
• Update Kubeflow Training Operator documentation to include instructions for running JAX jobs.
• Provide example configurations and training scripts.
• Ensure backward compatibility with existing Kubeflow components.
• Release the updated Kubeflow Training Operator with JAX support.

Milestone Due Date

TBD

Assignees

@sandipanpanda

/area gsoc