Training loop

The training loop runs every iteration of RL: sample prompts, generate rollouts, score them, update the policy, and repeat with the new weights. In Flyte, the loop is a Python for loop in a driver task. The driver holds no GPU and does no heavy work; it calls tasks in the rollout, reward, and train environments, and keeps track of where the run is.

Define the environments

Each part of the loop gets its own environment. The trainer needs a GPU only while it updates the policy, and the driver is a small CPU task:

rl_grpo_lora.py
                
train_env = flyte.TaskEnvironment(
    name="rl-grpo-train",
    image=image,
    resources=flyte.Resources(cpu=4, memory="24Gi", gpu=flyte.GPU("L4", 1), shm="auto"),
    secrets=[HF_SECRET],
)

rl_grpo_lora.py
                
driver_env = flyte.TaskEnvironment(
    name="rl-grpo-driver",
    image=image,
    resources=flyte.Resources(cpu=1, memory="2Gi"),
    depends_on=[rollout_env, reward_env, train_env],
)

depends_on lists every environment the driver calls, so deploying the driver deploys the whole loop. For the rollout and reward environments, see Inference rollouts and Reward setup.

Write the loop

The driver runs one iteration per loop pass:

rl_grpo_lora.py
                
                    
                
            
                
            

                
                    
                
            
@driver_env.task(report=True)
async def train_rl(base: flyte.io.Dir, num_iterations: int = NUM_ITERATIONS) -> flyte.io.Dir:
    """Own the GRPO loop: fan out rollouts, score as they finish, take one GRPO step, repeat.

    Loop state (iteration, current adapter, and the report history) is checkpointed each iteration so a
    preempted driver resumes mid-run — including the accumulated report rows — instead of restarting.
    Progress is published to a live HTML report (``report=True``) after every iteration.
    """
    ctx = flyte.ctx()
    cp = ctx.checkpoint if ctx is not None else None

    start_iter = 0
    adapter: flyte.io.Dir | None = None
    adapter_version = 0
    history: list[IterationMetrics] = []

    # Resume from a prior driver attempt, if any.
    if cp is not None:
        prev = await cp.load()
        if prev is not None:
            state = json.loads(prev.read_text())
            start_iter = state["iteration"] + 1
            adapter_version = state["adapter_version"]
            adapter = flyte.io.Dir.from_existing_remote(state["adapter_path"])
            history = [IterationMetrics(**row) for row in state.get("history", [])]
            logger.info("Resumed from checkpoint at iteration %d (adapter v%d)", start_iter, adapter_version)

    # Cold start: mint a fresh LoRA adapter (version 0).
    if adapter is None:
        adapter = await init_adapter(base)
        adapter_version = 0

    await _publish_report(history, status="running")

    for it in range(start_iter, num_iterations):
        with flyte.group(f"iter-{it}"):
            prompts = _sample_prompts(it)

            # Launch every rollout group at once on the warm replicas.
            rollout_futs = [
                asyncio.create_task(generate(base, q, a, adapter, adapter_version, group_id=gid))
                for gid, (q, a) in enumerate(prompts)
            ]

            # Score each group the instant its rollout finishes — reward overlaps in-flight rollouts.
            # One reward task per group (not per rollout): see score_group.
            flat_rollouts: list[Rollout] = []
            reward_futs: list[asyncio.Task[list[float]]] = []
            for fut in asyncio.as_completed(rollout_futs):
                group = await fut
                flat_rollouts.extend(group)
                reward_futs.append(asyncio.create_task(score_group(group)))

            group_rewards = await asyncio.gather(*reward_futs)  # aligned with append order
            rewards = [r for gr in group_rewards for r in gr]  # flatten → aligned with flat_rollouts
            mean_reward = sum(rewards) / len(rewards) if rewards else 0.0
            logger.info("iter %d: %d rollouts, mean reward %.3f", it, len(rewards), mean_reward)

            # One GRPO step → next adapter version.
            new_version = adapter_version + 1
            adapter, mean_loss, contributing = await train_step(
                base, flat_rollouts, rewards, adapter, new_version
            )
            adapter_version = new_version

            # Record metrics for the report (accuracy/format derived directly from the rollouts).
            n = len(flat_rollouts)
            correct = sum(1 for r in flat_rollouts if _extract_answer(r.completion) == r.answer)
            formatted = sum(1 for r in flat_rollouts if "####" in r.completion)
            best_idx = max(range(n), key=lambda i: rewards[i]) if n else None
            history.append(
                IterationMetrics(
                    iteration=it,
                    adapter_version=adapter_version,
                    num_rollouts=n,
                    mean_reward=mean_reward,
                    max_reward=max(rewards) if rewards else 0.0,
                    accuracy=correct / n if n else 0.0,
                    format_rate=formatted / n if n else 0.0,
                    mean_loss=mean_loss,
                    contributing=contributing,
                    sample_question=flat_rollouts[best_idx].question if best_idx is not None else "",
                    sample_completion=flat_rollouts[best_idx].completion if best_idx is not None else "",
                    sample_reward=rewards[best_idx] if best_idx is not None else 0.0,
                )
            )
            await _publish_report(history, status="running")

            # Persist loop state so a preempted driver resumes here (with its report history).
            if cp is not None:
                state = {
                    "iteration": it,
                    "adapter_version": adapter_version,
                    "adapter_path": adapter.path,
                    "history": [vars(m) for m in history],
                }
                await cp.save(json.dumps(state).encode())

    await _publish_report(history, status="complete")
    assert adapter is not None
    return adapter

Each part of the loop uses one construct:

Fan-out. asyncio.create_task launches every rollout group at once.
Overlapping reward with generation. asyncio.as_completed hands each group to score_group the moment it finishes, so scoring runs while slower groups are still generating.
The update. train_step takes the rollouts and rewards and returns the next adapter as a flyte.io.Dir, which the next iteration’s rollouts receive.
One node per iteration in the UI. flyte.group(f"iter-{it}") collects an iteration’s rollout, reward, and train actions, so you can open iteration 12 and see exactly which completions it trained on.
A live report. report=True and flyte.report.replace.aio publish reward, accuracy, and loss after every iteration.
Resuming. The driver saves the iteration number, the adapter’s path, and the report history to flyte.ctx().checkpoint after every iteration. If a later attempt of the driver starts, it loads that state and continues from the next iteration instead of from the start. Set retries on the driver task so that a failed or preempted driver gets another attempt. The checkpoint belongs to the task’s attempts within one run, so a new run starts from iteration 0.

Compute the update

The trainer turns rewards into a learning signal. GRPO normalizes each completion’s reward against the other completions of the same prompt:

rl_grpo_lora.py
                
def _group_normalized_advantages(rollouts: list[Rollout], rewards: list[float]) -> list[float]:
    """GRPO advantage: within each prompt group, ``(r - mean) / (std + eps)``."""
    import statistics
    from collections import defaultdict

    by_group: dict[int, list[int]] = defaultdict(list)
    for i, r in enumerate(rollouts):
        by_group[r.group_id].append(i)

    advantages = [0.0] * len(rollouts)
    for idxs in by_group.values():
        group_rewards = [rewards[i] for i in idxs]
        mean = statistics.fmean(group_rewards)
        std = statistics.pstdev(group_rewards) if len(group_rewards) > 1 else 0.0
        for i in idxs:
            advantages[i] = (rewards[i] - mean) / (std + 1e-4)
    return advantages

A completion that scored above its group’s mean gets a positive advantage, and the update makes it more likely. A group where every completion scored the same gives no signal and is skipped. The policy update itself is ordinary PyTorch inside train_step. Libraries such as TRL work there too.

Scale the trainer across nodes

When the policy is too large to update on one node, make the trainer a flyte.clustered.ClusteredTaskEnvironment:

rl_grpo_lora_clustered.py
                
train_env = ClusteredTaskEnvironment(
    name="rl-grpo-train-cluster",
    image=image,  # same image as Part 1 — torch/peft/transformers already in it
    resources=_train_resources,  # per-POD resources (not per-JobSet)
    replicas=REPLICAS,
    nproc_per_node=NPROC_PER_NODE,
    runtime=TorchRun(rdzv_backend="static", max_restarts=0),  # in-pod torchrun; restarts are JobSet-level
    failure_policy=ClusterFailurePolicy(max_restarts=2),
    secrets=[HF_SECRET],
    env_vars=_PROFILE_ENV_VARS,
)

resources applies to each pod, and the train task runs on every rank, as in Pre-training. Rank 0 saves the new adapter and broadcasts its path, so every rank returns the same values of the declared types. Rank 0’s return value becomes the task’s output, so the driver’s loop doesn’t change. A multi-node trainer has a cold start of its own, because every pod has to be scheduled before the step begins. The next section hides it.

Overlap generation with training

In a sequential loop, the GPUs for generation wait while the trainer runs, and the trainer waits while generation runs. Start generating the next batch before training on the current one. The body of the pipelined loop, simplified from the multi-node example:

        
    
for it in range(start_iter, num_iterations):
    rollouts, rewards = await pending

    # Generate batch N+1 with the current adapter while batch N trains.
    if it + 1 < num_iterations:
        with flyte.group(f"iter-{it + 1}-generate"):
            pending = asyncio.create_task(generate_and_score(it + 1, adapter, version))

    version += 1
    with flyte.group(f"iter-{it}-train"):
        adapter, opt_state, metrics = await train_step_clustered(
            base, rollouts, rewards, adapter, version, opt_state
        )

Before the loop, the driver starts pending for the first batch. Each trained adapter gets a new version, which the rollout tasks use to tell adapters apart.

The trainer’s time, including its cold start, now runs under generation, and each iteration takes about as long as generation alone. The cost is that batch N+1 is generated by a policy one step behind the one it trains. Correct for this in the loss with a clipped importance ratio, using the sampling log-probabilities recorded during generation:

rl_grpo_lora_clustered.py
                
                    
                
            
                
            

                
                    
                
            
def grpo_sample_loss(
    new_logp: Any,
    old_logp: Any,
    ref_logp: Any,
    advantage: float,
    clip_eps: float = CLIP_EPS,
    kl_beta: float = KL_BETA,
) -> tuple[Any, dict[str, float]]:
    """Token-level clipped-ratio GRPO loss for ONE completion.

    All inputs are fp32 tensors of shape (T,) over the completion tokens; only ``new_logp``
    requires grad. ``old_logp`` is the *behavior* policy — vLLM's sampling-time logprobs, i.e.
    the adapter version that actually generated the rollout (one step stale when pipelining).
    ``ref_logp`` is the frozen base with the adapter disabled, for the KL anchor.

    loss = −( mean_t min(r_t·Â, clip(r_t, 1±ε)·Â) − β·mean_t KL_k3 ),  r_t = exp(new−old).

    The k3 estimator ``exp(δ) − δ − 1`` (δ = ref − new) is non-negative and low-variance
    (Schulman). At r = 1 the surrogate's *gradient* equals REINFORCE's (its value does not).
    """
    import torch

    ratio = torch.exp(new_logp - old_logp)
    clipped = torch.clamp(ratio, 1.0 - clip_eps, 1.0 + clip_eps)
    surrogate = torch.minimum(ratio * advantage, clipped * advantage).mean()
    delta = ref_logp - new_logp
    kl = (torch.exp(delta) - delta - 1.0).mean()
    loss = -(surrogate - kl_beta * kl)
    stats = {
        "loss": float(loss.detach()),
        "ratio": float(ratio.mean().detach()),
        "kl": float(kl.detach()),
        "clip_frac": float(((ratio - 1.0).abs() > clip_eps).float().mean().detach()),
    }
    return loss, stats

def ddp_shard_step(ddp_model: Any, sample_loss_fns: list, sync_forward: Any) -> None:
    """Accumulate per-sample gradients, then synchronize with EXACTLY ONE collective backward.

    DDP fires its gradient all-reduce inside ``backward()``, so all ranks must execute the same
    number of collective-bearing backwards — but shards are data-dependent and may be empty on
    some ranks. The documented escape hatch: backwards inside ``no_sync()`` accumulate grads
    locally with NO collectives (and free each sample's graph immediately), and the first
    forward-backward AFTER the context syncs everything. Every rank therefore runs one dummy
    1-token forward + zero-scaled backward, unconditionally — the collective count is identical
    on every rank *by construction*, empty shards included.

    (A bare ``(0.0 * sum(p.sum())).backward()`` would NOT work: it never passes through
    ``DDP.forward``, so the reducer is not prepared for the iteration.)
    """
    with ddp_model.no_sync():
        for fn in sample_loss_fns:
            fn().backward()  # local grad accumulation; per-sample graph freed here
    (sync_forward().float().sum() * 0.0).backward()  # the ONE synchronizing backward

old_logp comes from the rollout, new_logp from the policy being trained, and ref_logp from the base model with the adapter turned off. The ratio of new to old is clipped, so a stale sample can’t move the policy far, and the KL term keeps the policy close to the base model. See Keep rollouts on-policy.

Run held-out evaluation off the critical path too. Launch the evaluation of step N with asyncio.create_task and await it after step N+1 trains, so it never sits between iterations. See Evaluation.

For a GRPO loop on Qwen3-8B with a two-node, eight-GPU trainer, overlapping cut each iteration from 434 seconds to 298 in steady state. See We ran multi-node GRPO on 8 GPUs and the trainer cost us nothing.

Make each step durable

Checkpointing the driver protects the loop’s position. To also protect the work inside an iteration, make each step a pure function of its inputs:

Set retries on rollout, reward, and train tasks, so a transient failure reruns one task, not the iteration.
Seed each rollout from its prompt and iteration, for example with vLLM’s SamplingParams(seed=...), so a retried rollout produces the same result. Without a seed, sampling at a nonzero temperature gives a different result on every attempt, and caching the rollout isn’t sound.
Pass weights and optimizer state in and out of the train task as flyte.io.Dir, instead of keeping them in a long-lived process. A replayed step then recomputes the same outputs.
Set cache="auto" on steps that are pure, so a rerun of the loop returns finished steps from the cache.

Keep rollouts on a volume

Weights and optimizer state move best as a flyte.io.Dir, because every task that receives them reads all of them. See Move weights from the trainer to the generators. Rollouts are different: they accumulate across iterations, and later work often reads only part of them. An off-policy trainer samples from a replay buffer, an analysis reads the completions from a few iterations, and an evaluation rescores them with a new verifier.

Store them on a volume that each iteration forks and adds to:

        
    
import json

import flyte
from flyteplugins.union.io import ROVolume, allow_volumes

buffer_env = flyte.TaskEnvironment(
    name="rollout-buffer",
    image=image.with_pip_packages("flyteplugins-union"),
    pod_template=allow_volumes(cache_size="20Gi"),
    resources=flyte.Resources(cpu=2, memory="4Gi"),
)


@buffer_env.task
async def add_rollouts(
    buffer: ROVolume, it: int, completions: list[str], rewards: list[float]
) -> ROVolume:
    rw = await buffer.fork(name=f"rollouts-iter-{it}")
    data = await rw.mount()
    rows = (json.dumps({"iteration": it, "completion": c, "reward": r}) for c, r in zip(completions, rewards))
    (data / f"iter-{it:05d}.jsonl").write_text("\n".join(rows))
    return await rw.finalize(message=f"iteration {it}")

Each version is cheap. A fork is copy-on-write, so each iteration’s version stores only its own new file and shares the rest with the version before it.
Readers mount only what they need. A trainer task that receives the ROVolume mounts it and reads the iterations it samples from, without downloading the whole buffer.
Every version stays addressable. The driver keeps the ROVolume for each iteration, and any later task or run can mount the buffer as it was after any iteration.

Create the first, empty version with Volume.new() in a task before the loop, and pass the version each call returns to the next one, so the buffer is a chain of versions. Each fork has one writer, so writes never collide.

RL for LLMs with GRPO and LoRA: the loop on this page, including the multi-node, pipelined version.

Durable agentic RL with SkyRL and Harbor: an agentic loop with sandboxed episodes and a judge.
Intra-task checkpoints: the checkpoint store the driver uses.
Grouping actions: how flyte.group organizes the run’s UI.