Lingze Personal website for life and research

Miles (Post-Training Framework)

Q&A with claude Fable 5.1 in reading paper. A simple record.

Paper: Miles v0.1: Production-Level Post-Training

Code: https://github.com/radixark/miles

Two problems: sync between rollout and trainer (efficiency), and mismatch between them (correctness). Breaking the sync makes the mismatch worse.

The loop

Loop

Sync

Breaking the lockstep

Q: The loop looks ordered: rollout, train, sync weights, rollout. How can it be async?

The only real constraint is per trajectory: it must be generated before it is trained on. It does not have to come from the latest weights. Drop that assumption and the loop becomes a pipeline: the trainer works on step N while the engines generate for N+1 under weights from N−1.

  • Rollout keeps a fixed number in flight and refills a slot whenever one trajectory finishes.
  • Finished groups land in a buffer. The trainer pulls whatever is done.
  • The only barrier is the weight install pause, which takes seconds. That is why the transports matter.
  • Needs separate GPU pools. Colocation forces strict turn-taking.

Memory under async

Q: Under async both sides are busy. When is the actor offloaded?

It isn’t. Actor offload (evict the whole training process between steps) only helps when something else needs those GPUs, which means colocation or PPO’s critic. Under async and disaggregated it is off by default.

Q: Then how do we free training memory under async?

Optimizer streaming. Optimizer state lives on disk permanently. Forward and backward never read it, so it is not on the GPU during those. It is loaded one bucket at a time only inside step().

def step(self):
    for bucket in self.buckets:
        state = load_from_disk(bucket)    # → HBM
        for p in bucket:
            adam_update(p, p.grad, state)
        save_and_free(state)              # → disk, HBM released

This works because Adam is elementwise (parameter i needs only grad[i], m[i], v[i]) and all gradients exist before step() starts, so any subset can be updated independently.

State is 12 bytes per parameter: FP32 master weights plus two Adam moments. Moments can drop to BF16; master stays FP32. This is separate from the BF16 working weights, which stay on the GPU.


LoRA

Q: Is LoRA another way to cut memory?

Yes. Gradients, optimizer state, and weight-transfer bytes all scale with trainable parameters, so freezing the base collapses the 12 B/param term. Base weights and activations still stay resident. It is not Megatron built-in: each supported model is a contract between Megatron’s adapter wrapper, Miles’ module-name export, and SGLang’s serving, so support is per validated recipe.


Mismatch

Trajectory vs training sequence

Q: What is the difference between what rollout produces and what the trainer consumes?

Rollout’s unit is the trajectory: a whole episode with task, model turns, tool calls, and environment replies. It is scheduled, buffered, and rewarded as a unit, and the model generates only some spans of it. The trainer’s unit is a token sequence: prompt plus output, with a loss on the output positions.

Difference

The conversion between them is where mismatch enters. A conventional pipeline rebuilds the sequence from messages afterwards. Miles records the sequence while generating it, so no rebuild happens.


How the trainer uses a sequence

Q: We already have the rollout output. Why does the trainer run the actor again?

For gradients. Rollout hands over the sequence plus, at every generated position, the probability p the engine assigned to the sampled token. Those are plain numbers with no graph. The trainer runs the actor once over the whole sequence (one teacher-forced forward pass, not decoding) and reads off its own probability q for the same sampled token.

position t:   rollout distribution  [p1, p2, p3, ...]   sampled index 2 → record p2
              actor  distribution   [q1, q2, q3, ...]   read off      q2

              r_t = q2 / p2          p2 is a stored constant, q2 is a function of θ

Both are needed. p says how off-policy the sample is; q is the only one with a gradient. The policy gradient is valid only if q is scored on exactly the tokens the policy sampled. Any gap shows up as r drifting from 1. This is silent; nothing raises an error.


Sources of mismatch

Q: Where does the gap come from?

Source Treatment
Token mismatch: trainer rebuilds a different sequence eliminate: TITO (token-in-token-out)
Numerical mismatch: kernels, precision, MoE routing eliminate where possible: R3 (rollout routing replay), true-on-policy
Staleness: data from older weights (async only) bound + reweight

r can’t tell them apart, so the first two are removed structurally so r measures staleness only.


TITO (token-in-token-out)

Q: How does TITO avoid rebuilding?

The session server tokenizes the prompt and calls the engine, so it already holds the exact input IDs, and it keeps the output IDs, log-probs, and routed experts the engine returns. The trajectory becomes one contiguous sequence with a loss mask. Model turns are trained; prompt and tool output are masked. Each model turn is one LLM call.

The invariant is append-only: turn N+1’s tokens are turn N’s plus a suffix. This is verified per model family with a CPU check on the template and a GPU check with a real parser. No auto-detection.


Sessions

Q: Linear vs branching sessions?

Linear stores a chain: tail-append or one-step retry only, anything else rejected, one sequence. Branching stores an append-only tree: attach at the deepest matching prefix, the rest forks a new branch, nothing is deleted, one sequence per leaf.

Branching exists because Claude Code compacts context mid-task.

Q: Agent compacts context and does a turn. New branch?

Yes, attached near the root. The old branch survives and yields its own sequence.

Q: Two branches share one reward. Why not merge them into one trajectory?

Post-compaction tokens were sampled conditioned on the summary, not the original history. Splicing them after the original history would make their recorded p wrong. A branch is a complete trajectory sharing a prefix, not a segment.


The matcher

Q: Step 1’s tool call becomes step 2’s prefix, but the harness may replay it differently. Then?

The matcher decides which stored checkpoint a replayed message extends. The server always keeps its own tokens; the replay is only a lookup key. Too strict and cosmetic differences waste a valid prefix. Too loose and two turns with different tool calls but the same empty visible text collapse, so the trainer keeps a call that never happened. Default is strict.

Q: Why not use the session id for this?

The id says which tree, not which node. Append, retry, compact, and fork all arrive on the same id. And ordering is weaker than append-only: a compacted request is correctly ordered but rewrites the prefix.


Staleness

Q: How is staleness handled?

Staleness exists only under async. It is the strictness given up in exchange for throughput. Miles doesn’t fix weights; it reweights the gradient using r above.

  • TIS (truncated importance sampling): clamp r to [0, 2]. Outlier stays, damped.
  • Clip-or-pop: outlier’s weight set to 0. Dropped.

The staleness limit drops a group if current version minus the oldest version in the group exceeds the limit. One trajectory can span several weight versions, since retraction redoes only the in-flight request and keeps completed turns.


Loss

Q: What is the loss?

loss_t = −r · A · log q_t, with r and A detached. A is the trajectory’s reward centered within its group, broadcast to every unmasked token. Only log q_t carries the gradient.


OPD (on-policy distillation)

Q: The vocabulary is huge. Does the teacher return a full distribution?

No. One log-prob per position by default, for the token the student sampled. Top-K is optional. The teacher runs one teacher-forced forward pass over the student’s finished sequence; it does not decode. Must share the student’s tokenizer.

Q: The example uses the same model 5 RL steps ahead as teacher. Why does length drop?

Reverse KL is mode-seeking. On a sharpened teacher the student learns where to stop, not what to conclude. Reported: 56% shorter, accuracy unchanged. Useful when the teacher can’t be deployed, or as a dense signal beside a sparse reward.