perf(runtime): upload stacked-tensor shards concurrently - #2292
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughChangesStacked-tensor upload handling
Estimated code review effort: 2 (Simple) | ~10 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
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 |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 80ad527baa
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| return self.alloc_tensor( | ||
| tuple(host.shape[1:]), | ||
| host.dtype, | ||
| init=host[i].contiguous(), | ||
| worker_id=w, |
There was a problem hiding this comment.
When host is not shared or registered, each submitted shard reaches Worker.alloc_tensor, which performs malloc before _prepare_init rejects the host buffer. Because the pool starts all shards at once, an invalid large stacked tensor can transiently allocate every shard across the group and surface backend OOM/free failures instead of the intended shared-memory ValueError; validate the stacked host/shard buffer before submitting the concurrent uploads.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Valid, and confirmed in the code: _require_forked_host_buffer is only reached from _prepare_init, which alloc_device_tensor runs after malloc(nbytes). Serially that cost one shard's malloc before raising; concurrently all N commit device memory first, so a bad 8 × 11.5 GB stacked tensor would surface an allocator failure instead of the shared-memory ValueError.
Fixed in a43d66d by validating the stacked host once before any upload starts:
self._require_forked_host_buffer(host, "alloc_stacked_tensor(host=...)", "read")The check is storage-based and host[i] is a view of the same storage, so validating the stacked tensor covers every shard; alloc_tensor's per-shard check stays as-is. This also matches what copy_stacked_from already does for read-backs.
Verified it does not reject legitimate inputs: a full 8-card DeepSeek V4 Flash W8A8 run (whose stacked weights come in through inherited_host_tensors, not share_memory_()) still uploads clean — 346 GB in 9.3s, generated output unchanged, no errors.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
python/pypto/runtime/distributed_runner.py (1)
1857-1863: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winValidate
hostbefore starting the thread pool.
alloc_stacked_tensorrequires a CPU, contiguous, pre-fork host buffer. The current path checks this only after each task callshost[i].contiguous(). An invalid non-contiguous or non-CPU input can therefore create one temporary copy per shard before all tasks fail.Call
_require_forked_host_buffer(host, "alloc_stacked_tensor(host=...)", "read")before creating the pool. Valid inputs keep the same behavior.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/pypto/runtime/distributed_runner.py` around lines 1857 - 1863, In alloc_stacked_tensor, validate host with _require_forked_host_buffer(host, "alloc_stacked_tensor(host=...)", "read") before creating the thread pool or submitting _upload_shard tasks. Preserve the existing per-shard allocation behavior for valid CPU, contiguous inputs.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/pypto/runtime/distributed_runner.py`:
- Around line 1865-1879: Update the upload flow around uploaded, pool.submit,
and shards.extend so every submitted future is drained and successful tensors
remain rollback-tracked until ownership reaches shards. Handle submission and
future failures, including BaseException, with cleanup in finally; preserve and
re-raise the original upload error after rollback. Remove the inline Ruff noqa
and configure the intentional broad-exception ignore per file instead.
---
Nitpick comments:
In `@python/pypto/runtime/distributed_runner.py`:
- Around line 1857-1863: In alloc_stacked_tensor, validate host with
_require_forked_host_buffer(host, "alloc_stacked_tensor(host=...)", "read")
before creating the thread pool or submitting _upload_shard tasks. Preserve the
existing per-shard allocation behavior for valid CPU, contiguous inputs.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 5e0322e5-da90-4553-85ce-13fc1448ae49
📒 Files selected for processing (1)
python/pypto/runtime/distributed_runner.py
alloc_stacked_tensor uploads shard i to worker i in a serial loop, so a rank-stacked resident weight moves at single-chip H2D bandwidth no matter how many chips the group spans. Each shard targets a different chip worker and nothing orders them, so drive them from a thread pool. Rolling back needs a little more care than the serial loop: a concurrent failure can land anywhere in the group, so the successes are no longer a prefix of ids. Collect them by index and free them against their own worker before re-raising, instead of zipping shards with ids positionally. Measured on 8 x 910B2 uploading DeepSeek V4 Flash W8A8's 346 GB of rank-stacked weights (per-shard 11.5 GB): upload 57.1s -> 10.4s (6.0 -> 33.2 GB/s) startup 82.0s -> 35.1s This needs the matching Simpler change (hw-native-sys/simpler#1702) to pay off: Simpler holds one process-wide lock across the native half of malloc / copy_to, which serializes the group regardless of how the caller issues it. Without that change this commit is a no-op, not a regression.
…upload Two problems the concurrent upload introduced, both from review: `alloc_tensor` only rejects an unreachable host buffer from `_prepare_init`, which runs *after* its device malloc. Serially an invalid host cost one malloc before raising; concurrently every shard commits device memory first, so a bad 8x11.5 GB stacked tensor would surface an allocator failure instead of the intended ValueError. Validate the stacked host once before any upload starts, mirroring what copy_stacked_from already does for read-backs. The collection loop also stopped draining at the first failed future, and `except Exception` let a BaseException (or a failure while submitting) escape with successful uploads still unowned — the rollback below walks `shards`, which is only populated once the whole group succeeded. Submit defensively, drain every future that was submitted, and roll back by index against the owning worker before re-raising the first error. The inline `# noqa: PERF203` is dropped rather than moved to a per-file ignore: neither PERF nor BLE is in this repo's ruff select, so it was suppressing nothing. Re-validated on 8 x 910B2: DeepSeek V4 Flash W8A8 resident upload 9.3s (346 GB), generated output unchanged, no errors; isolated harness threaded 22.7 / 30.7 GB/s vs serial 10.5 / 6.2 GB/s.
a43d66d to
d4ef2a9
Compare
1 participant
What
DistributedWorker.alloc_stacked_tensoruploads shard i to worker i in a serial loop, so a rank-stacked resident weight moves at single-chip H2D bandwidth no matter how many chips the group spans. Each shard targets a different chip worker and nothing orders them, so this drives them from a thread pool.Measured
8 × 910B2, DeepSeek V4 Flash W8A8, 346 GB of rank-stacked resident weights (11.5 GB per shard), warm start:
33 GB/s is what 8 independent processes reach on that node doing a raw cold H2D from a shared mapping, so the upload is now hardware-bound.
Depends on hw-native-sys/simpler#1702
Simpler holds one process-wide lock (
_child_prov_lock) across the native half ofmalloc/copy_to, which serializes the whole group regardless of how the caller issues it — with the serial loop replaced by a thread pool and that lock still in place, the measurement is unchanged (57s). Simpler#1702 narrows it to a per-worker lock.Without that change this commit is a no-op, not a regression — the threads simply queue on the same lock — so the merge order does not matter for correctness.
Rollback semantics
The serial loop could roll back by zipping
shardswithidspositionally, because the successes were always a prefix. A concurrent failure can land anywhere in the group, so successes are collected by index and freed against their own worker before the first error is re-raised. The pre-existingexceptclause remains as the backstop for anything that fails after the group is complete;shardsis either empty or complete when it runs, so nothing is freed twice.Threads only call
alloc_tensor, whosemalloc/copy_tobindings release the GIL, and each targets a distinctworker_id, so no new shared state is introduced on this path.Validation
ruff checkpasses.