Train models with Flyte, from multi-node pretraining and LLM fine-tuning to classical machine learning, with every stage cached, recoverable, and sized to its own hardware.
Model training
A training job is more than a training loop. Data has to be prepared, the loop needs GPUs, the model has to be evaluated, and the result has to reach the people or apps that use it. Flyte runs each of these stages as a task, so each one gets its own hardware, its own retries, and its own cache.
The shape of a training pipeline
Most pipelines in this section follow the same structure. Separate task environments hold each kind of work, and a lightweight driver task calls them in plain Python:
flowchart LR
D["driver<br/>CPU"] --> P["prepare data<br/>CPU, cached"]
P --> T["train<br/>GPU or multi-node"]
T --> E["evaluate<br/>small GPU or CPU"]
E --> S["serve or publish"]
data_env = flyte.TaskEnvironment(name="data", resources=flyte.Resources(cpu=8, memory="32Gi"), cache="auto")
train_env = flyte.TaskEnvironment(name="train", resources=flyte.Resources(gpu="A100:8", shm="16Gi"))
driver_env = flyte.TaskEnvironment(name="driver", depends_on=[data_env, train_env])The driver is an ordinary async function. It can loop, branch on a metric, or fan out a sweep with asyncio.gather, and every task it calls shows up as a separate action in the run.
What Flyte adds to a training job
| Need | Construct |
|---|---|
| Put each stage on the right hardware | One flyte.TaskEnvironment per stage, with flyte.Resources for CPU, memory, GPU type and count, and shared memory. See Resources. |
| Train across several nodes | flyte.clustered.ClusteredTaskEnvironment runs torchrun across a group of pods. See Pre-training. |
| Skip work that already ran | cache="auto" on data prep and training. A rerun with the same inputs and code returns the stored outputs. See Caching. |
| Resume after a crash or preemption | retries plus flyte.Checkpoint restarts the loop from its last checkpoint. See Durable training. |
| Pass datasets, checkpoints, and weights between stages | flyte.io.Dir, flyte.io.File, and DataFrames store data in your object store and pass a reference. See Files and directories. |
| Watch a run while it trains | report=True and flyte.report publish live charts to the run’s UI. See Reports. |
| Run many trials at once | asyncio.gather or flyte.map fans out one task per trial. See Hyperparameter optimization. |
Choose a page
| You want to | Read |
|---|---|
| Train a model from scratch, or on more GPUs than one node has | Pre-training |
| Adapt a pretrained model with SFT, LoRA, or preference tuning | Post-training |
| Train scikit-learn or XGBoost models | Classical machine learning |
| Sweep hyperparameters and stop weak trials early | Hyperparameter optimization |
| Measure a model on held-out data and compare it with a baseline | Evaluation |
| Keep a long job from losing progress | Durable training |
To optimize a model with rewards instead of labels, see Reinforcement learning. To serve the model you train, see Deploy trained models.
Related
- Model training tutorials: complete pipelines for fine-tuning and hyperparameter optimization.
- Distributed LLM pretraining: a GPT model pretrained on multiple H200 nodes, end to end.
- Computer vision tutorials: fine-tuning detectors and vision-language models.