feat(sc): report trajectory age in the SingleController train pump - #3759
Open
tianyi-zhang-02 wants to merge 5 commits into
Open
tianyi-zhang-02 wants to merge 5 commits into
tianyi-zhang-02 wants to merge 5 commits into
Conversation
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
tianyi-zhang-02
force-pushed
the
feat/sc-trajectory-age
branch
4 times, most recently
from
September 13, 2026 02:52
a9590b9 to
c821273
Compare
The pre-SC async GRPO path logs ``avg_trajectory_age`` from ``ReplayBuffer.sample``. The SC path replaced that buffer with ``TQReplayBuffer`` plus the sampler policies and never carried the metric over: ``start_weight_list`` is still maintained, but only for selection filtering, bookkeeping and the checkpoint snapshot. So moving a recipe to the Single Controller silently drops it. The two counters that sound like they would cover it -- ``evicted_stale_prompt_groups`` and ``aborted_stale_inflight_groups`` -- count what was discarded, and both are structurally zero under ``ready_first``: its ``evict`` returns 0, and it inherits ``should_abort_inflight`` -> False (``WindowedSampler`` overrides that; ``_GatedSampler`` does not). That matters because ``ReadyFirstSampler.select`` puts no lower bound on ``start_weight`` -- deliberately, so late stragglers are never thrown away -- while ``WeightFifoSampler`` bounds it with ``min_valid_version``. A group can therefore be arbitrarily many versions old, and the importance ratio spans all of them, with nothing reporting how far. ``_finalize_selection`` now records ``current_train_weight - start_weight`` per selected group, just before ``remove()`` while the indices are still valid. Every built-in sampler routes through it, so all four report uniformly; all four also filter ``start_weight <= current_train_weight`` (``InOrderSampler`` via ``target_step``, the dispatch index, which is never below the trainer version at dispatch), so the value is non-negative throughout. Exposed as ``last_selection_trajectory_ages`` and deliberately kept off the ``PromptGroupSampler`` Protocol: a sampler loaded by FQN from outside the repo needn't implement it, and the train pump reads it with ``getattr``, so it degrades to reporting nothing instead of failing. The pump accumulates across the several selects that assemble one step, then emits ``avg_trajectory_age`` -- same name and quantity as the async GRPO path, so one dashboard covers both -- and ``max_trajectory_age``, which is the one that shows a ready_first tail. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
The parametrized case covered the three samplers whose select filters on start_weight <= current_train_weight. InOrderSampler filters on target_step instead, so it was the one built-in the tests did not reach -- and the one whose selection does not itself bound the age. Adds it, plus a case pinning that a group newer than the trainer reports a negative age rather than a clamped zero. I had claimed non-negativity held for all four; tracing it, the gate only bounds dispatch_index from above (dispatch_index <= trainer_version + gate_window) and reserve stamps the generation worker's weight version, so that is an invariant of how the pumps interact, not one the sampler enforces. Better to report the inversion than to round it into something that looks fine. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
The existing tests are all on the sampler. Nothing asserted that the two
keys ever reach step_metrics, and the two things the pump itself does were
both unpinned:
- accumulation across selects. A step is assembled from several selects on
the streaming path, so an assignment instead of an extend would report
only the last one and pass any single-select test.
- the empty guard. Without it the step epilogue does sum([]) / len([]) at
step close, which kills the run -- and that is the path every sampler
without the property takes, since the metric is deliberately not on the
Protocol.
Mutation-tested: swapping extend for an assignment, and removing the guard,
each turn one of these red (the second as a ZeroDivisionError).
Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
tianyi-zhang-02
force-pushed
the
feat/sc-trajectory-age
branch
from
September 18, 2026 23:51
c821273 to
a868781
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Reports
avg_trajectory_ageandmax_trajectory_agefor trajectories consumed by each Single Controller train step.Built-in samplers return a typed
SamplerSelectionwith selected metadata, group count, and trajectory ages. The train pump accumulates ages across all selections contributing to a step. External samplers returning the existing two-tuple remain compatible.Age is
current_train_weight - selected_group.start_weight, intentionally unclamped so invalid negative ages remain visible.Validation
Refreshed September 18: head
a8687819be2f51922d78ae619224f80495e95102, based on upstreammainatb7a4d95d9099bae2b7f3713cc402aed3f21aef8d.The unchanged
TestPeriodicRolloutCheckpoint::test_logs_raw_and_canonical_rollout_throughputassertion expects20.0and observes6.666666666666667. It fails identically on unmodified current main in the same local harness. It was included in the complete run, not deselected.Environment: macOS arm64, Python 3.12.2, pytest 8.4.2, torch 2.8.0 CPU, Ray 2.51.1. This is not locked Linux worker or maintainer-CI validation. Rebase conflicts were limited to adjacent imports; range-diff preserves the feature/test commits. One test import ordering correction is included.
Coverage includes selected-only accounting, unclamped ages, reset behavior, every built-in sampler, external two-tuple compatibility, multi-selection accumulation, rollout metrics, and checkpoint restore. Main's eviction-before-selection ordering remains intact.
Earlier GPU evidence
Pre-refresh head
b48548c6dd692d10dab3a831f70d74abac6eb33dcompleted a 2×H100 run with Qwen3-0.6B, Megatron policy, real vLLM rollout, native TQ,ready_first, importance-sampling correction, andmax_staleness_versions=3:avg_trajectory_agemax_trajectory_ageIt finished at
train_steps=6,trainer_version=6, with no evicted, aborted, dropped, or replaced groups. This was not rerun at the current head; current-head GPU confirmation remains pending.