NVIDIA-NeMo / NVIDIA-NeMo/RL

SingleController async path does not report avg_trajectory_age

Open
#3,758 1 comment 0 reactions 1 assignee Claimed by @terrykong View on GitHub
accuracy community-request Feature waiting-on-maintainers
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

### 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

```python
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.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.