mirror of https://github.com/hanzo-ml/trainer.git synced 2026-09-24 00:31:34 +00:00

Files

T

Andrey VelichkevichandGitHub c4f4e431c4 chore(docs): Move KEPs to the proposals directory (#3629 )

* chore(docs): Move KEPs to the proposals directory

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

* Add KEPs from communtiy repo

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

* Update AGENTS.md

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

---------

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

2026-06-22 19:44:47 +00:00

12 KiB

Raw Permalink Blame History

KEP-2432: GPU Testing for LLM Blueprints

Project Page: GPU Testing for LLM Blueprints

Issue: https://github.com/kubeflow/trainer/issues/2432

Mentors: @andreyvelich, @varodrig

Project Size: 350 hrs

Summary

This project aims to use self-hosted runners to run GPU-intensive tasks like LLM blueprint or (planned) AI Playground. The necessary infra is provided by Oracle, plan is to use Oracle Kubernetes Engine (OKE) with NVIDIA GPUs for this task. Any code or sample that requires GPU-intensive resources will be transferred to OKE infra instead of generic GitHub infra for faster and more efficient execution.

For gpu runner testing, main idea to verify the below example by running it on GPU infra —

TorchTune: master/examples/torchtune/llama3_2/alpaca-trainjob-yaml.ipynb

The workflow will monitor the label on PRs and, upon approval from maintainer, execute the notebooks using gpu enabled self-hosted runner on the OKE infrastructure. The GitHub Action will be triggered only when a specific label (e.g., ok-to-test-gpu-runner) is applied to a pull request. For security, this workflow requires the label to be added from a maintainer before execution.

A monitoring dashboard will be established to track metrics, resource usage, and identify bottlenecks. While this setup is designed for OKE, the approach is platform-agnostic and can be adapted to any Kubernetes cluster with adequate GPU resources.

Motivation

Kubeflow Trainer is a core component of the Kubeflow ecosystem, responsible for managing and executing distributed AI/ML training jobs. With the growing adoption of Large Language Models (LLMs), reliable GPU-based training workflows have become essential.

Currently, the CI infrastructure for Kubeflow Trainer does not include automated GPU test coverage for LLM training blueprints. This leads to several limitations:

Limited GPU workload validation – GPU-specific regressions in frameworks such as PyTorch may only be detected after deployment.
Gaps in LLM-scale testing – CPU-only environments cannot replicate the performance, dependencies, and runtime behavior of GPU workloads.

Introducing self-hosted GPU runners into the CI pipeline will address these gaps by enabling end-to-end GPU testing for LLM workloads. This enhancement will deliver multiple benefits:

Validated GPU workflows – Ensure that LLM training scenarios run reliably in production-like conditions.
Live showcase capability – Demonstrate the full Kubeflow stack, including GPU-enabled LLM training, at events like KubeCon.
Improved contributor efficiency – Accelerate development with automated GPU-specific CI feedback.

By closing this testing gap and enabling impactful live demonstrations, this project will improve Kubeflow Trainer’s technical quality, accelerate development, and enhance its visibility within the AI/ML community.

Goals

Use sample LLM Blueprint
Configure GPU nodes on OKE
Establish ACR on the OKE Cluster for deploying the LLM
Create a GitHub Action for manual triggers and runners on the OKE cluster
Implement metrics and analytics for the GPU Cluster

Non-Goals

The GPU cluster for production deployment should be provided by Oracle. For testing purposes, I have a sufficiently powerful personal machine (Ryzen 7 8600G, 32GB RAM, Nvidia RTX 4060) to conduct tests.
Once the infrastructure for the self-runner is set up, running the AI Playground will require minimal setup. The primary focus of this project is to establish the infrastructure for running the LLM blueprint on OKE. The AI Playground is a secondary priority for this GSoC project, but I will continue working on it if it is not completed within the GSoC period.

Proposal

TechStack

GitHub Actions (and ARC), Kubernetes, Oracle Cloud, PyTorch, Python, Linux

User Stories

Run GPU intensive tasks on self hosted gpu infra (here OKE) instead of default CPU based infra
[Out of scope for this KEP] Run AI playground during KubeCon or other events leveraging full potential of KubeFlow Ecosystem

The idea is to automate and setup sample models where user can just Open Kubeflow Jupyter Notebook -> select Kubeflow LLM blueprint -> fine-tune model with Kubeflow Trainer -> serve it with Kubeflow KServe. This will help us to show full potential of KubeFlow ecosystem.

Design Details

Use sample LLM Blueprint

To use the same LLM blueprint that can be triggered based on admin approval. We have already 2 samples on in trainer repo,DeepSpeed and TorchTune. I have tested a sample project for running on my local system. For our usecase, we are targeting 2 samples which requires GPU.

These samples would be used to cover a range of scenarios and configurations, thereby enhancing the versatility and applicability of the LLM blueprint. This approach will not only facilitate thorough testing but also provide valuable insights into optimizing the deployment and execution of LLMs on the OKE infrastructure.

Github Action

Create a GitHub action GPU E2E Test which add functionality to run and validate the changes on the gpu based self runner. Once the maintainer adds the (`ok-to-test-gpu-runner) label, the torchtune example is executed in the gpu enabled self-runner. Assuming it takes some time and resources, we will implement queuing so that resources don't get flooded with requests. We will maintain a queue for requests, and report the result back to CI accordingly.

Here is the branch which demonstrats the running of self runner on gpu enabled infra -test-on-oci-vm

Setup and access control of OKE Cluster with GPU

In this milestone, the aim is to setup an OKE Cluster with GPU node. System: Ubuntu 22.04 LTS The GPU image has the GPU drivers pre-installed.

Cluster Architecture: The OKE cluster will be configured with:

• GPU Node Pool utilizing NVIDIA GPUs with pre-installed drivers

• Standard Node Pool for regular workloads (Metrics and Queue)

• Managed Control Plane by OCI

Accessing Cluster

OKE can be accessed with kubeconfig file
Or via Bastian

Access control for the OKE cluster

Cluster Administrators: Full cluster management rights
Maintainers: Limited administrative access
CI/CD Systems: Restricted service account access

Images for NVIDIA shapes

Reference

https://github.com/oracle-quickstart/oci-hpc-oke

https://blogs.oracle.com/java/post/create-k8s-clusters-and-deply-to-oci-from-vscode

Setup GitHub Actions Runner Controller (ARC)

Actions Runner Controller (ARC) is a Kubernetes operator that orchestrates and scales self-hosted runners for GitHub Actions. This is advanced phase of our project where we use k8s operator that is useful to scale and orchestrate pods based on the action CI.

OKE Monitoring

For admins, we also need to maintain monitoring to see the metrics and resource utilisation of the OKE infra. Oracle already provides an open-source sample for OKE Monitoring, so we can leverage that. Out of various options, installation via Helm is sufficient for our basic needs.

Estimated monthly cost: $0/month

Metrics needed


Avg CPU Usage	Avg Queue timing
Avg GPU Usage	Avg Build timing
Peak GPU Usage

Test Plan

During the GSoC period, until OKE infra is donated to KubeFlow. I will be testing local machine with 32GB RAM, Nvidia RTX 4060 GPU, Ryzen 7 8700G. I will have a demo during mid term evaluation in community call, once that is finalized by mentors.

I will be using OKE to deploy after production. To make sure there is no unnecessary usage of infra while testing, i will be putting certain guardrails on prod OKE.