Skip to content

feat: add deterministic fully-async staleness control - #2278

Draft
0z5a wants to merge 3 commits into
THUDM:mainfrom
0z5a:codex/policy-version-metrics
Draft

feat: add deterministic fully-async staleness control#2278
0z5a wants to merge 3 commits into
THUDM:mainfrom
0z5a:codex/policy-version-metrics

Conversation

@0z5a

@0z5a 0z5a commented Aug 17, 2026

Copy link
Copy Markdown

Summary

Add a narrow, opt-in policy-version staleness path for fully-async rollout:

  • publish monotonically increasing policy versions around successful actor weight updates and stamp admitted samples with the scheduling version;
  • add --max-policy-version-lag N, accepting a completed group exactly when current_version - group_version <= N;
  • reject groups outside that deterministic window and requeue pristine admission-time prompt snapshots instead of relabeling old fragments;
  • expose fixed-cardinality lag buckets and stale/requeue, abort/requeue, queue-size, and current-version metrics through the existing rollout logging path.

The default lag budget is unset, so existing fully-async behavior is unchanged.

This draft is a deliberately smaller alternative to #1781 for discussion under #1800. It excludes partial-rollout resume, legacy_backpressure, version-window eviction, and benchmark scripts. The implementation remains split into three commits so the lifecycle, deterministic rejection policy, and observability layers can be reviewed or separated independently.

Test plan

  • uvx pre-commit run --files examples/fully_async/README.md slime/ray/rollout.py slime/rollout/fully_async_rollout.py slime/utils/arguments.py slime/utils/types.py tests/test_fully_async_rollout.py tests/test_megatron_argument_validation.py tests/test_sample.py tests/test_train_async_policy_version.py train_async.py
    • All hooks passed: repository checks, Ruff, autoflake, isort, and Black.
  • python -m pytest -q tests/test_fully_async_rollout.py tests/test_megatron_argument_validation.py tests/test_sample.py tests/test_train_async_policy_version.py
    • 49 passed.

Coverage includes monotonic lifecycle publication, failure behavior, disabled-budget compatibility, exact lag-boundary acceptance, stale rejection and pristine requeue, negative argument validation, bounded metric keys, reset semantics, and abort/requeue accounting.

Compatibility

  • No behavior change unless --max-policy-version-lag is set.
  • No partial rollout fragment is assigned a newer policy version.
  • No sample IDs or policy-version values are used as metric labels.

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant