Pre-training
Pretraining runs one training loop across more GPUs than a single machine holds. Flyte launches the nodes as one group, wires up torch.distributed, and runs the data pipeline and the driver on cheaper hardware next to it.
Split the job by hardware
A pretraining run has three kinds of work with very different needs. Give each its own environment:
import flyte
from flyte.clustered import ClusteredTaskEnvironment, TorchRun
image = flyte.Image.from_debian_base().with_pip_packages("torch", "datasets")
data_env = flyte.TaskEnvironment(
name="data",
image=image,
resources=flyte.Resources(cpu=16, memory="64Gi", disk="200Gi"),
cache="auto",
)
train_env = ClusteredTaskEnvironment(
name="pretrain",
image=image,
resources=flyte.Resources(cpu=64, memory="512Gi", gpu="H200:8", shm="16Gi"),
replicas=4,
nproc_per_node=8,
runtime=TorchRun(),
)
driver_env = flyte.TaskEnvironment(
name="driver",
image=image,
resources=flyte.Resources(cpu=2, memory="4Gi"),
depends_on=[data_env, train_env],
)- Data loading gets CPUs, memory, and disk for tokenization, and no GPU.
cache="auto"means a second training run on the same dataset reuses the tokenized shards instead of rebuilding them. - Training gets the GPUs: 4 nodes with 8 GPUs each. Set
shmexplicitly, because NCCL and PyTorch data loaders use shared memory, and the container default is small. - The driver gets two CPUs.
depends_ondeploys the other two environments with it, so a singleflyte rundeploys all three.
Run torchrun across nodes
flyte.clustered.ClusteredTaskEnvironment runs a task on replicas identical pods, each starting nproc_per_node processes with torchrun:
# The torch wheel from PyPI bundles CUDA + NCCL, so the same image runs on CPU
# and GPU nodes. `flyte` itself provides the `clustered` runtime entrypoint that
# bootstraps torchrun inside each pod — no extra runtime library is needed.
image = flyte.Image.from_debian_base().with_pip_packages("torch", "numpy")
resources = (
flyte.Resources(cpu=(2, 4), memory=("4Gi", "8Gi"), gpu="L4:1")
if USE_GPU
else flyte.Resources(cpu=(1, 2), memory=("1Gi", "2Gi"))
)
env = ClusteredTaskEnvironment(
name="ddp_env",
image=image,
resources=resources,
replicas=REPLICAS, # number of pods (nodes) in the JobSet
nproc_per_node=NPROC_PER_NODE, # processes (one per GPU) per pod
runtime=TorchRun(rdzv_backend="static", max_restarts=0),
failure_policy=ClusterFailurePolicy(max_restarts=1),
)
resources applies to each pod, so 2 replicas with gpu="H100:8" and nproc_per_node=8 is a 16-GPU job. Inside the task, torchrun has already set RANK, WORLD_SIZE, and MASTER_ADDR, and flyte.ctx() exposes the same topology to your code:
@env.task
async def train_ddp(steps: int = 50, lr: float = 0.05) -> float:
"""Run DDP training across the cluster and return the final (rank-0) loss."""
import torch
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
ctx = flyte.ctx()
# Bind this rank to its local GPU BEFORE init_process_group so NCCL binds
# the right device.
if _BACKEND == "nccl" and torch.cuda.is_available():
torch.cuda.set_device(ctx.local_rank or 0)
device = torch.device(f"cuda:{ctx.local_rank or 0}")
else:
device = torch.device("cpu")
# torchrun has already populated RANK / WORLD_SIZE / MASTER_ADDR / MASTER_PORT.
dist.init_process_group(backend=_BACKEND)
rank = dist.get_rank()
world_size = dist.get_world_size()
print(
f"[rank {rank}/{world_size}] device={device} "
f"node_rank={ctx.node_rank} nnodes={ctx.nnodes} master_addr={ctx.master_addr}",
flush=True,
)
# Tiny model the workers train cooperatively: learn y = x · [1,1,1,1].
torch.manual_seed(0)
model = nn.Linear(4, 1).to(device)
ddp = DDP(model, device_ids=[device.index] if device.type == "cuda" else None)
opt = torch.optim.SGD(ddp.parameters(), lr=lr)
loss_fn = nn.MSELoss()
# Each rank trains on its own shard of synthetic data.
g = torch.Generator().manual_seed(rank)
x = torch.randn(64, 4, generator=g).to(device)
y = x.sum(dim=1, keepdim=True)
last_loss = 0.0
for step in range(steps):
opt.zero_grad()
loss = loss_fn(ddp(x), y)
loss.backward()
opt.step()
last_loss = float(loss.detach())
if rank == 0 and step % 10 == 0:
print(f"[rank 0] step {step:3d} loss {last_loss:.5f}", flush=True)
dist.barrier()
dist.destroy_process_group()
print(f"[rank {rank}] done — final loss {last_loss:.5f}", flush=True)
return last_loss
Every rank runs the same function. Return a value of the declared type from every rank; Flyte keeps rank 0’s output. Frameworks that read the torchrun environment, such as Lightning, Hugging Face Accelerate, and DeepSpeed, work unchanged. For FSDP and Lightning versions of this task, and the prerequisites on the cluster, see Clustered task environment. For JAX, set runtime=JaxRun() and nproc_per_node=1: JAX runs one process per node that owns every local GPU.
Save checkpoints to your object store
A pretraining run can take days, so save progress as it trains. Write to flyte.ctx().checkpoint from rank 0 only, so the ranks don’t write the same file:
# --- Checkpoint: gather a FULL state dict onto rank-0 and persist it. -----
# For models too large to gather, switch to StateDictType.SHARDED_STATE_DICT
# and have EVERY rank save its own shard under a rank-suffixed key instead.
save_cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
with FSDP.state_dict_type(fsdp_model, StateDictType.FULL_STATE_DICT, save_cfg):
full_state = fsdp_model.state_dict() # populated only on rank-0
cp = ctx.checkpoint
if rank == 0 and cp is not None:
buf = io.BytesIO()
torch.save(full_state, buf)
await cp.save(buf.getvalue())
print(f"[rank 0] saved FSDP checkpoint to {cp.path}", flush=True)
The checkpoint lives in your object store, not on the pod, so it survives a restart of the whole group and a retry of the task. Load it at the start of the task to continue from the last save. See Durable training.
Schedule all nodes together
A multi-node job can’t start until every node is available. If pods start one at a time, the first nodes sit idle while the rest wait for capacity. Admit the job as one gang through a Kueue queue on the cluster:
gang = ClusteredTaskEnvironment(
name="gang_scheduling_demo_workers",
image=image,
env_vars=knobs,
resources=flyte.Resources(cpu=2, memory="8Gi", gpu=GPU) if GPU else flyte.Resources(cpu=1, memory="2Gi"),
replicas=WORKERS,
nproc_per_node=1,
runtime=TorchRun(rdzv_backend="static", max_restarts=0),
failure_policy=ClusterFailurePolicy(max_restarts=0),
# This label is the whole opt-in: the job is admitted as one gang by the cluster's queue.
pod_template=flyte.PodTemplate(labels={"kueue.x-k8s.io/queue-name": QUEUE}),
)
If a higher-priority job can preempt yours after it starts, set failure_policy=ClusterFailurePolicy(resume_after_preemption=True), so the job waits to be readmitted and resumes from its checkpoint without using a retry. See Durable training.
Profile a training step
To find out whether a step is bound by compute, communication, or data loading, capture an NVIDIA Nsight Systems profile from the task. See the Nsight integration.
Related
- Distributed LLM pretraining: a GPT model pretrained on multiple H200 nodes, end to end.
- Fine-tuning a VLM: multi-node DeepSpeed training with checkpoints you can resume from.
- PyTorch: the
Elasticplugin config, an alternative to the clustered environment. - Ray: run a Ray cluster from a task.
- Weights & Biases distributed training: log from every rank.