Training infrastructure

51 systems in the training infrastructure category of the AI Agent Infrastructure Landscape, grouped by maturity tier.

Tier 1 — battle-tested (22)

  • Accelerate (Hugging Face) Distributed training abstraction

    Hugging Face's abstraction over distributed-training backends (DDP, FSDP, DeepSpeed, Megatron). Minimal code change to scale a PyTorch script.

  • AWS SageMaker AWS ML platform

    AWS's managed ML platform — Studio, Training, HyperPod (large-scale clusters), Endpoints. Industry-standard for enterprise AWS-hosted training.

  • Azure Machine Learning Azure ML platform

    Microsoft Azure's managed ML platform — pipelines, AutoML, Designer, prompt flow (LLM). Heavy integration with Azure OpenAI and Azure AI Studio.

  • Dagster Data + ML pipeline orchestrator

    Data orchestrator with first-class ML pipeline support — software-defined assets model; competes with Airflow/Prefect for ML pipelines.

  • DeepSpeed ZeRO-based distributed training

    Microsoft's distributed-training library — ZeRO optimizer states/grads/params partitioning; ZeRO-Inference for serving. Foundation under many OSS LLM training stacks.

  • DVC Data version control

    Open-source data and ML model version control — git-like commands for datasets/models stored in cloud blob storage. Foundational tool in MLOps stack.

  • FSDP (PyTorch) Fully Sharded Data Parallel

    Native PyTorch fully-sharded data parallel implementation — Meta's response to DeepSpeed ZeRO; FSDP2 in 2024 with improved API.

  • Google Vertex AI Training GCP ML training

    Google Cloud's managed training service — TPU and GPU clusters, hyperparameter tuning, model registry. Part of Vertex AI umbrella.

  • Hugging Face Datasets Dataset library + Hub

    Library for accessing and processing 100k+ public datasets on the Hub. Apache Arrow-backed; default dataset layer for the open ML ecosystem.

  • Label Studio (HumanSignal) Data labelling platform

    Open-source data-labelling platform supporting text, image, audio, video, time-series. Commercial Enterprise edition + Cloud (HumanSignal).

  • Labelbox Data factory for AI

    Enterprise data-labelling platform — annotation tools + 'Frontier' offering (expert RLHF labelers + Boost compute). Pivoting toward frontier-lab data services.

  • Megatron-LM Tensor/pipeline parallel LLM training

    NVIDIA's reference framework for training large transformer models with tensor, pipeline, sequence, and expert parallelism. Underpins many frontier-lab training stacks.

  • MLflow Open ML lifecycle platform

    Databricks-incubated open platform for ML lifecycle — tracking, projects, models, model registry. Standard ML-ops baseline. MLflow 3.0 (2024) added LLM-eval features.

  • PEFT (Hugging Face) Parameter-efficient fine-tuning

    Hugging Face's PEFT library — LoRA, AdaLoRA, IA3, Prefix Tuning, P-Tuning. Standard interface for adapter-style fine-tuning.

  • Prefect Workflow orchestration

    Python-native workflow orchestrator — widely used for ML pipelines and ETL. Prefect Cloud + open-source Prefect Core.

  • Roboflow Computer-vision data platform

    End-to-end computer-vision platform — dataset hosting, annotation, training, deployment. Universe public-dataset library; widely used for production CV.

  • Scale AI Data labelling at frontier scale

    Data-labeling powerhouse for AI training — RLHF data, fine-tuning corpora, frontier-lab partnerships. Meta invested $14B in Scale in 2025 for a 49% stake; Wang joined Meta as Chief AI Officer.

  • Snorkel AI Programmatic data labelling

    Programmatic data labelling and weak supervision platform — built on Snorkel research from Stanford. Enterprise platform for high-volume data programming.

  • Stable-Baselines3 Reliable RL algorithm implementations

    PyTorch implementations of RL algorithms (PPO, SAC, DQN, etc.) — research-grade reliability; widely used as a reference benchmark.

  • TRL (Hugging Face) Transformer Reinforcement Learning library

    Hugging Face's open-source library for fine-tuning and aligning transformer models with RLHF, DPO, ORPO, KTO and PPO trainers. The default RLHF stack on top of HF Transformers.

  • Unsloth Fast/efficient LLM fine-tuning

    Fast and memory-efficient LLM fine-tuning library — claims 2-5× speedup, 60-80% memory savings on a single GPU vs HF baseline through custom Triton kernels.

  • Weights & Biases ML experiment tracking + ops

    Experiment-tracking and ML-ops platform — log metrics, artifacts, datasets, sweeps. Acquired by CoreWeave in 2025 for $1.7B.

Tier 2 — production-ready (18)

  • Argilla Data-quality + labelling for LLMs

    Open-source data-quality platform for LLM training data — human-in-the-loop labelling, RLHF data collection, dataset curation. Acquired by Hugging Face 2024.

  • Axolotl YAML-driven LLM fine-tuning

    Open-source LLM fine-tuning library — single YAML config drives full / LoRA / QLoRA / DPO training across most popular base models.

  • ClearML ML experiment + orchestration

    Open-source ML experiment tracking, orchestration, and data management. Allegro AI commercial; widely used in research labs.

  • ColossalAI All-in-one large-scale training

    HPC-AI Tech's distributed training framework — ZeRO + tensor/pipeline parallel + heterogeneous training (CPU offload). HPC-AI raised $50M.

  • Comet ML ML experiment tracking + LLM eval

    ML experiment tracking + LLM evaluation platform (Opik for LLM eval). Comet competes head-to-head with Weights & Biases.

  • Composer (MosaicML) PyTorch training library

    Open-source training library from MosaicML (Databricks) — Trainer + recipes for fast convergence; underlies the MPT model training stack.

  • Flyte Data + ML workflow engine

    Kubernetes-native workflow engine — type-safe Python tasks; Lyft origin; Union AI commercial. Strong adoption in regulated industries.

  • lakeFS Git-like data lake

    Git-like branching, commits, and rollback for object-store data lakes (S3/GCS/Azure). Treeverse commercial; OSS + Cloud edition.

  • LLaMA-Factory Unified LLM fine-tuning toolbox

    Unified efficient fine-tuning of 100+ LLMs — SFT, DPO, PPO, KTO, ORPO with LoRA/QLoRA/full; CLI + WebUI. Popular in the Chinese OSS scene.

  • Megatron-Core Library-form of Megatron

    PyTorch-library form of Megatron-LM — distributed training building blocks (tensor/pipeline/expert parallel) packaged for embedding into custom training stacks.

  • Nemotron Synthetic Data (NVIDIA) Synthetic SFT/RLHF data pipelines

    NVIDIA's open synthetic-data generation pipelines (HelpSteer2, NemoSkills, Nemotron-4 340B). Used by frontier labs to bootstrap RLHF reward modeling data.

  • OpenRLHF Distributed RLHF training framework

    Distributed RLHF training framework built on Ray + DeepSpeed + vLLM — supports PPO/DPO/REINFORCE++/GRPO. Targeted at 70B+ models with efficient sample throughput.

  • Outerbounds Metaflow-based ML platform

    Commercial platform behind Metaflow (originally Netflix OSS) — ML workflow orchestration with managed compute, focused on data scientists shipping production pipelines.

  • Pachyderm (HPE) Data versioning + pipelines

    Data-versioning and pipeline platform — Pachyderm File System tracks lineage across pipeline runs. Acquired by HPE in 2023.

  • Ray Train Distributed training on Ray

    Ray's distributed training library — multi-node PyTorch/TF jobs; integrates Ray Tune and Ray Serve. Used at scale by OpenAI/Uber/Anyscale.

  • Surge AI Premium RLHF data service

    Premium human-labelling company specialised in RLHF data and difficult expert tasks — Surge prides itself on bootstrapped, no-VC; major frontier-lab supplier.

  • torchtune (PyTorch) Native PyTorch LLM fine-tuning

    Pure-PyTorch library for fine-tuning LLMs — Llama/Mistral/Phi/Gemma recipes; native FSDP2 + memory-efficient training. From Meta PyTorch team.

  • ZenML ML pipeline orchestration

    Open-source ML pipeline orchestration framework — connects sklearn/PyTorch/HF pipelines to MLflow/W&B/cloud infra. Pivoting toward LLM-ops.

Tier 3 — emerging (9)

  • Bonito Conditional task generator for instruction tuning

    Open model that generates instruction-tuning data conditioned on unstructured text — adapts pretrained LLMs to specialised tasks without manual labels (Brown Univ).

  • Distilabel Synthetic-data pipelines for LLMs

    Argilla's framework for synthetic data generation — pipelines of LLM steps that produce SFT/DPO training data; integrates with Hub.

  • DPO Direct Preference Optimization

    Stanford method that converts RLHF into a supervised-learning loss over preference pairs — no separate reward model needed; widely adopted alignment recipe.

  • GRPO Group Relative Policy Optimization

    DeepSeek-Math's GRPO — group-relative advantage estimation replacing PPO's critic. Used in DeepSeek-R1 reasoning post-training; widely adopted in OSS reasoning RLHF.

  • LoRA Low-rank adapter fine-tuning

    Microsoft Research's Low-Rank Adapters method — trains rank-r matrices added to attention weights; 10000x parameter reduction. Foundation of all adapter-style fine-tuning.

  • QLoRA Quantized LoRA

    4-bit quantized LoRA fine-tuning — train 65B models on a single 48GB GPU. UW (Tim Dettmers).

  • RL4LMs AllenAI RL-for-LM library

    Allen AI's research library for RL training of language models — NLPO algorithm, NLP-task RL benchmarks; predates the TRL/OpenRLHF surge.

  • RSL-RL Fast RL for robotics

    ETH RSL's lightweight PyTorch RL library — optimised for GPU rollouts in robotics simulation (Isaac/MuJoCo). Used in many legged-robot RL papers.

  • Verl Volcano RLHF framework

    ByteDance Volcano-Engine's open RLHF framework — flexible HybridFlow architecture; supports PPO/GRPO at large scale.

Tier 4 — early / experimental (2)

  • Cleanba Reproducible RL baselines

    Distributed PPO/IMPALA reference implementations from CleanRL author — research-friendly, reproducible.

  • RLHFlow RLHF reward modelling toolkit

    Research codebase for RLHF reward modelling — pairwise / Bradley-Terry / pointwise reward heads; release of RewardBench-winning models.