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.