Skip to content

SingleController async path does not report avg_trajectory_age #3758

Open

Description

@tianyi-zhang-02

What

The SingleController async path doesn't report avg_trajectory_age. The pre-SC async GRPO path does, and under sampler.name: ready_first there is no other metric that would reveal how old the data a step trains on actually was.

The parity gap

On the async GRPO path, ReplayBuffer.sample computes

avg_trajectory_age = current_weight_version - sum(sampled_weights) / len(sampled_weights)

and grpo.py puts it into metrics, so it reaches the logger.

The SC path replaced that buffer with TQReplayBuffer plus the sampler policies in staleness_sampler.py. start_weight_list is still maintained, but nothing derives an age from it — it is used only for selection filtering, buffer bookkeeping and the checkpoint snapshot. So switching a recipe to the Single Controller silently drops a metric that was there before.

Why the two staleness counters don't cover it

step_metrics carries evicted_stale_prompt_groups and aborted_stale_inflight_groups. Both count what was discarded, not what was used. For ReadyFirstSampler both are always 0, structurally rather than because the run was healthy:

  • evict returns 0 unconditionally — "Every already-admitted group remains selectable; nothing is stale."
  • it inherits BaseSampler.should_abort_inflight, which returns False. WindowedSampler overrides this; _GatedSampler and ReadyFirstSampler do not.

Why the number matters, not just the parity

ReadyFirstSampler.select keeps every ready group whose start_weight <= current_train_weight. There is no lower bound, and the docstring is explicit that this is intended: "including late stragglers outside the admission window, so no rollout is ever discarded." Compare WeightFifoSampler, which computes min_valid_version = max(0, current_train_weight - self.max_staleness_versions).

So a group generated at version 5 can be trained on at version 50, and π_θ / π_gen for it spans 45 updates. That is what the truncated-IS knobs exist to bound — the icepop reference range of 0.5–5 assumes modest off-policyness. Right now you cannot tell whether clipping is occasional or dominant, because token_mult_prob_error reports the symptom while its main cause is unmeasured.

There is also a naming trap worth a docstring line either way: max_staleness_versions means "lookahead + selectable weight window" for weight_fifo and only "how far generation may run ahead" for ready_first. Switching name: weight_fifo → name: ready_first removes the staleness bound while the knob keeps its name and value.

Proposal

Record current_train_weight - start_weight for each prompt group a step actually selects, and report avg_trajectory_age (same name and same quantity as the async GRPO path, so one dashboard covers both) plus max_trajectory_age.

start_weight — stamped before rollout — is the right one: it identifies the policy that produced the tokens, which is the distribution the ratio corrects against. (end_weight measures how far the trainer moved during generation, a different quantity.)

Shape I have working and tested:

  • BaseSampler._finalize_selection computes the per-group ages just before self._buffer.remove(...), while the indices are still valid. Every built-in sampler goes through it, so all four report uniformly — for the bounded ones the max just confirms the window is holding.
  • The value is left unclamped. WindowedSampler, ReadyFirstSampler and WeightFifoSampler all filter start_weight <= current_train_weight, so they cannot produce a negative. InOrderSampler filters on target_step, and the admission gate only bounds the dispatch index from above (dispatch_index <= trainer_version + gate_window) while reserve stamps the generation worker's weight version — so non-negativity there is a property of how the two pumps interact, not something the sampler enforces. A negative would mean a stamping inversion, which seems more useful to surface than to round to zero.
  • Exposed as last_selection_trajectory_ages, deliberately not on the PromptGroupSampler Protocol: a sampler loaded by FQN from outside the repo needn't implement it, and the train pump reads it with getattr, degrading to reporting nothing rather than failing.
  • The train pump accumulates across the several select calls that assemble one step.

~155 lines including tests, no config change, no new dependency, nothing on the hot path.

Two calls I'd rather you made

  1. Is mean/max enough, or do you want a count of groups beyond max_staleness_versions? For ready_first that count is what answers "is the unbounded tail actually biting?", but I didn't want to add a third key unasked.
  2. Should the max_staleness_versions naming divergence be addressed here (a line on ReadyFirstSamplerConfig saying it gates admission only), or separately?

Happy to just send the PR if you'd rather read code. Flagging it first because this area is moving fast (#3582, #3589, #3590, #3660, #3665) and I'd rather not collide with something in flight.

Activity

  1. tianyi-zhang-02 commented on Aug 21, 2026

    @tianyi-zhang-02
    ContributorAuthor

    Implementation is up in #3759 — figured code was easier to react to than prose. Happy to close either this or the PR if you would rather have just one thread.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions