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:

agentic_rl_durable_oss.py
                
# Separate environments so each tier's resources, retries, and timeouts are its own.
sandbox_env = flyte.TaskEnvironment(
    name="agentic-rl-oss-sandbox",
    image=image,
    resources=flyte.Resources(cpu="500m", memory="1Gi"),
)

judge_env = flyte.TaskEnvironment(
    name="agentic-rl-oss-judge", image=image, resources=flyte.Resources(cpu="500m", memory="512Mi")
)

trainer_env = flyte.TaskEnvironment(
    name="agentic-rl-oss-trainer", image=image, resources=flyte.Resources(cpu="500m", memory="512Mi")
)

driver_env = flyte.TaskEnvironment(
    name="agentic-rl-oss-driver",
    image=image,
    resources=flyte.Resources(cpu=1, memory="1Gi"),
    depends_on=[sandbox_env, judge_env, trainer_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:

agentic_rl_durable_oss.py
                
                    
                
            
                
            

                
                    
                
            
@driver_env.task(report=True)
async def train(
    n_steps: int = 3,
    n_worlds: int = 3,
    prompts_per_step: int = 3,
    n_samples: int = 2,
    rubric: Rubric = Rubric(),
    crash_driver_at_step: int = -1,
    flaky_trial_rate: float = 0.0,
    judge_flake_rate: float = 0.0,
) -> list[StepResult]:
    attempt = int(os.environ.get("FLYTE_ATTEMPT_NUMBER", "0"))
    started = flyte.durable.now()  # recorded once; replayed on retry
    world_ids = [f"w{i}" for i in range(n_worlds)]

    # worlds are cached; initial weights are deterministic, so step-0 generate inputs
    # hash identically across runs (which is what lets a re-scoring run reuse them)
    worlds = dict(zip(world_ids, await asyncio.gather(*[build_world(w) for w in world_ids])))
    weights = Weights(version=0, doc_prior={w: [1.0] * N_DOCS for w in world_ids})

    history: list[StepResult] = []
    tab = flyte.report.get_tab("training")
    for step in range(n_steps):
        with flyte.group(f"step-{step}"):
            # [L3] fan out trials; [L2b] every trial carries the current weights as an input
            specs = [
                TrialSpec(
                    trial_id=f"s{step}-p{p}-n{s}", step=step, world_id=world_ids[p % n_worlds], prompt="find the secret"
                )
                for p in range(prompts_per_step)
                for s in range(n_samples)
            ]
            gens = [asyncio.create_task(generate(sp, weights, worlds[sp.world_id], flaky_trial_rate)) for sp in specs]

            # pipeline: judge each trajectory the moment it lands; a failed trial is skipped, not fatal
            judges = []
            for fut in asyncio.as_completed(gens):
                try:
                    traj = await fut
                except Exception as e:  # noqa: BLE001 — skip_failed_rollouts
                    print(f"trial failed after retries, skipping: {e}", flush=True)
                    continue
                judges.append(asyncio.create_task(verify_and_judge(traj, rubric, judge_flake_rate)))
            results = await asyncio.gather(*judges, return_exceptions=True)
            # canonical order: as_completed order differs between attempts, and the step's inputs
            # must hash identically for the replayed step to be reused
            rewards = sorted((r for r in results if isinstance(r, Reward)), key=lambda r: r.trial_id)

            # one durable step; its output weights are the checkpoint
            result = await train_step(step, weights, rewards)
            weights = result.weights
            history.append(result)

            tab.log(
                f"<p><b>step {step}</b> · trials {result.n_trials}/{len(specs)} · mean reward "
                f"{result.mean_reward} · mean turns {result.mean_turns} · weights v{weights.version} · "
                f"driver attempt {attempt}</p>"
            )
            await flyte.report.flush.aio()

            if attempt == 0 and step == crash_driver_at_step:
                raise flyte.errors.RuntimeSystemError("simulated", f"driver crash after step {step} on attempt 0")

    tab.log(f"<p>started {started.isoformat()} · finished on attempt {attempt}</p>")
    await flyte.report.flush.aio()
    return history

Each part of the loop uses one construct:

Fan-out. asyncio.create_task launches every episode at once, and each one carries the current weights as an input.
Overlapping reward with generation. asyncio.as_completed hands each trajectory to verify_and_judge the moment it finishes. An episode that still fails after its retries is skipped, not fatal to the step.
Stable inputs. The rewards are sorted by trial ID before training, because as_completed returns them in a different order on each attempt. Sorted inputs hash the same way, so a replayed step can be reused.
The update. train_step is a pure task: it takes the weights and rewards and returns the next weights, which are the step’s checkpoint.
One node per step in the UI. flyte.group(f"step-{step}") collects a step’s episodes, judging, and update.
A live report. report=True and tab.log publish the reward and weights version after every step.
Resuming. If the driver fails and is retried, the actions it already completed are reused, so it reaches the step where it stopped without rerunning episodes. flyte.durable.now() records the start time once and returns the same value on a retry.

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.

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.

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.

Training loop

Define the environments

Write the loop

Compute the update

Scale the trainer across nodes

Overlap generation with training

Make each step durable

Related