Module 12. RL infrastructure for LLMs: memory, rollouts, and asynchrony

The four roles are:

1. the trainer, which computes losses, gradients, and optimizer steps;

2. the rollout engine, which samples policy actions and owns the KV cache;

3. the verifier or environment, which produces an outcome and operational status;

4. the orchestrator, which selects prompts, assembles groups, and manages queues, versions, and recovery.

They need not be separate processes. One GPU may alternate among all roles, while one logical role may span many workers on a cluster.

The invariant is that each training record unambiguously associates actions with the behavior policy that produced them and reward with the verifier or environment version that produced it.

A string and scalar reward omit tokenization, the policy-token mask, behavior log-probabilities, weight version, decoding parameters, stop reason, and reward provenance. The loss and the failure cannot be reconstructed reliably.

1. Placement is disaggregated because generation runs on a separate server.

2. The logical protocol is synchronous because every batch ends at a barrier.

3. Asynchrony is defined by production and consumption order, not by a network address. A separate server can serve a strictly sequential loop.

4. Colocation with buffering: one graphics processing unit (GPU) alternates between rollout and training while the CPU prepares the next prompt batch. Disaggregation with a barrier: a dedicated vLLM pool generates a batch, then waits until training and weight loading finish.

A common 16-byte profile is:

• 2 bytes for a bfloat16 parameter;
• 2 bytes for a bfloat16 gradient;
• 8 bytes for two float32 Adam moments;
• 4 bytes for a separate float32 master copy.

The profile changes with gradient precision, whether a master copy exists, optimizer type, state quantization, CPU offload, adapter-only training, sharding, and implementation details.

Dividing by the FSDP group size ignores replicated objects, temporary module gathering, communication buffers, activations, logits, allocator fragmentation, and transfer phases.

LoRA reduces parameters with gradients and optimizer state. It does not remove base weights, forward activations, the KV cache, temporary buffers, or the requirement to publish a consistent adapter version to the rollout engine.

1. For Group Relative Policy Optimization (GRPO), useful payload grows by a factor of four when concurrent sequences grow from 8 to 32.

2. Actual growth depends on early completion, length distribution, allocator pages, tensor parallelism, reserve, and prefix reuse.

3. PagedAttention reduces managed-cache fragmentation; prefix caching or RadixAttention reuses the common prompt prefix.

4. Keys and values depend on model parameters. A cache computed by the old weights is not a cache of the new policy.

The record assigns one version to a completion generated by a mixture of two weight sets. The behavior-policy denominator and replay semantics are therefore ambiguous.

Safe regimes include:

• let in-flight completions finish under version 120 and switch only new requests;
• record a version and behavior log-probabilities per action or turn when version changes are an explicit part of the protocol.

A mid-completion update requires token-level versions, actual log-probabilities after decoding transforms, an atomic update order, and an action mask.

Temperature, top-p, and logit masks transform the distribution after the weights. Identical transformers can therefore define different behavior policies.

The first problem is training–inference mismatch (TIM): the declared version is the same, but numerical probabilities differ.

The second is policy lag: the data were generated by an older version.

The control is a zero-lag comparison immediately after synchronization: score the same token sequence with the engine and trainer and compare token log-probabilities. Then vary version age separately.

ESS measures weight concentration. It does not prove that the denominator is correct, clipping is unbiased, masks are valid, or the chosen off-policy estimator matches the target objective.

1. The condition $\lambda<\mu$ is violated: 14 million tokens/min arrive and only 10 million are consumed.

2. Backpressure slows or pauses rollout workers when a queue threshold is reached.

3. Records can differ by orders of magnitude in length. Record count is a weak proxy for memory, compute, and age; token or cost limits follow physical load more closely.

4. Long episodes can be systematically harder. Removing them shifts the task distribution toward short and easy cases.

This is a new run initialized from old weights, not a continuation of the previous optimization process. Adam moments, scheduler state, and reward contract changed.

A semantically complete checkpoint commonly includes:

• model and adapter weights;
• optimizer state;
• learning-rate scheduler;
• global step and policy version;
• random-number generator states;
• dataset cursor;
• curriculum or prompt-scheduler state;
• reward and advantage normalization;
• tokenizer and chat-template revisions;
• verifier and environment versions;
• decoding configuration;
• rollout schema;
• distributed layout and queue-recovery policy.

Recovery should be tested by saving, terminating processes deliberately, restoring, executing a control batch, and checking versions, tokenization, finite losses, and expected statistics.

Frameworks optimize different levels: a ready-made trainer, role orchestration, a fully asynchronous dataflow, agent environments, or a particular distributed backend. There is no one-dimensional ranking without model size, scale, synchrony, environment type, fault-tolerance requirements, and team expertise.

Bytes per trainable parameter are

Full training requires

Base bfloat16 weights require

With LoRA, base weights still take 14 GB. Each of 20 million trainable parameters requires the adapter weight itself, a gradient, two moments, and a master copy:

The extra amount is

for a total of

The peak also contains activations, temporary buffers, logits, KV cache, fragmentation, replicas, and weight-transfer phases.

One cached token requires

That is

The total token count is

Memory is

Ideal uniform sharding across four participants gives

This is useful payload without pages, block tables, alignment, or temporary tensors.

One sequential batch takes

The steady-state period is

Twenty sequential batches take

The ideal pipeline takes

Full-run speedup is

and asymptotic period speedup is

The stages are unbalanced. Ten seconds of training cannot hide 90 seconds of rollout work.

Total work is

The synchronous group ends after

Capacity is

Utilization is

and idle fraction is

After safely truncating the final rollout at eight seconds, work is 43 GPU-s, wall time is 8 s, and capacity is 64 GPU-s:

Utilization improves, but if long episodes are systematically hard, truncation changes the observed task and reward distribution.

Weights are

Their sum is 13 and squared sum is

Thus

Normalized ESS is

Four equal weights have ESS 4 and normalized ESS 1.

After clipping the final log ratio at $\log2$, weights are $(1,2,2,2)$:

with normalized ESS about 0.942.

Clipping reduces concentration but changes the estimator. ESS describes weight geometry; it does not prove that the objective or policy quality is correct.

The minimum among these candidates is $K=4$.

The stationary lag cycle is $(0,1,2,3)$, so maximum lag is 3 and mean lag is

For ten updates:

with exact finite-run mean $13/10=1.3$.

The checkpoint interval is

The model omits restart time, asynchronous writes, correlated failures, storage and network cost, partially lost queues, distributed-layout incompatibility, and time-varying risk.

import math

REQUIRED = {
    "sample_id", "prompt_id", "policy_version",
    "prompt_tokens", "response_tokens", "behavior_logprobs",
    "decoding", "reward", "stop_reason", "verifier_version",
}


def rollout_contract_errors(record):
    if not isinstance(record, dict):
        return ["record must be a dict"]

    errors = []
    missing = sorted(REQUIRED - set(record))
    if missing:
        errors.append("missing fields: " + ", ".join(missing))

    version = record.get("policy_version")
    if isinstance(version, bool) or not isinstance(version, int) or version < 0:
        errors.append("policy_version must be a non-negative integer")

    for key in ("prompt_tokens", "response_tokens"):
        tokens = record.get(key)
        if not isinstance(tokens, list) or any(
            isinstance(token, bool) or not isinstance(token, int) or token < 0
            for token in (tokens or [])
        ):
            errors.append(f"invalid {key}")

    logps = record.get("behavior_logprobs")
    if not isinstance(logps, list) or any(
        not isinstance(value, (int, float))
        or not math.isfinite(value)
        or value > 0
        for value in (logps or [])
    ):
        errors.append("invalid behavior_logprobs")
    elif isinstance(record.get("response_tokens"), list) and len(logps) != len(record["response_tokens"]):
        errors.append("logprobs do not align with response_tokens")

    return errors

In Module 11, token_ids, policy_mask, and turn_ids describe a whole trajectory, while policy_versions identifies its turns. Here select one completion generated by one behavior-policy version: prompt_tokens is its generation context, response_tokens contains the policy-selected tokens in that completion, and behavior_logprobs aligns with those tokens. The single numeric policy_version requires an explicit version registry; do not silently coerce a version string into a number. Assign prompt_id when constructing the prompt. Mapping rewards to reward and choosing stop_reason require declared credit-assignment and stopping rules for this completion; do not automatically copy the episode outcome to every turn. Environment IDs, actions, observations, token provenance, and trajectory-level flags remain in the original log linked to the infrastructure record. This is a projection onto a different data unit, not a replacement for the full agentic contract.

def diagnose_pipeline(metrics):
    required = {
        "trainer_idle_fraction",
        "rollout_idle_fraction",
        "verifier_fraction",
        "weight_sync_fraction",
        "queue_growth_per_minute",
        "normalized_ess",
        "oom_events",
    }
    missing = required - set(metrics)
    if missing:
        raise ValueError("missing metrics: " + ", ".join(sorted(missing)))

    bounded = required - {"queue_growth_per_minute", "oom_events"}
    for key in bounded:
        if not 0 <= metrics[key] <= 1:
            raise ValueError(f"{key} must lie in [0, 1]")
    if metrics["oom_events"] < 0:
        raise ValueError("oom_events must be non-negative")

    hypotheses = []
    if metrics["trainer_idle_fraction"] > 0.25 and metrics["queue_growth_per_minute"] <= 0:
        hypotheses.append("check rollout, verifier, and input supply")
    if metrics["verifier_fraction"] > 0.30:
        hypotheses.append("profile the verifier critical path")
    if metrics["weight_sync_fraction"] > 0.20:
        hypotheses.append("profile weight transfer and conversion")
    if metrics["normalized_ess"] < 0.50:
        hypotheses.append("separate policy lag from training-inference mismatch")
    if metrics["oom_events"] > 0:
        hypotheses.append("measure phase-specific peaks, KV cache, activations, and transfer buffers")
    if not hypotheses:
        hypotheses.append("inspect a stage-level trace")
    return hypotheses

The important design choice is that outputs are hypotheses that can be tested with a targeted measurement. A single aggregate metric rarely identifies a root cause.