From High Performance Computing To AI Workloads on Kubernetes: M... Andrey Velichkevich, & Yuki Iwai

From High Performance Computing To AI Workloads on Kubernetes: M... Andrey Velichkevich, & Yuki Iwai

Speakers: Andrey Velichkevich, Kubeflow Trainer Contributor; Yuki Iwai, Kubeflow Trainer Contributor

Conference: KubeCon EU

YouTube: https://www.youtube.com/watch?v=Fnb1a5Kaxgo

Overview

This talk, "From High Performance Computing To AI Workloads on Kubernetes," presented by Andrey Velichkevich and Yuki Iwai, addresses a critical challenge in modern machine learning operations: abstracting the inherent complexity of Kubernetes and distributed training from data scientists. The core problem highlighted is the significant "infrastructure tax" that data scientists pay when trying to scale their machine learning code, forcing them to grapple with Docker, compute configurations, data access, and Kubernetes APIs instead of focusing on model development. The talk introduces Kubeflow Trainer V2, a project designed to simplify this process, particularly focusing on its new MPI (Message Passing Interface) runtime.

The presentation demonstrates how Kubeflow Trainer V2 empowers data scientists to use native ML libraries like MLX and DeepSpeed for distributed training and fine-tuning large language models (LLMs) on Kubernetes clusters, completely abstracting away the underlying infrastructure. By providing a unified Python SDK and a clear separation of concerns between data scientists (the Train service) and DevOps engineers (the Training Runtime), Kubeflow Trainer aims to accelerate the adoption of advanced AI workloads on cloud-native platforms. This initiative is particularly pertinent in the era of generative AI, where models are becoming exponentially more complex and data volumes are immense, demanding highly scalable and efficient training environments without imposing an overwhelming learning curve on practitioners.

Background

▶ Watch: Challenges for data scientists and Kubeflow Trainer solution (0:00)

The journey of a data scientist often begins with a desire for simplicity, flexibility, and scalability in their ML workflows. They want to write code using familiar native libraries such as PyTorch, JAX, or Apple's MLX, and then effortlessly scale these workloads. However, the reality of distributed machine learning, especially on platforms like Kubernetes, presents a formidable barrier. Data scientists frequently encounter a steep learning curve involving tasks like environment setup, Docker image configuration, managing compute resources, configuring data access, understanding Kubernetes APIs, and even delving into HPC technologies or resource queuing mechanisms. This extensive overhead, termed the "infrastructure tax," diverts their focus from core ML research and development.

The problem is exacerbated by the rapid evolution of the ML ecosystem. The emergence of generative AI and large language models (LLMs) means models are larger, data sets are massive, and new frameworks and APIs are constantly appearing. This dynamic landscape creates a significant disparity between the fast-paced development of ML frameworks and the comparatively slower, more complex process of establishing robust infrastructure. Traditionally, high-performance computing (HPC) environments have been used for large-scale scientific simulations and, more recently, for initial ML training. However, HPC-dedicated infrastructures often rely on disparate job schedulers and orchestrators, leading to increased maintenance costs and fragmented user experiences across the entire ML lifecycle—from data preparation to model serving.

The speakers advocate for leveraging Kubernetes as a unified platform to manage the entire ML lifecycle. Kubernetes offers automated infrastructure management, self-healing capabilities, and a comprehensive API, which can mitigate maintenance costs. However, migrating existing ML jobs or adopting new state-of-the-art libraries to a Kubernetes-native environment often requires data scientists to adapt their training code, a task they neither desire nor are equipped for. This "gap" between the data scientist's focus on training code and the DevOps engineer's responsibility for infrastructure configuration is what Kubeflow Trainer aims to bridge, building on the foundation of earlier projects like the original training operator, which began in 2017.

Key Findings

▶ Watch: New API Runtime and MLX Distributed Training Demo (2:30)

The talk presents several key findings and contributions through the evolution of Kubeflow Trainer, particularly with the introduction of Kubeflow Trainer V2 and its MPI runtime:

• Abstraction of Kubernetes Complexity: The primary finding is the successful abstraction of Kubernetes infrastructure specifics from data scientists. Kubeflow Trainer V2 allows data scientists to write and scale native ML code (e.g., MLX, DeepSpeed) without needing to interact directly with kubectl commands, YAML manifests, or intricate cluster configurations.
• Separation of Concerns with Train and Training Runtime: A crucial architectural innovation is the clear division of responsibilities. The Train service provides a simplified Python interface for data scientists to submit jobs, while the Training Runtime acts as a blueprint for platform engineers to define and configure the underlying infrastructure, compute resources, and pre-installed packages. This enables rapid iteration for data scientists and standardized, secure environments for DevOps.
• Introduction of MPI Runtime for Distributed Communication: The talk highlights the new MPI runtime as a significant advancement. This runtime leverages Message Passing Interface (MPI), a standard for parallel computing, to orchestrate distributed communication for ML frameworks on Kubernetes. This is particularly beneficial for frameworks like MLX that inherently use MPI for scaling, and for DeepSpeed/PyTorch workloads requiring robust multi-node, multi-GPU training.
• Kubeflow SDK for Pythonic Job Submission: The development of a dedicated Kubeflow SDK provides a high-level Python interface, allowing data scientists to submit training jobs in a familiar Pythonic style. This SDK simplifies interactions with the Kubernetes cluster, replacing complex API calls with intuitive function invocations.
• Decoupling Training Code from Infrastructure Changes: A core benefit demonstrated is that data scientists' training code remains largely untouched, even when the underlying infrastructure (e.g., migrating from a traditional HPC scheduler to Kubernetes) changes. The MPI runtime, combined with the Kubeflow SDK, fills the gaps between different job schedulers, ensuring code portability and reducing adaptation overhead.
• Enhanced Debugability for Distributed Training: The speakers emphasize MPI's advantage in debugability. Tools like MPI run or MPI exec can perform arbitrary commands, enabling DevOps engineers to verify infrastructure integrity (e.g., using Nvidia NCCL test for GPU connectivity) independently of the ML training code. This allows for clear isolation and easier resolution of infrastructure-related errors, which are common in large-scale model training.
• Bridging HPC and AI Workloads: The MPI runtime facilitates a smoother transition for organizations moving from traditional HPC environments to cloud-native AI workloads on Kubernetes. By embracing MPI, Kubeflow Trainer offers a familiar paradigm for orchestrating parallel tasks, making it easier to adapt existing HPC expertise to AI training challenges.

Technical Deep Dive

▶ Watch: Simplifying MLX distributed training with Kubeflow SDK (4:00)

Kubeflow Trainer V2's architecture is meticulously designed to abstract the complexities of distributed ML training on Kubernetes. At its core, it leverages a clear separation between the data scientist's interface and the underlying infrastructure orchestration, primarily through the Kubeflow SDK and the Training Runtime concept.

For data scientists, the interaction begins by initializing a Trainer Client, which points to their Kubernetes cluster (e.g., a local Minikube instance or a larger GPU cluster). The actual training code, written in native ML frameworks like MLX or DeepSpeed, remains largely unmodified. The magic happens when the data scientist invokes a simple Python API, typically train(), provided by the Kubeflow SDK. This API call encapsulates the training function, its arguments, and crucial scaling parameters like the desired number of nodes.

The Training Runtime is a critical component configured by platform or DevOps engineers. It acts as a template or blueprint, defining the execution environment for training jobs. These runtimes specify the entry point for the training process (e.g., MPI run), the required compute resources (CPUs, GPUs), and pre-installed packages (e.g., scikit-learn, pandas, specific versions of DeepSpeed). Examples include MLX distributed runtime, DeepSpeed distributed runtime, and Torch distributed runtime. The runtime packages API allows data scientists to inspect the installed software within a given runtime, ensuring compatibility and reproducibility.

Behind the scenes, when a data scientist submits a job via the SDK, Kubeflow Trainer orchestrates an entire MPI cluster on Kubernetes. This involves complex tasks such as generating host files, setting up SSH keys for secure inter-node communication, and creating a JobSet. The JobSet is a Kubernetes custom resource that manages a collection of pods as a single unit, crucial for distributed workloads. Within the JobSet, Kubeflow Trainer typically provisions a launcher node and multiple worker nodes. The launcher node is responsible for initiating the training by executing the MPI run command, which then distributes the workload across the worker nodes. In the MPI world, each GPU is treated as a "slot," allowing for fine-grained resource allocation and parallel processing across all available devices. For instance, a job using 8 GPUs across 2 nodes would effectively utilize 8 MPI slots.

The choice of MPI as a core communication mechanism is deliberate. The speakers contrast MPI with torch.distributed, PyTorch's native distributed training module. MPI is fundamentally a launcher-based or center-oriented mechanism, where a central process initiates and manages the distributed job. This approach often provides superior debugability, as errors can be more easily traced back to specific nodes or communication failures using standard MPI tools. This is particularly valuable in large-scale model training environments where infrastructure errors (e.g., network issues, GPU problems) are common. While torch.distributed is a launcher-less or distributing-oriented approach, potentially offering better fault tolerance out-of-the-box, the speakers argue that MPI's debugability is paramount for resolving complex infrastructure issues. They also note that MPI's fault tolerance can be augmented with tools like Horovod.

Constructing an MPI environment on Kubernetes presents its own challenges. The talk discusses two primary approaches: the kubectl execution pattern and the SSH pattern. The kubectl pattern uses the Kubernetes API server for MPI initialization, simplifying setup but potentially leading to control plane performance issues due to high API server load. The SSH pattern, while requiring more intricate initial setup for SSH key management and connectivity, is generally safer for the Kubernetes control plane. Kubeflow Trainer mitigates these complexities by implementing the SSH-based MPI environment setup, abstracting it entirely from both data scientists and DevOps engineers. This ensures a robust and scalable distributed training environment without burdening users with the underlying networking and security configurations.

Kubeflow Trainer V2 supports multiple runtimes today, including Torch, DeepSpeed, and MLX, with ongoing development for JAX and TensorFlow. Its feature set includes multi-node training, fine-tuning, elastic training, and gang scheduling, which ensures that all resources required for a distributed job are available simultaneously before the job starts, preventing deadlocks and improving resource utilization. This comprehensive approach allows Kubeflow Trainer to connect various ML libraries seamlessly on top of Kubernetes, providing a powerful, unified platform for diverse AI workloads.

Demo / Proof of Concept

▶ Watch: Observing MLX distributed training progress and scalability (5:00)

The speakers presented two compelling live demonstrations to showcase the capabilities of Kubeflow Trainer V2 and its MPI runtime, highlighting how data scientists can scale complex ML workloads without Kubernetes expertise.

Demo 1: Distributed Training with MLX

The first demonstration focused on scaling a CNN model for image classification using MLX, Apple's machine learning framework designed for Apple Silicon. MLX inherently leverages MPI for distributed communication, making it an ideal candidate for showcasing Kubeflow Trainer's MPI runtime.

1. Environment Setup: The demo ran on a local Minikube instance, simulating a small Kubernetes cluster with three nodes, each configured with two CPU devices.
2. Data Scientist Workflow: The presenter assumed the role of a data scientist. The first step involved initializing the trainer client and pointing it to the local Minikube platform.
3. Native MLX Code: A key highlight was the use of entirely native MLX code. This included defining the model (a simple MLP for image classification), handling data distribution (MNIST dataset split across multiple partitions), and specifying the loss function and training loop. Crucially, the code incorporated MLX's native APIs for distributed communication, such as world_size and rank, without any Kubernetes-specific adaptations. The model was exported to disk upon completion.
4. Job Submission via Kubeflow SDK: To scale this MLX code, the data scientist simply retrieved the MLX distributed runtime (which uses MPI run as its entry point). Then, using the Kubeflow SDK, a single train() API call was made. This call passed the MLX training function, arguments (like the model path), the desired number of nodes (three), and any additional Python packages required (e.g., scikit-learn).
5. Execution and Monitoring: The system generated a job ID. The presenter demonstrated listing active jobs and viewing job steps, showing one launcher node and two worker nodes, each utilizing two CPUs. Logs revealed the distributed training in progress, with 60,000 samples distributed across the three nodes (20,000 images per worker), performing DDP (Distributed Data Parallel). The training completed, and the saved model was then used for evaluations, demonstrating its classification accuracy.
6. Key Insight: The demo effectively illustrated that native MLX code could be scaled across a Kubernetes cluster by merely specifying the number of nodes, with all Kubernetes and MPI orchestration hidden by the Kubeflow SDK.

Demo 2: Fine-tuning T5 LLM with DeepSpeed on GPUs

The second, more advanced demonstration involved fine-tuning a T5 transformer model for text summarization using DeepSpeed, a PyTorch-based optimization framework for distributed training, on a multi-GPU cluster.

1. Environment Setup: This demo utilized a GPU cluster with eight Nvidia V100 GPUs distributed across two nodes, with four GPUs per node.
2. Data Scientist Workflow: Similar to the first demo, the trainer client was initialized, this time pointing to the GPU cluster.
3. Native DeepSpeed and PyTorch Code: The training function involved defining a distributed environment, loading the WikiHow dataset, tokenizing it, downloading the T5 model and tokenizer, and configuring DeepSpeed parameters (optimizer, schedulers, microbatches) using DeepSpeed's native APIs. The model checkpoints were saved to an S3 bucket.
4. Runtime and Package Inspection: The presenter showed how to retrieve the DeepSpeed distributed runtime, which explicitly indicated the availability of four GPUs per node. A useful feature, runtime_packages, was demonstrated (though briefly glitching in the live demo, a pre-recorded snippet confirmed its functionality), allowing data scientists to inspect the installed packages and their versions (e.g., DeepSpeed 4.1, Python 3.8) within the runtime. This is crucial for reproducibility and debugging.
5. Job Submission: The train() API was invoked again, passing the DeepSpeed training function, the desired two nodes, and arguments such as the dataset URL and S3 bucket name.
6. Execution and Monitoring: The system returned a job ID. The job's nodes were listed, confirming two nodes each with four GPUs. A Grafana dashboard, leveraging DCGM exporter, was presented to visualize GPU utilization, showing spikes as the training progressed. Logs confirmed multi-node, multi-GPU training, processing 160 samples per worker across eight devices.
7. Results and Evaluation: The fine-tuning completed in approximately 69 seconds. The presenter then demonstrated loading the fine-tuned T5 model into a Hugging Face transformer pipeline and using it to summarize text from Kubeflow documentation, showcasing its practical application.
8. Key Insight: This demo powerfully demonstrated the ability to fine-tune large language models on multi-node, multi-GPU clusters using advanced frameworks like DeepSpeed, all orchestrated seamlessly by Kubeflow Trainer without direct Kubernetes interaction. The unified train() API remained consistent across different frameworks and hardware configurations.

Both demonstrations underscored the core message: Kubeflow Trainer V2 successfully abstracts the complexities of Kubernetes and distributed computing, enabling data scientists to focus purely on their ML code and achieve rapid iteration and scalable model development.

Defensive Implications

▶ Watch: Next Demo: Fine-tuning T5 model with DeepSpeed (6:50)

The advancements presented by Kubeflow Trainer V2 and its MPI runtime offer significant defensive implications for both platform/DevOps engineers and data scientists operating in MLOps environments. These implications revolve around streamlining operations, enhancing security, and optimizing resource utilization.

For Platform/DevOps Engineers:

• Standardized Training Environments: Kubeflow Trainer's Training Runtime concept allows engineers to define standardized, pre-configured environments. This ensures consistency across training jobs, reduces configuration drift, and simplifies troubleshooting. By controlling the runtime, engineers can enforce specific package versions, security policies, and resource allocations.
• Robust Distributed Training Setup: The project abstracts the complex setup of MPI environments, including SSH key management and host file generation. This is a crucial defensive measure as manual MPI setup is error-prone and can introduce security vulnerabilities if not handled carefully. Kubeflow Trainer's automated SSH-based pattern provides a more secure and stable foundation for multi-node communication compared to relying solely on kubectl for MPI initialization, which can strain the Kubernetes control plane.
• Resource Management and Monitoring: The integration with monitoring tools like Grafana via DCGM exporter provides critical visibility into GPU utilization and overall cluster health. This allows engineers to proactively identify bottlenecks, optimize resource allocation, and debug infrastructure-related issues effectively, preventing costly resource wastage and job failures.
• Decoupling Infrastructure from ML Code: By providing a clear boundary between infrastructure and ML code, DevOps teams can manage and update the underlying Kubernetes infrastructure without forcing data scientists to rewrite their training scripts. This reduces friction, accelerates infrastructure upgrades, and minimizes potential regressions in ML workflows.
• Security Best Practices: When configuring Training Runtimes, engineers can embed security best practices, such as specifying base images from trusted registries, enforcing network policies for inter-pod communication, and integrating with secrets management for sensitive credentials (e.g., S3 access keys for model checkpoints).
• Scalability and Elasticity: Leveraging Kubernetes JobSet for orchestrating distributed jobs, combined with features like gang scheduling, ensures that training resources are allocated efficiently and jobs can scale elastically without manual intervention, preventing resource starvation or over-provisioning.

For Data Scientists:

• Focus on Core ML Development: The most significant defensive implication for data scientists is the ability to entirely abstract away Kubernetes complexity. This allows them to focus their energy on model architecture, data preprocessing, and hyperparameter tuning, rather than grappling with infrastructure configuration, which is not their core expertise.
• Framework Agnosticism (at the user level): Data scientists can use their preferred native ML frameworks (MLX, DeepSpeed, PyTorch, JAX, TensorFlow) without needing to learn Kubernetes-specific adaptions for each. This promotes rapid experimentation and adoption of new state-of-the-art libraries.
• Reproducibility: The Training Runtime concept, along with the ability to query runtime_packages, helps data scientists understand and reproduce their environments. This is vital for debugging, auditing, and ensuring that models can be reliably retrained.
• Simplified Distributed Training: The single train() API call from the Kubeflow SDK simplifies the daunting task of setting up and running distributed training jobs. Data scientists can easily scale their workloads by merely adjusting the number of nodes or GPUs, democratizing access to high-performance computing for AI.
• Reduced Error Surface: By not directly interacting with Kubernetes YAMLs or kubectl commands, data scientists are less likely to introduce configuration errors that could lead to job failures or security vulnerabilities.

In summary, Kubeflow Trainer V2 provides a robust framework that defends against the "infrastructure tax" on data scientists, streamlines MLOps for engineers, and fosters a more secure, scalable, and efficient environment for developing and deploying AI models on Kubernetes.

Key Takeaways

• Kubeflow Trainer V2 abstracts Kubernetes complexity, enabling data scientists to focus on native ML code (e.g., MLX, DeepSpeed, PyTorch) without needing kubectl or YAML expertise.
• The new MPI runtime within Kubeflow Trainer simplifies distributed training for frameworks like MLX and DeepSpeed on Kubernetes, providing a robust communication layer.
• A user-friendly Python SDK offers a consistent train() API, allowing data scientists to submit and scale jobs in a Pythonic style, completely hiding underlying infrastructure details.
• Kubeflow Trainer effectively bridges the gap between HPC and AI workloads, offering a unified platform on Kubernetes that accommodates the rapid evolution of ML frameworks while providing stable infrastructure.
• The architecture introduces a clear separation of concerns with the Train service for data scientists and Training Runtime for DevOps, facilitating rapid iteration and standardized, scalable ML environments.
• MPI's strong debugability is highlighted as a key advantage for large-scale model training, allowing engineers to diagnose infrastructure errors independently of ML code, which is crucial in complex multi-node, multi-GPU setups.

About the Speaker(s)

Andrey Velichkevich is a key contributor to the Kubeflow Trainer project. His work focuses on abstracting the complexities of Kubernetes and distributed computing environments from data scientists, enabling them to efficiently scale their machine learning and AI workloads. He has been involved with the project since its earlier iteration, the training operator, which began in 2017, demonstrating a long-standing commitment to improving the MLOps experience on Kubernetes.

Yuki Iwai contributes to the Kubeflow Trainer initiative with a specific focus on the transition from traditional High-Performance Computing (HPC) environments to cloud-native AI workloads on Kubernetes. His expertise lies in understanding the challenges faced by both data scientists and DevOps engineers during this migration and developing solutions, such as the MPI runtime, that bridge the gap between rapidly evolving ML frameworks and the underlying infrastructure. His insights are crucial for ensuring that Kubeflow Trainer addresses real-world pain points in large-scale machine learning deployments.