DOCSEARCHbyAlgolia
Submit Docs

Kubeflow Trainer Documentation

repository·master·Indexed 24 days ago

https://github.com/kubeflow/trainer

A Kubernetes-native platform for distributed AI training and LLM fine-tuning. It supports frameworks including PyTorch, JAX, and XGBoost through high-performance orchestration. The platform includes the kubeflow_trainer_api Python package for programmatic workload management, a Helm-based installation process for managing ClusterTrainingRuntimes, and a dedicated Data Cache component.

Tokens
132K
Snippets
267
Records
509
Agent score
79%

What's inside Kubeflow Trainer

  1. What is Kubeflow Trainer?

    master

    Kubeflow Trainer is a Kubernetes-native distributed AI platform designed for scalable large language model (LLM) fine-tuning and training. It enables multi-node, multi-GPU distributed jobs across HPC clusters by bringing MPI capabilities to Kubernetes.

    It supports a wide range of frameworks, including:

    • PyTorch
    • MLX
    • HuggingFace
    • DeepSpeed
    • JAX
    • XGBoost

    Key technical features include:

    • Distributed Data Caching: Uses Apache Arrow and Apache DataFusion for zero-copy tensor streaming directly to GPU nodes to maximize performance.
    • Single CRD Orchestration: Use a single Kubernetes Custom Resource Definition (CRD), TrainJob, to manage workloads across any supported framework.
    • Native Kubernetes Integration: Works with ecosystem tools like Kueue (topology-aware scheduling), JobSet/LeaderWorkerSet (orchestration), Volcano/YuniKorn (batch scheduling), and KAI Scheduler (GPU-aware gang scheduling).

  3. Overview of Kubeflow Training Operator V1

    master

    The Training Operator is a Kubernetes-native project designed for fine-tuning and scalable distributed training of machine learning (ML) models. It implements a centralized Kubernetes controller to orchestrate distributed training jobs across various ML frameworks.

    Key Capabilities:

    • Distributed Training: Scale model training from a single machine to large-scale distributed Kubernetes clusters.
    • Framework Support: Native support for PyTorch, TensorFlow, XGBoost, JAX, and MPI. It can also orchestrate libraries like HuggingFace, DeepSpeed, or Megatron-LM.
    • Extensibility: Can be deployed on any Kubernetes-enabled cloud and integrated with custom ML frameworks written in any language.
    • Kubernetes Integration: Supports advanced scheduling techniques like Kueue, Volcano, and YuniKorn to optimize resource costs.
    • HPC Support: Supports Message Passing Interface (MPI) via the MPIJob resource, enabling high-performance computing tasks on Kubernetes.

    Note: This information pertains to Training Operator V1. For the latest features, use the Kubeflow Trainer V2 documentation.

  4. Overview of Kubeflow Trainer

    master

    Kubeflow Trainer is a Kubernetes-native platform designed for distributed AI model training and Large Language Model (LLM) fine-tuning at scale. It provides a unified Python SDK and a single TrainJob Custom Resource Definition (CRD) that works across multiple frameworks.

    Key Capabilities:

    • Multi-Framework Support: Use a single API for PyTorch, JAX, DeepSpeed, MLX, HuggingFace, Megatron, and XGBoost.
    • Distributed Training: Automatically handles the setup of DDP, FSDP, parameter servers, and gang-scheduling across multi-node GPU clusters.
    • Hybrid Environments: Develop and test locally using Docker or Podman, then deploy the exact same TrainJob to a Kubernetes cluster without code changes.
    • LLM Fine-Tuning: Native support for LoRA, QLoRA, and full fine-tuning via TorchTune, allowing you to use HuggingFace models and dataset URIs.
    • High-Performance Data Caching: Features distributed data caching using Apache Arrow and Apache DataFusion for zero-copy tensor streaming directly to GPU nodes.
    • Extensible Runtimes: Built on a plugin architecture that allows platform teams to implement custom TrainingRuntimes for specialized scheduling, networking, or resource management.

  5. Overview of Kubeflow LLM Trainer V2

    master

    Kubeflow LLM Trainer V2 is a specialized implementation designed to simplify Large Language Model (LLM) fine-tuning on Kubernetes. It leverages torchtune as its low-level runtime to provide a modular, PyTorch-native approach to fine-tuning.

    Key features include:

    • Recipe-Config Design: Uses torchtune recipes and YAML configurations for flexible training.
    • Distributed Training: Supports distributed strategies like FSDP2.
    • Abstraction: Hides complex Kubernetes infrastructure and distributed strategy configurations (like data/model parallelism) behind a simplified Python SDK.
    • Integration: Built to work with Kubeflow Trainer V2 TrainingRuntimes and ClusterTrainingRuntime.

  6. Overview of Kubeflow Trainer V2 API

    master

    Kubeflow Trainer V2 is a redesigned API built on top of Kubernetes JobSet to support cloud-native machine learning training. It aims to reduce redundancy by leveraging Kubernetes batch workload features (like PodFailurePolicy) and provides a simplified interface for different user personas.

    Key improvements include:

    • Reduced Complexity: Moves away from low-level Pod/Service management to JobSet-based orchestration.
    • Standardized Runtimes: Uses blueprints for common training tasks (e.g., PyTorch, MPI, LLM fine-tuning).
    • Extensibility: Introduces a RuntimePatches API to allow multiple controllers (like Kueue or webhooks) to configure job metadata without conflicts.

  7. Distributed Training Frameworks in Kubeflow Trainer

    master

    Kubeflow Trainer supports distributed training for a wide variety of machine learning frameworks. You can configure workloads for:

    • PyTorch: Including FSDP (Fully Sharded Data Parallel) and DDP (Distributed Data Parallel).
    • PyTorch on AMD ROCm: For distributed training on AMD GPUs.
    • JAX: Using jax.distributed for distributed workloads.
    • JAX on TPU: Specifically for distributed training on Google Cloud TPUs.
    • DeepSpeed: Utilizing ZeRO optimization for large-scale training.
    • XGBoost: Distributed training running on Kubernetes.
    • Megatron: Using Megatron-Core with Tensor Parallelism for large transformer models.
    • MLX: Training on Apple Silicon hardware.
    • Flux: Integration for HPC (High-Performance Computing) workloads.

  8. Local Development for Kubeflow Trainer

    master

    Before deploying to a Kubernetes cluster, you can run TrainJobs locally using different backends to iterate quickly:

    • Local Execution Overview: General guidance on running jobs locally.
    • Docker Backend: Execute training jobs within Docker containers.
    • Podman Backend: Execute training jobs using Podman (a Docker alternative).
    • Process Backend: Run training jobs as local processes for the fastest possible iteration loop.

  10. PyTorch training examples available in Kubeflow Trainer

    master

    The following PyTorch model training examples are available for exploration using the Kubeflow Trainer SDK. These examples cover various tasks including image classification, NLP, speech, and audio, as well as data caching patterns.

    | Task | Model | Dataset | Notebook |
| :--- | :--- | :--- | :--- |
| Image Classification | CNN | Fashion MNIST | [mnist.ipynb](./image-classification/mnist.ipynb) |
| Question Answering | DistilBERT | SQuAD | [fine-tune-distilbert.ipynb](./question-answering/fine-tune-distilbert.ipynb) |
| Speech Recognition | Transformer | Speech Commands | [speech-recognition.ipynb](./speech-recognition/speech-recognition.ipynb) |
| Audio Classification | CNN (M5) | GTZAN | [audio-classification.ipynb](./audio-classification/audio-classification.ipynb) |
| Data Caching | [Model] | [Dataset] | [data-cache-example.ipynb](./data-cache/data-cache-example.ipynb) |

  11. Run TrainJobs locally with different backends

    master
    The Kubeflow SDK allows you to run TrainJobs on your local machine without a Kubernetes cluster. This is useful for prototyping, testing training scripts, or learning in environments where Kubernetes is unavailable. You can switch between different execution backends (Local Process, Docker, Podman, or Kubernetes) by simply changing the backend_config passed to the TrainerClient.

  12. Integrate JAX with Kubeflow Training Operator

    master

    The KEP-2145 proposal introduces support for JAX distributed training on Kubernetes via the Kubeflow Training Operator. This integration enables high-performance numerical computing and accelerated training using the Single Program, Multiple Data (SPMD) paradigm.

    Key components include:

    • A new JAXJob Custom Resource Definition (CRD).
    • An updated Training Operator controller to manage JAXJob resources.
    • An extended Training Operator Python SDK for simplified job management.
    • Support for CPU-based distributed training using the Gloo backend.