Skip to content

feat(miles): NCCL broadcast weight-update transport (#42) - #43

Open
zhenyulincs wants to merge 1 commit into
zhenyu/miles-mvp-e2efrom
zhenyu/miles-nccl-broadcast
Open

feat(miles): NCCL broadcast weight-update transport (#42)#43
zhenyulincs wants to merge 1 commit into
zhenyu/miles-mvp-e2efrom
zhenyu/miles-nccl-broadcast

Conversation

@zhenyulincs

@zhenyulincs zhenyulincs commented Aug 16, 2026

Copy link
Copy Markdown
Collaborator

Closes #42.

What

Unlocks the NCCL broadcast weight-update transport for MILES RLix mode, as a user-selected mode — plus the concurrency/robustness fixes that 20-rollout under-load runs surfaced. Companion miles PR: rlops/miles#37 carries the sender-side NCCL + driver topology support; this PR is the rlix half.

Transport modes (RLIX_MILES_UPDATE_TRANSPORT)

mode behavior
cpu_serialize (default) pre-#42 behavior, zero change until opt-in
auto topology split: engines whose GPUs are disjoint from the pipeline's train pool receive via NCCL broadcast; colocate engines via cpu_serialize (the only mixing mode — NCCL cannot form a group containing the same physical GPU twice)
broadcast strict all-NCCL; classification-time fail-fast if any target engine shares a GPU with the train pool

Key mechanisms

  • Per-GPU NCCL rank cursor (world_size = 1 + per_engine × len(broadcast_set)) matching SGLang's rank_offset semantics; per-engine GPU count derived from rollout_num_gpus_per_engine, backed by a startup uniformity guard rejecting SGLang server-group overrides that diverge (heterogeneous engines are explicitly out of scope).
  • Port-claim protocol made real in miles mode: _PortClaimStore detached fallback actor (ROLL's SharedStorage doesn't exist there — claims were silently skipped, and concurrent pipeline syncs deterministically collided on get_free_port's 20000 base, cross-wiring NCCL TCP stores). Collision retry re-picks with start_port=port+1; claim release only after sender teardown-ack; wedged-sender path leaks-and-logs instead of deleting under a possibly-live TCP store.
  • Sync-under-load bracket (sync_base_weights_to_active): under the fully-async rollout, the router keeps dispatching during finalize's /flush_cache wait → guaranteed 60s timeout under real load. The bracket mirrors shrink_engines' quiesce ordering (router unregister → abort → sync), re-admits via the new atomic RolloutManager.register_router_if_active (state check + idempotent /add_worker in one serialized manager call — no TOCTOU vs concurrent shrink), escalates when re-registration fails after a successful sync, and probes the manager capability at registration so new-rlix/old-miles version skew fails before any engine is quiesced.
  • C9 smoke assertion (RLIX_ASSERT_MIXED_TRANSPORT=1): overlap smokes must produce mixed sessions per pipeline or be declared INVALID.

Evidence (details in plans/miles-nccl-broadcast-plan.md §10-§15)

  • 20-rollout dual overlap run (4×A100, mem 0.8, mode auto): 45 NCCL group creations (backend=nccl, world_size=3), 48 NCCL receives + 46 cpu legs (mixed every sync), concurrent syncs on distinct ports (20000×27 / 20001×18 — claim store verified under load), zero OOM / zero flush timeouts / EXIT 0, dual-smoke 7-condition pass-bar PASS.
  • Strict-mode rejection e2e: overlap + broadcast → actionable classification error, exit 1.
  • NVLink counter forensics: broadcast syncs move byte-exact GiB over NVLink (sender ΔTx == receivers ΔRx); cpu path moves zero.
  • Review trail: plan hardened over 6 codex adversarial rounds, implementation over 16 (all findings fixed; final verdicts approve).

Test plan

  • pytest tests/test_miles_service_broadcast.py — rank cursor, claim ownership incl. post-resolution cancel + wedged-sender leak, collision re-pick, fail-fast storage lookup (12 tests)
  • pytest tests/test_transport_mode_classification.py — mode matrix, strict rejection, uniformity guard, C9 assertion (18 tests)
  • pytest tests/test_sync_under_load_bracket.py + test_sync_bracket_concurrency.py — bracket AST ordering + behavioral concurrent-shrink / escalation / original-error precedence (7 tests)
  • Full rlix suite: 105 passed (1 pre-existing ray-2.54 connect() API failure in test_orchestrator_death_is_benign, fails identically on base — not introduced here)
  • E2E on 4×A100: 2-rollout smokes + 20-rollout under-load run, evidence above

Known limitations (documented in plan)

  • Dual + fully-disjoint + shared infer cards deadlocks (no train pressure → no GEN rotation on shared GPUs) — scheduler-regime follow-up (O7); on 4 GPUs use overlap+auto (dual) or single+disjoint (all-NCCL).
  • tms hook must be torch for broadcast + offload_train (preload leaves stale allocator pointers at sender staging).
  • All-broadcast positive e2e requires disjoint topologies; heterogeneous per-engine GPU counts out of scope (guard-enforced).

- unlock Path B in MilesModelUpdateService: per-GPU NCCL rank cursor
  (world_size = 1 + per_engine * len(broadcast_set)), teardown-ack
  port-claim release, wedged-sender leak-and-log, _PortClaimStore
  fallback backend for miles mode + collision re-pick via start_port
- user-selected transport mode env RLIX_MILES_UPDATE_TRANSPORT
  (cpu_serialize default / broadcast strict no-mix / auto topology
  split), C9 mixed-transport startup assertion, engine-GPU-count
  uniformity startup guard
- sync-under-load bracket in sync_base_weights_to_active: quiesce
  (router unregister + abort) before finalize, fixing the /flush_cache
  timeout under fully-async load; atomic manager re-admission via
  register_router_if_active; escalation when re-register fails after a
  successful sync; registration-time capability probe for version skew
- harness: env-parameterized sglang mem fraction (default preserved)
- tests: service broadcast + claim rules, transport-mode matrix,
  uniformity guard, bracket AST ordering + behavioral concurrency
- plan + full evidence trail (codex plan r1-r6, impl r1-r16) in plans/
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