Support MLX on Kubernetes with Kubeflow

[andreyvelich]

MLX is a new ML framework specifically designed to run on Apple silicon: https://github.com/ml-explore/mlx

It has some differences compare to PyTorch with mps backend: ml-explore/mlx#12 (comment)

It would be nice to integrate MLX in Kubeflow ecosystem for distributed capabilities, and provide a way to run MLX models on Kubernetes.

For example, we can leverage Kubeflow Training Operator for MLX Model Training and Fine-Tuning, and Kubeflow Katib for HyperParameter optimization.
Since Kind cluster supports ARM arch, we should explore if we can use M-series GPUs for MLX model training with Kind in the future.

In addition to that, I saw examples how folks run Kubernetes on multi-VMs with MacOS machines and kubeadm.
That might be useful when a single machine can't handle very large ML model.

cc @kubeflow/wg-training-leads @awni

[gaocegege]

In addition to that, I saw examples how folks run Kubernetes on multi-VMs with MacOS machines and kubeadm.
That might be useful when a single machine can't handle very large ML model.

Interesting. Does MLX support multi-node training?

[awni]

Not yet. We are working on it. Probably makes sense to follow up on this once we have some basic support there.