perf(comm): coalesce non-Partial chunks by default; Partial stays per-chunk - #109
perf(comm): coalesce non-Partial chunks by default; Partial stays per-chunk#109junjzhang wants to merge 1 commit into
Conversation
|
Warning Review limit reached
More reviews will be available in 58 minutes and 23 seconds. Learn how PR review limits work. Your organization has run out of usage credits. Purchase more in the billing tab. ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans include higher PR review limits than trial, open-source, and free plans. In all cases, reviews become available again over time. During sustained high-volume PR review activity, CodeRabbit may temporarily slow when the next review becomes available. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (6)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
…-chunk Coalescing many small chunks into one wire op is a big win on small tensors (up to 3x) and never worse on large ones. But Partial transfers can't coalesce: their per-cell all_reduce must stay aligned across the reduce group, and send/recv must bucketize symmetrically — a sender-only signal deadlocks or corrupts (see bench/BUCKET_TUNING.md for the two failed attempts). Fix: a Chunk.no_coalesce flag set from M2MMap.source_partial_reductions (the same map on source and target ranks, so both ends agree → symmetric). Partial chunks stay single-entry (per-cell all_reduce, aligned) = their empirical optimum; everything else coalesces. Agent default bucket_size 1 -> 8MB. 16-GPU sweep: bucket >= chunk everywhere, no deadlock, no mismatch. Verified by CPU unit tests + vLLM weight-sync end-to-end (production Partial MoE).
fe94326 to
bf0eecb
Compare
|
Failed to generate code suggestions for PR |
There was a problem hiding this comment.
Pull request overview
This PR changes the tensor-bus communication pipeline to enable chunk coalescing by default for non-Partial transfers while ensuring Partial transfers remain per-chunk (non-coalesced) to preserve correctness and avoid deadlocks/corruption scenarios.
Changes:
- Introduces
Chunk.no_coalesceand sets it symmetrically fromM2MMap.source_partial_reductionsso Partial transfers never coalesce. - Turns on coalescing-by-default in the agent by defaulting
bucket_sizeto an 8MB threshold when unset. - Updates the benchmark to sweep multiple
bucket_sizevalues in a single run and adds a tuning write-up.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| src/etha/tensor_bus/agent.py | Enables default coalescing threshold (8MB) when client does not specify bucket_size. |
| src/etha/comm/ir.py | Adds Chunk.no_coalesce flag to support symmetric “never coalesce” semantics for Partial transfers. |
| src/etha/comm/get_chunks.py | Derives and propagates no_coalesce to all chunks based on M2MMap.source_partial_reductions. |
| src/etha/comm/get_buckets.py | Ensures no_coalesce chunks always become single-entry buckets. |
| bench/transfer_benchmark.py | Adds a bucket_size sweep timing path and adjusts reporting accordingly. |
| bench/BUCKET_TUNING.md | Documents the investigation/tuning process and rationale for the final policy. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # intermediate; only buckets are executed. Partial chunks self-exclude from | ||
| # coalescing (Chunk.no_coalesce); everything else coalesces up to the bytes | ||
| # threshold. | ||
| coalesce_bytes = bucket_size or DEFAULT_COALESCE_BYTES |
| ## M2 — sweep bucket_size (IN PROGRESS) | ||
| `benchmark_single_shape` now sweeps `BUCKET_SIZE_SWEEP = {1, 2MB, 8MB, 32MB, 256MB}` | ||
| in one run (1 = chunk method). Each value reuses the same chunk list; the result | ||
| dict carries `sweep_throughput_gb_s`. (Removed the now-unused per-method profiling | ||
| blocks; profiling args kept but unwired, marked noqa.) |
What did you do
Makes the bucket (coalesced) transfer path never lose to per-chunk across the
benchmark matrix, then turns coalescing on by default.
Investigation (full write-up in
bench/BUCKET_TUNING.md) found two orthogonal rootcauses for why coalescing sometimes lost:
bucket_sizeis wrong for some tensor size; butchunk_to_bucket_opsalready sends any chunk ≥bucket_sizeon its own, so amoderate 8MB threshold coalesces the small chunks (the win) and leaves large ones
alone (no loss).
must stay aligned across the reduce group, and send/recv must bucketize
symmetrically. Two attempts failed (documented): a batched all_reduce deadlocks
(reduce-only per-cell vs shipping batched mismatch sizes), and a sender-only
"don't coalesce Partial" corrupts (P2P send/recv buffers desync).
Fix: a
Chunk.no_coalesceflag set fromM2MMap.source_partial_reductions— thesame map on source and target ranks, so both ends agree and bucketize symmetrically.
Partial chunks stay single-entry (per-cell all_reduce, aligned) = their empirical
optimum; everything else coalesces up to the threshold. Agent default
bucket_size1 → 8MB (
DEFAULT_COALESCE_BYTES).bench/transfer_benchmark.pygains abucket_sizesweep to measure this.New test cases
No new feature/bug; this is a perf + safety change. The Partial-coalescing hazards
are now structurally excluded (
no_coalesce). Covered by the 3-layer verification.Test results
Other comments
Full milestone log of the investigation (including the two reverted dead ends) is in
bench/BUCKET_TUNING.md.