trainything 0.1

A visual graph editor for composing and running multi-stage ML training pipelines.

Trainything lets you compose training pipelines visually: drop a dataset, drop a model, wire them together, configure, hit Start. Watch it train. Drag the checkpoint into the next stage. The goal is not to beat a hand-tuned Python script -- it's to make the entire multi-stage process legible, explorable, and fast to iterate on.

The primary audience is people learning how modern ML training works: how vision encoders connect to language models, how VLAs consume demonstration data, how RL rollouts feed back into training. If you can see the graph, you can build intuition.

Graph Model

The graph is bipartite (Petri-net style): artifact nodes (ovals) hold passive data at rest, process nodes (rectangles) run active computation. Every data edge crosses the artifact/process boundary -- never same-type to same-type.

A second edge type -- module edges -- wire process nodes together without implying execution order. For example, a Score node wired to a Train node provides a scoring function that Train calls internally during its loop.

Node types

What you can build

Supervised learning — Image classification, regression, fine-tuning pretrained vision models. CIFAR-10 → ResNet-18 → Train (supervised) → Evaluate

LLM pre-training — Train a GPT-2 from scratch on text data with next-token prediction. Tiny Shakespeare → Tokenize → GPT-2 (from config) → Train (SFT) → Generate

LLM fine-tuning — SFT or LoRA on instruction/chat datasets. Dataset → Model (LoRA) → Train (SFT) → Evaluate (generate)

R1-style reasoning — Teach a language model to reason step-by-step using GRPO with correctness and format rewards. GSM8K → GPT-2 → Train (GRPO-text) → Generate Answers

Image diffusion — Train a DDPM denoising model from scratch, then generate samples. CIFAR-10 → UNet2D → Train (diffusion) → Evaluate (generate_images)

Image generation — Run inference with pretrained diffusion pipelines (Stable Diffusion, etc.). Model (diffusion) → Evaluate (generate_images) → Generated Images

RL control — DQN, PPO, or GRPO on gymnasium environments (CartPole, LunarLander, Atari, MuJoCo). Environment → Score → Train (PPO/DQN) ← Model (MLP/CNN)

Multi-stage RL — Behavior cloning → GRPO → PPO, with checkpoint warmstarting between stages. Demos → BC → GRPO warmstart → PPO polish

Imitation learning — Clone expert demonstrations into a policy network. Dataset (demos) → Model (MLP) → Train (behavior_cloning) → Evaluate

Robot learning — Evaluate vision-language-action models (LeRobot/SmolVLA) on manipulation tasks. Model (lerobot) → Evaluate → Metrics

Getting Started

Requires Docker or Podman.

make build   # build the container image (CUDA PyTorch by default)
make run     # auto-detects GPU (NVIDIA, AMD, or CPU-only)

Open http://localhost:8000 in your browser.

For AMD ROCm or CPU-only PyTorch, rebuild with a different torch index:

make build TORCH_INDEX=https://download.pytorch.org/whl/rocm6.2   # AMD
make build TORCH_INDEX=https://download.pytorch.org/whl/cpu        # CPU

Install from PyPI

pip install trainything
trainything serve

To install with all optional environments (Atari, retro, diffusers, LeRobot, etc.):

pip install trainything[all]

For NVIDIA GPU support, install PyTorch with CUDA first:

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124
pip install trainything

On macOS (Apple Silicon), PyTorch automatically uses the MPS backend -- no extra steps needed.

Development setup

pip install -e ".[dev]"
cd frontend && npm ci && npm run build && cd ..
trainything serve --dev

Headless execution

trainything run examples/mnist_train.json
trainything run --dry-run examples/lunar_lander_ppo.json

Architecture

• Frontend: React + React Flow for the node editor. Side panel for node configuration, bottom panel for live logs and metrics.
• Backend: Python + FastAPI. REST for graph CRUD, WebSocket for streaming progress events.
• Execution: PyTorch, in-process via asyncio.to_thread. Each process node's execute() is the same code a user would write by hand.

Why Python backend

The entire ecosystem this tool orchestrates is Python: PyTorch, HuggingFace transformers, datasets, tokenizers, schedulers. A Rust backend would mean FFI or subprocess management for every node, and would break the key educational feature: "View Code." Each node's execute() method should be the same code a user would write by hand.

A Rust inference node (e.g. via Meganeura) can exist as a node type that calls out to a Rust binary -- the backend doesn't prevent this.

Acknowledgments

Inspired by ComfyUI's node-based approach to ML pipelines, Kedro's data-centric pipeline model, and the general frustration of managing multi-stage training with shell scripts and notebooks.