Skip to content

bench: add generic weight-sync policy sweep - #2354

Open
0z5a wants to merge 3 commits into
THUDM:mainfrom
0z5a:codex/weight-sync-benchmark
Open

bench: add generic weight-sync policy sweep#2354
0z5a wants to merge 3 commits into
THUDM:mainfrom
0z5a:codex/weight-sync-benchmark

Conversation

@0z5a

@0z5a 0z5a commented Sep 4, 2026

Copy link
Copy Markdown

Draft update — 2026-09-05

Adds fail-closed timing provenance to the benchmark. Historical 2/4-rank results below demonstrate benchmark execution, not production engine loading or a universal speedup.

AI assistance: OpenAI Codex. Remains Draft; implementation, CPU/Gloo checks and production GPU validation are separate acceptance gates. No new GPU results were generated during publication.

The original submission description and validation history follow.


Summary

  • add a standalone torch.distributed benchmark for trainer-to-rollout-engine weight synchronization
  • sweep NCCL broadcast and P2P transport, bucket/byte credits, engine-wave shape, buffer size, and phase stride
  • report slowest-rank latency, component timings, observed in-flight bytes/groups, rank skew, idle time, and optional per-transfer traces as JSON
  • add focused CPU unit coverage for topology construction, scheduling constraints, metrics, validation, and buffer reuse

Scope: benchmark only

This PR adds only tools/benchmark_weight_sync.py and tests/test_weight_sync_benchmark.py. It does not modify slime's production weight updater, Ray/SGLang lifecycle, or select a runtime scheduling policy. all_at_once, windowed, and serialized are controlled benchmark scenarios for measuring tradeoffs, not new production defaults or policy decisions.

General topology support

Rank 0 is the trainer and all remaining launched ranks are derived into logical engine groups. Homogeneous groups use --engine-group-size; heterogeneous layouts use --engine-group-sizes, whose sum must cover every non-trainer rank exactly once. The implementation supports any valid world size of at least two and contains no hard-coded five-rank, six-rank, or other topology special case.

Each logical engine receives its own trainer-plus-engine data and Gloo control groups, so TP/PP-style multi-rank and heterogeneous groups can be represented without changing the benchmark core.

Validation and evidence

  • CPU unit coverage exercises size parsing, arbitrary/heterogeneous world-size partitioning, all three wave policies, bucket/byte/group bounds, fail-fast byte validation, preallocated buffers, percentiles, slowest-rank aggregation, and observed concurrency reporting.
  • A 2-rank NVIDIA L20 CUDA/NCCL smoke completed 4 configurations / 8 measured samples across nccl_broadcast and p2p, covering the minimum valid trainer-plus-one-engine topology.
  • A 4-rank NVIDIA L20 matrix completed 48 configurations / 480 measured samples: two transports (nccl_broadcast, p2p), two message sizes (16 MiB, 64 MiB), bucket credits 1/2, all three engine-wave policies, and phase strides 0/500 us, with 2 warmups and 10 measured iterations per configuration.
  • The 4-rank report observed maximum logical engine-group concurrency of 3 / 2 / 1 for all-at-once / windowed / serialized, respectively, matching the planned policies for three rollout engine groups.
  • git diff --check passes for both commits.

These numbers demonstrate benchmark execution and policy observability; they are not a claim that one policy is universally faster. The engine-load stage is a synthetic device-to-device copy, and the benchmark intentionally excludes a real model plus Ray/SGLang lifecycle. Production decisions require repeating the matrix on the target topology and workload.

Duplicate check

I searched current open issues and PRs for weight-sync benchmarks, communication-phase benchmarks, policy sweeps, and the exact benchmark filename. No substantive duplicate was found. #2351 and #2352 are related runtime proposals; this PR remains an independent, policy-neutral measurement tool and does not depend on or include either implementation.

AI assistance and required review

This implementation, tests, commit messages, and PR text were prepared with OpenAI Codex assistance. A human contributor/maintainer must review every changed line, validate the synchronization and timing methodology, and reproduce the relevant 2-rank and 4-rank evidence before this draft is considered ready to merge.

kl3-s and others added 3 commits September 4, 2026 20:47
Benchmark trainer-to-engine NCCL broadcast and P2P synchronization across bucket credits, engine waves, and phase strides.

Support arbitrary homogeneous or heterogeneous engine-group layouts and report slowest-rank latency, component timings, inflight bytes, skew, and optional traces.

Assisted-by: OpenAI Codex
Signed-off-by: jiachun <2609024259@qq.com>
Record the maximum number of engine groups present in each generated wave so benchmark results can validate all-at-once, serialized, and windowed policies directly.

Assisted-by: OpenAI Codex

Signed-off-by: jiachun <2609024259@qq.com>
Signed-off-by: 0z5a <dinadegttyarenko6@gmail.com>
@0z5a
0z5a marked this pull request as ready for review September 7, 2026 09:36
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.

2 participants