Module 12. RL infrastructure for LLMs: memory, rollouts, and asynchrony
Read the explanations and code in the existing notebooks, or download the files. Code does not run on this site.
Download course files (.zip) 382.1 KBThe archive preserves the source folder structure, both languages, and supporting files.
Module 12 Lab: Memory, Queues, Staleness, and Recovery
Main module notebook
Download .ipynb module-12-rl-infrastructure.en.ipynb · 23.5 KB
Show notebook text
Module 12 Lab: Memory, Queues, Staleness, and Recovery
This notebook turns the infrastructure contracts in the lecture into one reproducible CPU workflow. It does not emulate a particular cluster cycle by cycle. Instead, it uses a discrete-event model with variable-length rollouts, a token-bounded queue, token-budgeted training batches, delayed weight publication, and an explicit maximum-lag policy. A separate failure drill shows why a resumable checkpoint is wider than model weights.
All numbers below belong to the course-owned synthetic environment. They explain mechanisms; they are not a benchmark of TRL, verl, OpenRLHF, NeMo RL, or any other stack.
1. Reproducible profile
The smoke profile runs on CPU. The full profile increases simulated time while preserving the same deterministic event model.
Python
from __future__ import annotations
import os
import random
import sys
import time
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from IPython.display import display
def find_repo_root(start: Path) -> Path:
"""Find the course root regardless of the notebook working directory."""
start = start.resolve()
for candidate in (start, *start.parents):
if (candidate / "modules/module-12/module_12_reference.py").is_file():
return candidate
raise FileNotFoundError("Could not locate the RL for LLM repository root")
ROOT = find_repo_root(Path.cwd())
MODULE_DIR = ROOT / "modules/module-12"
if str(MODULE_DIR) not in sys.path:
sys.path.insert(0, str(MODULE_DIR))
from module_12_reference import (
BYTES_PER_GIB,
checkpoint_cadence,
completion_token_mask,
effective_sample_size,
human_bytes,
ideal_sharded_memory,
kv_cache_bytes,
model_state_bytes,
optimal_checkpoint_interval,
peak_memory,
pipeline_report,
rollout_contract_errors,
straggler_report,
synchronization_schedule,
synchronization_transfer_time,
)
from module_12_pipeline_sim import (
PipelineConfig,
checkpoint_as_dict,
make_checkpoint,
recovery_drill,
result_row,
simulate_pipeline,
)
PROFILE = os.getenv("RL_LLM_PROFILE", "smoke").strip().lower()
RUN_REAL_MODEL = False
if PROFILE not in {"smoke", "full", "real-model"}:
raise ValueError("RL_LLM_PROFILE must be smoke, full, or real-model")
DURATION = 240.0 if PROFILE == "full" else 120.0
SEED = 12016
np.random.seed(SEED)
random.seed(SEED)
print(f"Course root: {ROOT}")
print(f"Profile: {PROFILE}; simulated duration={DURATION:.0f} s; seed={SEED}")2. Memory ledger
Begin with logical model state and useful KV-cache payload. The numbers still exclude activations, temporary materialization, fragmentation, and library overhead.
Python
profiles = {
"7B / 12-byte full-training profile": model_state_bytes(
7_000_000_000,
parameter_bytes=2,
gradient_bytes=2,
optimizer_state_bytes=8,
master_parameter_bytes=0,
),
"7B / 16-byte full-training profile": model_state_bytes(
7_000_000_000,
parameter_bytes=2,
gradient_bytes=2,
optimizer_state_bytes=8,
master_parameter_bytes=4,
),
"7B base + 1% trainable adapters": model_state_bytes(
7_070_000_000,
trainable_parameters=70_000_000,
parameter_bytes=2,
gradient_bytes=2,
optimizer_state_bytes=8,
master_parameter_bytes=4,
),
}
memory_rows = []
for name, profile in profiles.items():
memory_rows.append({
"profile": name,
"logical bytes": profile.total_bytes,
"GiB": profile.total_bytes / BYTES_PER_GIB,
"human": human_bytes(profile.total_bytes),
})
display(pd.DataFrame(memory_rows))
kv_16_gib = kv_cache_bytes(
num_sequences=64,
tokens_per_sequence=4096,
num_layers=32,
num_kv_heads=4,
head_dim=128,
element_bytes=2,
)
print("KV-cache teaching configuration:", human_bytes(kv_16_gib))
assert kv_16_gib == 16 * BYTES_PER_GIB
3. Phase peaks, overlap, and stragglers
Colocation requires a phase-level maximum rather than a sum of every possible component. We then compare a sequential loop, ideal overlap, and the loss caused by one slow member of a synchronous group.
Python
phase_report = peak_memory({
"rollout": [14 * BYTES_PER_GIB, kv_16_gib, 2 * BYTES_PER_GIB],
"training": [14 * BYTES_PER_GIB, 8 * BYTES_PER_GIB, 28 * BYTES_PER_GIB],
"weight publication": [14 * BYTES_PER_GIB, 14 * BYTES_PER_GIB, 2 * BYTES_PER_GIB],
})
print("Peak phase:", phase_report.phase, human_bytes(phase_report.peak_bytes))
display(pd.DataFrame([
{"phase": name, "GiB": value / BYTES_PER_GIB}
for name, value in phase_report.phase_totals.items()
]))
pipe = pipeline_report(rollout_seconds=60.0, train_seconds=20.0, batches=10)
stragglers = straggler_report([3, 3, 4, 4, 5, 5, 6, 18])
print("Two-stage pipeline:", pipe)
print("Synchronous straggler group:", stragglers)
assert np.isclose(pipe.sequential_total, 800.0)
assert np.isclose(pipe.pipeline_total, 620.0)
assert np.isclose(stragglers.utilization, 1 / 3)
4. Sequential and bounded-asynchronous loops
The next experiment runs a complete synthetic system: rollout workers, a token queue, a trainer, weight publication, and an explicit maximum policy lag.
Python
profile_configs = {
"sequential": PipelineConfig(
mode="sequential",
duration_seconds=DURATION,
seed=SEED,
),
"bounded async / sync every update": PipelineConfig(
mode="bounded_async",
duration_seconds=DURATION,
max_queue_tokens=7000,
max_policy_lag=2,
sync_every_updates=1,
seed=SEED,
),
"deeper async / sync every 3 updates": PipelineConfig(
mode="bounded_async",
duration_seconds=DURATION,
max_queue_tokens=7000,
max_policy_lag=6,
sync_every_updates=3,
seed=SEED,
),
}
results = {name: simulate_pipeline(cfg) for name, cfg in profile_configs.items()}
summary = pd.DataFrame([
{"profile": name, **result_row(result)}
for name, result in results.items()
])
display(summary)
# Profile rankings are observations, not invariants under parameter changes.
assert all(np.isfinite(r.trained_tokens_per_second) and r.trained_tokens_per_second >= 0 for r in results.values())
5. Queue depth and backpressure
Stopping admission does not erase in-flight work. Rollouts already running may complete after the threshold is crossed, so observed depth can overshoot briefly while new work is no longer launched.
Python
fig, ax = plt.subplots(figsize=(9.5, 4.8))
for name, result in results.items():
if result.config.mode != "bounded_async":
continue
x = [row["time"] for row in result.timeline]
y = [row["queue_tokens"] for row in result.timeline]
ax.plot(x, y, label=name)
ax.axhline(7000, linestyle="--", linewidth=1.5, label="new-rollout admission threshold")
ax.set_xlabel("time, s")
ax.set_ylabel("tokens in queue")
ax.set_title("A bounded queue converts overload into observable backpressure")
ax.grid(alpha=0.25)
ax.legend()
plt.show()
queue_table = pd.DataFrame([
{
"profile": name,
"maximum observed queue, tokens": result.max_queue_tokens_observed,
"backpressure, s": result.backpressure_seconds,
"discarded rollouts": result.discarded_rollouts,
}
for name, result in results.items()
])
display(queue_table)
6. Staleness and weight concentration
Throughput without statistical diagnostics is incomplete. Deeper overlap in this teaching model usually increases lag and reduces normalized ESS; that relationship is part of the simulator, not a universal law.
Python
batch_rows = []
for name, result in results.items():
for batch in result.batches:
batch_rows.append({
"profile": name,
"update": batch.update,
"mean lag": batch.mean_lag,
"max lag": batch.max_lag,
"normalized ESS": batch.normalized_ess,
})
batch_df = pd.DataFrame(batch_rows)
display(batch_df.groupby("profile").agg(
updates=("update", "count"),
mean_lag=("mean lag", "mean"),
p95_lag=("mean lag", lambda x: x.quantile(0.95)),
mean_nESS=("normalized ESS", "mean"),
min_nESS=("normalized ESS", "min"),
))
fig, ax = plt.subplots(figsize=(8.6, 4.6))
for name, part in batch_df.groupby("profile"):
if name == "sequential":
continue
ax.scatter(part["mean lag"], part["normalized ESS"], s=20, alpha=0.65, label=name)
ax.set_xlabel("mean policy lag in a batch, versions")
ax.set_ylabel("normalized ESS")
ax.set_ylim(0, 1.03)
ax.set_title("The simulator exposes a speed–freshness trade-off, not a theorem")
ax.grid(alpha=0.25)
ax.legend()
plt.show()
7. The rollout contract
Validate records before they enter an optimizer batch. In a production design, the same check belongs at a queue or batch-assembly boundary.
Python
good_record = {
"sample_id": "run12-p003-r01",
"prompt_id": "p003",
"policy_version": 17,
"prompt_tokens": [101, 42, 17],
"response_tokens": [314, 159, 2],
"behavior_logprobs": [-0.31, -1.07, -0.22],
"decoding": {"temperature": 0.8, "top_p": 1.0},
"reward": {"correctness": 1.0, "format": 0.0},
"stop_reason": "eos",
"verifier_version": "arith-v4",
}
broken_record = dict(good_record)
broken_record.pop("behavior_logprobs")
broken_record["policy_version"] = -1
print("Valid record errors:", rollout_contract_errors(good_record))
print("Broken record errors:", rollout_contract_errors(broken_record))
assert rollout_contract_errors(good_record) == []
assert rollout_contract_errors(broken_record)
8. Weight synchronization
Publication cost depends on payload, bandwidth, transport efficiency, and fixed setup overhead. Less frequent synchronization saves wall time while increasing policy lag.
Python
full_bf16_payload = 7_000_000_000 * 2
adapter_payload = 70_000_000 * 2
rows = []
for name, payload in (("full 7B bf16 weights", full_bf16_payload), ("1% bf16 adapter", adapter_payload)):
rows.append({
"payload": name,
"GiB": payload / BYTES_PER_GIB,
"transfer over 25 GB/s at 80% efficiency, s": synchronization_transfer_time(
payload,
25.0,
setup_seconds=0.15,
efficiency=0.80,
),
})
display(pd.DataFrame(rows))
cadences = []
for every in (1, 2, 4, 8):
schedule = synchronization_schedule(sync_every=every, updates=16)
cadences.append({
"sync every K updates": every,
"synchronizations": schedule.synchronization_count,
"mean lag": schedule.mean_lag,
"max lag": schedule.max_lag,
})
display(pd.DataFrame(cadences))
9. Checkpoints and a failure drill
Estimate a teaching checkpoint interval, then preserve not only counters but also RNG state, data cursor, and component versions.
Python
checkpoint_seconds = 180.0
mtbf_seconds = 10 * 3600.0
interval = optimal_checkpoint_interval(checkpoint_seconds, mtbf_seconds)
cadence = checkpoint_cadence(
interval,
checkpoint_seconds=checkpoint_seconds,
mtbf_seconds=mtbf_seconds,
)
print("First-order Young interval, seconds:", interval)
print("Checkpoint overhead model:", cadence)
rng = np.random.default_rng(SEED)
_ = rng.integers(0, 10_000, size=8)
manifest = make_checkpoint(
optimizer_step=3,
current_version=3,
published_version=3,
dataset_cursor=8,
rng=rng,
)
display(pd.DataFrame([checkpoint_as_dict(manifest)]).drop(columns=["rng_state"]))
drill = recovery_drill(seed=SEED, before=8, after=8)
print("Uninterrupted:", drill.uninterrupted)
print("Correct resume:", drill.resumed, "match=", drill.exact_match)
print("Weights-only resume:", drill.broken_resume, "match=", drill.broken_match)
assert drill.exact_match and not drill.broken_match
10. Optional open-model branch
This disabled-by-default branch measures generation and a teacher-forced forward on the same fixed sequences. Before enabling it, install a reviewed environment and set RL_LLM_MODEL_REVISION to the model/tokenizer commit SHA. Loading and a warm-up forward are outside the timings. Left padding, the chat template, dtype, library versions, seed, and actual lengths through the first EOS are recorded. The second timer covers forward only, with no backward/optimizer step; one local batch is not a framework or cluster benchmark.
Python
if not RUN_REAL_MODEL:
print("Optional real-model branch disabled; review the environment and set RUN_REAL_MODEL=True.")
else:
import hashlib
import importlib.metadata
import re
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "Qwen/Qwen3-0.6B"
revision = os.environ.get("RL_LLM_MODEL_REVISION", "")
if not re.fullmatch(r"[0-9a-f]{40}", revision):
raise ValueError("Set RL_LLM_MODEL_REVISION to the reviewed model's immutable 40-character commit SHA")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = (torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16) if device.type == "cuda" else torch.float32
tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision, padding_side="left")
if tokenizer.pad_token_id is None:
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(model_id, revision=revision, torch_dtype=dtype).to(device)
model.eval()
prompts = [
"Explain why an unbounded rollout queue is dangerous.",
"Give one invariant for atomic rollout weight publication.",
"Distinguish policy lag from training-inference mismatch.",
]
formatted = [tokenizer.apply_chat_template(
[{"role": "user", "content": prompt}], tokenize=False,
add_generation_prompt=True, enable_thinking=False,
) for prompt in prompts]
encoded = tokenizer(formatted, add_special_tokens=False, return_tensors="pt", padding=True).to(device)
generation_config = dict(max_new_tokens=64, do_sample=True, temperature=0.8,
top_p=1.0, top_k=0, pad_token_id=tokenizer.pad_token_id)
def synchronize():
if device.type == "cuda":
torch.cuda.synchronize()
# Download/loading and this one warm-up forward are outside both timings.
with torch.inference_mode():
model(**encoded, use_cache=False)
torch.manual_seed(SEED)
synchronize()
started = time.perf_counter()
with torch.inference_mode():
generated = model.generate(**encoded, **generation_config)
synchronize()
generation_seconds = time.perf_counter() - started
prompt_width = encoded["input_ids"].shape[1]
suffix = generated[:, prompt_width:]
eos_ids = model.generation_config.eos_token_id
eos_ids = [] if eos_ids is None else ([eos_ids] if isinstance(eos_ids, int) else eos_ids)
suffix_mask = completion_token_mask(suffix.cpu().numpy(), eos_ids)
lengths = suffix_mask.sum(axis=1).tolist()
attention_mask = torch.cat((encoded["attention_mask"], torch.as_tensor(
suffix_mask, dtype=encoded["attention_mask"].dtype, device=device)), dim=1)
position_ids = attention_mask.long().cumsum(-1) - 1
position_ids.masked_fill_(attention_mask == 0, 0)
# Teacher forcing scores the fixed generated sequences, with prompt PAD and
# suffix PAD masked. Explicit positions preserve left-padding semantics.
synchronize()
started = time.perf_counter()
with torch.inference_mode():
scored = model(input_ids=generated, attention_mask=attention_mask,
position_ids=position_ids, use_cache=False)
synchronize()
forward_seconds = time.perf_counter() - started
assert scored.logits.shape[:2] == generated.shape
print({
"model": model_id, "revision": revision, "dtype": str(dtype), "seed": SEED,
"device": str(device), "padding_side": tokenizer.padding_side,
"template_sha256": hashlib.sha256(tokenizer.chat_template.encode()).hexdigest(),
"enable_thinking": False, "generation_config": generation_config,
"versions": {name: importlib.metadata.version(name) for name in ("torch", "transformers")},
"generation_seconds": generation_seconds, "teacher_forced_forward_seconds": forward_seconds,
"new_tokens_including_eos": lengths, "active_generated_tokens": sum(lengths),
"generated_tokens_per_second": sum(lengths) / generation_seconds,
"active_forward_tokens": int(attention_mask.sum()),
"forward_tokens_per_second": int(attention_mask.sum()) / forward_seconds,
})
print("One local batch; no backward pass or optimizer. This is not a cluster benchmark.")