Miles Deep Dive: Main Loop, Rollout, and Reward

All Miles code links in this post are pinned to commit 39a1580 of radixark/miles. The repository moves quickly; substitute main for 39a1580 if you want the latest code.

In the verl GRPO post, the architecture was Ray + FSDP + vLLM: an actor generates rollouts, a reward model scores them, and the actor trains on the generated tokens with GRPO. This post does the same source-level walk for Miles, which replaces vLLM rollout with SGLang and has a different opinion about where reward computation should live.

The short version:

JSONL prompts
  -> RolloutDataSource.get_samples()
  -> SGLang generate()
  -> async_rm() / custom_rm()
  -> convert_samples_to_train_data()
  -> FSDP or Megatron actor.train()
  -> actor.update_weights() pushes new weights back to SGLang

That is still RL. It is not supervised fine-tuning. The label field is not a teacher response. It is metadata consumed by the reward function. The actor is trained on its own generated tokens, weighted by reward-derived advantages.

Start here: the main training loop

If you only read one file first, read train.py. It is the synchronous driver used by the standard loop. train_async.py is the pipelined variant: it starts the next rollout before training on the current one, but it is a separate execution mode with its own constraints.

The high-level schedule is:

rollout_manager, _ = create_rollout_manager(args, pgs["rollout"])
actor_model, critic_model = await create_training_models(args, pgs, rollout_manager)

await actor_model.update_weights()  # seed SGLang with actor weights

for rollout_id in range(args.start_rollout_id, args.num_rollout):
    rollout_data_ref = await rollout_manager.generate.remote(rollout_id)

    await actor_model.train(rollout_id, rollout_data_ref)

    await offload_train()
    await actor_model.update_weights()

That loop is deliberately compact. The important consequence is that rollout_manager.generate() is not “just model generation”. It returns already rewarded and tensorized train data. Rollout generation, reward-model calls, dynamic filtering, rollout-logprob recomputation, and conversion to train tensors all happen before actor_model.train(...) starts.

For performance debugging, the relevant boundary is:

step_time
  ~= train_wait_time   # waiting for RolloutManager.generate()
   + train_time        # ref log-probs + actor log-probs + actor update

So if a Miles step is slow, first decide whether the time is in the wait for RolloutManager.generate() (rollout, reward, log-prob recomputation, and conversion) or in actor training.

The four moving pieces

Miles’ mental model is explicit in its docs and code: every job has a prompt dataset, rollout engines, a reward function, and a training actor.

Piece	Miles object	Source
Prompt dataset	RolloutDataSource	miles/rollout/data_source.py
Rollout engine	SGLang servers + router	miles/backends/sglang_utils/sglang_engine.py
Reward	async_rm() / custom_rm hook	miles/rollout/rm_hub/__init__.py
Trainer	Megatron or experimental FSDP actor	train.py

The controller is intentionally small. In synchronous mode, train.py does this:

for rollout_id in range(args.start_rollout_id, args.num_rollout):
    rollout_data_ref = await rollout_manager.generate.remote(rollout_id)

    if args.use_critic:
        ...
    else:
        await actor_model.train(rollout_id, rollout_data_ref)

    await offload_train()
    if args.offload_rollout:
        await rollout_manager.onload_weights.remote()
    await actor_model.update_weights()
    if args.offload_rollout:
        await rollout_manager.onload_kv.remote()

Read that loop carefully. rollout_manager.generate() returns already-scored training data. The actor then trains on it. After the step, update_weights() sends the new actor weights to the SGLang rollout engines. This is the whole closed loop.

Input data: labels are reward metadata, not target text

Miles’ default data loader reads JSONL and builds Sample objects. A sample has these important fields (miles/utils/types.py):

@dataclass
class Sample:
    prompt: str | list[dict[str, str]] = ""
    response: str = ""
    response_length: int = 0
    label: str | None = None
    reward: float | dict[str, Any] | None = None
    metadata: dict = field(default_factory=dict)

The prompt loader duplicates each prompt n_samples_per_prompt times so GRPO has a group to compare. The relevant code is RolloutDataSource.get_samples():

prompt_samples = self.dataset.samples[self.sample_offset : ...]

samples = []
for prompt_sample in prompt_samples:
    group = []
    for _ in range(self.args.n_samples_per_prompt):
        sample = copy.deepcopy(prompt_sample)
        sample.group_index = self.sample_group_index
        sample.index = self.sample_index
        self.sample_index += 1
        group.append(sample)
    self.sample_group_index += 1
    samples.append(group)
return samples

So if a JSONL file has one prompt and --n-samples-per-prompt 4, Miles creates four Samples with the same prompt and label, then asks SGLang to sample four independent responses. The label is carried along so reward can compare the generated response against ground truth. It is not used as a supervised target.

For a summarization task, a Miles-compatible row can look like this:

{
  "messages": [
    {"role": "system", "content": "...instructions..."},
    {"role": "user", "content": "...member context..."}
  ],
  "label": "{\"ground_truth\": {...}, \"extra_info\": {...}}",
  "metadata": {
    "data_source": "rl_summarization",
    "ability": "logic"
  }
}

The messages field is the actor prompt. The label field is reward metadata: held-out job, activity/profile/resume fields, and flags like has_activity or has_queries. Miles’ built-in loader can apply the model chat template before generation with --apply-chat-template.

Rollout: where SGLang enters

Rollout is implemented in miles/rollout/sglang_rollout.py. The default train rollout eventually calls generate_and_rm_group() for each prompt group. The core shape is:

tasks = []
for idx, sample in enumerate(group):
    current_sampling_params = sampling_params.copy()
    tasks.append(
        asyncio.create_task(
            generate_and_rm(args, sample, current_sampling_params, evaluation=evaluation)
        )
    )

group = await asyncio.gather(*tasks)

Inside generate_and_rm() (lines 244-305):

# 1. Generate with SGLang
sample = await generate(args, sample, sampling_params)

# 2. Score the generated sample
if sample.reward is None:
    sample.reward = await async_rm(args, sample)

return sample

That is the important architectural fact: reward happens inside the rollout function, immediately after generation. The RolloutManager does not first return raw responses and then call a separate reward phase. It returns already-rewarded samples.

The outer rollout loop keeps submitting groups until it has enough accepted groups:

target_data_size = args.rollout_batch_size
data = []
while len(data) < target_data_size:
    while state.remaining_batch_size < target_data_size:
        samples = data_source(args.over_sampling_batch_size)
        state.submit_generate_tasks(samples)

    done, state.pendings = await asyncio.wait(
        state.pendings, return_when=asyncio.FIRST_COMPLETED
    )
    for task in done:
        group = task.result()
        ...
        data.append(group)

At the end it recomputes rollout log-probs via SGLang prefill:

await recompute_samples_rollout_logprobs_via_prefill(
    args,
    [sample for group in data for sample in group],
    url=get_model_url(args, "default"),
    sampling_params=state.sampling_params,
)

That gives Miles the rollout-time log-probs needed for the PPO/GRPO ratio.

Reward: the hook point is deliberately simple

The built-in reward hub is small (miles/rollout/rm_hub/__init__.py):

async def async_rm(args, sample: Sample, **kwargs):
    if args.custom_rm_path is not None:
        rm_function = load_function(args.custom_rm_path)
        return await rm_function(args, sample, **kwargs)

    rm_type = (metadata.get("rm_type") or args.rm_type or "").strip()
    if rm_type == "remote_rm":
        return await remote_rm(args, sample)
    elif rm_type == "deepscaler":
        return get_deepscaler_rule_based_reward(response, label)
    ...

There are three reward modes:

1. --rm-type deepscaler, math, f1, etc. — rule-based reward.
2. --rm-type remote_rm --rm-url ... — POST {prompt, response, label} to an external service.
3. --custom-rm-path package.module.function — import your own async function.

The custom function receives the generated sample. At that point:

• sample.prompt is the original prompt.
• sample.response is SGLang’s generated text.
• sample.label is whatever label was loaded from JSONL.
• sample.metadata is auxiliary data.

So a task-specific reward hook is not a trainer modification. It is a Python function.

For example, a summarization reward hook can parse sample.label, build two OpenAI-compatible reward-model prompts (apply prediction and faithfulness), call a reward server via /v1/chat/completions, and return one scalar:

async def custom_rm(args, sample):
    ground_truth = json.loads(sample.label)["ground_truth"]
    response = sample.response

    format_score = compute_format_score(response, ...)
    length_score = compute_length_score(...)
    apply_score = await reward_model_apply(response, ground_truth)
    quality_pass = await reward_model_quality(response, ground_truth)

    if quality_pass is False:
        apply_score *= quality_multiplier
    return max(0.0, format_score + apply_score + length_score)

Again: the label is only used by custom_rm. The actor never sees the label as target tokens.

Timing boundaries: what Miles measures by default

Miles has a useful high-level timing boundary, but not a full subphase breakdown. RolloutManager.generate() records one end-to-end rollout timer around _get_rollout_data(...) and log_rollout_data(...):

async def generate(self, rollout_id):
    start_time = time.time()
    data, metadata, metrics = await self._get_rollout_data(rollout_id=rollout_id)
    log_rollout_data(rollout_id, self.args, data, metrics, time.time() - start_time)
    data = convert_samples_to_train_data(...)
    return split_train_data_by_dp(...)

That rollout_time includes the whole rollout function:

SGLang actor generation
reward model scoring
rollout-logprob prefill
postprocess/convert

because reward is embedded inside generate_and_rm() and the rollout-logprob prefill happens before the rollout function returns. The trainer then reports its own train-side timing separately: reference log-probs, actor log-probs, actor training, and total train time.

If you need to distinguish pure generation time from reward time, the next instrumentation to add is explicit timers around:

sample = await generate(args, sample, sampling_params)
sample.reward = await async_rm(args, sample)
await recompute_samples_rollout_logprobs_via_prefill(...)

Without those timers, perf/rollout_time is a useful end-to-end number but not enough to separate pure SGLang generation from reward.

From rewarded samples to train data

Once rollout returns groups of rewarded samples, RolloutManager.generate() does three things (rollout_manager.py:104-120):

data, metadata, metrics = await self._get_rollout_data(rollout_id=rollout_id)
save_debug_rollout_data(self.args, data, rollout_id=rollout_id, evaluation=False)
log_rollout_data(...)
data = convert_samples_to_train_data(
    self.args,
    data,
    metadata=metadata,
    custom_convert_samples_to_train_data_func=...,
    custom_reward_post_process_func=...,
)
return split_train_data_by_dp(self.args, data, self.train_parallel_config["dp_size"])

This is where Sample objects become tensors: prompt/response token IDs, response lengths, loss masks, rollout log-probs, and scalar rewards.

The training actor receives per-DP shards of this tensorized data. It does not go back to the JSONL file. The JSONL file is only the source of prompts and reward metadata.

GRPO advantage in Miles

Miles’ advantage dispatcher is in miles/backends/training_utils/loss_hub/advantages.py. For --advantage-estimator grpo:

if args.advantage_estimator in ["grpo", "gspo"]:
    rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device)
    returns = get_grpo_returns(rewards, kl)
    advantages = [r for r in returns]

Then the advantages are whitened across the data-parallel group:

all_advs = torch.cat(advantages)
all_masks = torch.cat(loss_masks)
whitened_advs_flat = distributed_masked_whiten(
    all_advs,
    all_masks,
    process_group=dp_group,
    shift_mean=True,
)
advantages = list(torch.split(whitened_advs_flat, chunk_lengths))

The exact math is a little different in shape from verl’s uid-grouped implementation because Miles converts rollout groups into its own train-data format first, but the intent is the same: use scalar reward differences to weight the generated response tokens, without training a critic.

Weight update: actor to SGLang

After actor_model.train(...), the main loop calls:

await actor_model.update_weights()

The rollout side exposes the SGLang endpoints Miles needs. In sglang_engine.py you can see the supported update calls:

def update_weights_from_disk(self, model_path: str, load_format: str | None = None):
    return self._make_request("update_weights_from_disk", payload)

def init_weights_update_group(...):
    return self._make_request("init_weights_update_group", ...)

def update_weights_from_distributed(...):
    return self._make_request("update_weights_from_distributed", payload)

In colocated FSDP experiments, the actor can push flattened tensor buckets to SGLang; in Megatron recipes Miles uses a more production-oriented weight-sync path. Either way, the architecture is:

actor trains
  -> actor publishes updated weights
  -> SGLang rollout engine receives updated weights
  -> next rollout samples from new policy

This is the RL loop closure.

Where Miles differs from verl

The most important difference is where reward lives.

verl

In the verl architecture, reward can be a separate vLLM resource pool. In a colocated variant, the system can serialize phases on the same GPUs:

actor rollout
  -> sleep/offload actor rollout engine
  -> wake reward model
  -> score rollouts
  -> sleep reward model
  -> train actor

That is why verl can run a “235B reward model colocated with the actor” topology: the trainer controls the phase boundary between generation and reward scoring.

Miles

Miles’ default rollout function does:

generate one sample
  -> immediately call async_rm/custom_rm for that sample

Reward is inside RolloutManager.generate(). The main training loop sees only already-scored samples. That means the main loop cannot naturally insert:

release rollout memory
wake reward model
score
sleep reward model
resume rollout

between generation and reward. By the time the loop regains control, reward has already happened.

This is fine for rule-based rewards, remote reward services, or reward models on separate GPUs. It is not enough for a single-node topology where a 235B reward model must reuse the same eight GPUs after rollout memory has been released.

The source-level gap for colocated 235B reward

To make Miles match a verl-style colocated reward setup, the clean framework change is to split rollout into two phases:

raw_samples = await rollout_manager.generate_only(rollout_id)
await rollout_manager.offload(tags=[kv_cache, weights, cuda_graph])

rewarded_samples = await reward_manager.score(raw_samples)
train_data = rollout_manager.convert_to_train_data(rewarded_samples)

await actor_model.train(rollout_id, train_data)

There are a few ways to implement this in Miles:

1. Add a post_generation_pre_reward_hook inside RolloutManager.generate(). This is the smallest patch, but still awkward because reward is scattered inside generate_and_rm().
2. Add a generate_only rollout function and move reward application into RolloutManager.generate() after a memory-management hook.
3. Add a first-class RewardServer role, parallel to SGLang rollout servers, with onload/offload methods and an OpenAI-compatible scoring API.

Option 3 is the architecture I would want long-term. It turns the four objects into five:

Prompt dataset
Rollout SGLang engines
Reward SGLang engines
Actor trainer
Reference policy

and makes the phase schedule explicit.

What worked in a small integration smoke

On an 8×H200 pod, a minimal Miles integration reached a full one-rollout path with the experimental FSDP backend:

VERL parquet -> Miles JSONL
Miles RolloutDataSource -> SGLang rollout
custom_rm fallback reward -> train data conversion
FSDP actor train -> weight update back to SGLang

Two practical issues showed up:

1. --colocate defaults offload_rollout=True. With the installed SGLang / torch-memory-saver combo, pauseable CUDA graphs require preload mode. For the small FSDP smoke, disabling rollout offload (--no-offload-rollout) avoided that path.
2. With DP=8, the global batch must be at least 8. Tiny smoke settings like global_batch_size=2 produce a zero local batch on some ranks and fail before training.

The successful smoke does not prove performance. It proves the control path: data, SGLang, reward hook, GRPO train, and weight update can all execute.

Megatron backend status

Miles’ production recipes are Megatron-first. The FSDP backend is explicitly experimental. To run the production backend you need:

• A Megatron-LM source tree on PYTHONPATH.
• Megatron Bridge / conversion tooling.
• Actor and reference checkpoints converted to Megatron torch_dist.
• Model architecture flags matching the checkpoint.

The catch is version alignment. Miles expects Megatron APIs such as megatron.training.tokenizer; current NVIDIA Megatron-LM main has drifted. Megatron Bridge is available as a package, but it brings its own dependency set including Transformer Engine and ModelOpt. In practice you want the Miles Docker image or the exact Megatron commit used by the Miles recipe, not arbitrary NVIDIA main.

This is not a conceptual blocker, but it is an environment-version blocker. Treat Megatron conversion as a separate bring-up:

1. Pin the Miles-tested Megatron-LM commit.
2. Convert HF actor/reference checkpoint to torch_dist.
3. Run a small dense-model Megatron smoke.
4. Only then scale to the 235B MoE recipe.

Summary

Miles is not SFT. Its training signal is:

generated response tokens + scalar rewards -> GRPO advantages -> policy update

The JSONL label is reward metadata, not target text.

The code path is clean and compact:

• RolloutDataSource.get_samples() builds prompt groups.
• generate_and_rm() calls SGLang and attaches reward.
• RolloutManager.generate() converts rewarded samples into train tensors.
• actor_model.train() runs GRPO.
• actor_model.update_weights() syncs the actor back to SGLang.

The main difference from verl is phase control. Miles’ reward is inside rollout; verl can make reward its own phase. If you need a 235B reward model colocated on the same GPUs as rollout, Miles needs a framework change to split generation and reward scoring. For external reward services or rule-based rewards, the current hook is already enough.