Repository navigation
Layer split: size device ranges by capacity and keep compute headroom under --max-vram - #32
danielhanchen wants to merge 1 commit into
Conversation
…mpute headroom under --max-vram The partitioner filled the first device up to its weight capacity before using the next one, so a module whose weights fit the main device alone was never split, and the compute buffers then had no room under the budget. The 2 GiB compute headroom was also taken only from free memory, never from an explicit --max-vram budget, so a per-device budget was filled entirely with weights. Ranges are now sized in proportion to each device's capacity (the lower of free memory and its --max-vram budget, minus the headroom), as docs/backend.md already described, with the fill-first plan kept as a fallback when the proportional ranges do not fit. The planner is a header-only function with a standalone test (scripts/layer_split_ranges_test.cpp).
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![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: 8c56b2b0c9
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| std::vector<double> caps(capacities.size(), 0.0); | ||
| for (size_t i = 0; i < capacities.size(); i++) { | ||
| // An unknown capacity (very large sentinel) counts as able to hold everything. | ||
| caps[i] = capacities[i] > 0 ? std::min((double)capacities[i], total) : 0.0; |
There was a problem hiding this comment.
When the model's total weight size is smaller than each device's capacity, clamping every capacity to total erases all capacity differences and produces an equal split. For example, 4 GB of segments with 4 GB and 12 GB capacities are assigned 2 GB to each device rather than roughly 1 GB/3 GB, placing twice the intended blocks and their compute workload on the smaller device; this defeats the new headroom strategy and can cause that device to run out of compute memory. Only the unknown-capacity sentinel should be normalized without discarding ratios between known capacities.
Useful? React with 👍 / 👎.
| @@ -0,0 +1,86 @@ | |||
| // Standalone check of the layer-split range planner: g++ -std=c++17 -I src scripts/layer_split_ranges_test.cpp && ./a.out | |||
There was a problem hiding this comment.
Remove this standalone test from the commit: CONTRIBUTING.md explicitly requires test code and test scripts to remain local and only their verification results to be reported, while the repository's agent workflow makes those PR expectations applicable here. Keeping this file in scripts/ violates that submission boundary even though it is not wired into the build.
AGENTS.md reference: AGENTS.md:L120-L122
Useful? React with 👍 / 👎.
This branch has not been deployed
1 participant
Summary
With
--backend "diffusion=cuda0&cuda1" --split-mode layer, the layer-split partitioner filled the first device up to its weight capacity before assigning anything to the next one. A DiT whose weights fit device 0 was therefore never split, and its compute buffers then had no room. Passing a per-device--max-vramdid force a split, but device 0 was still filled to its whole budget with weights and the run failed when the compute buffer was added.This change sizes the device ranges in proportion to each device's capacity and keeps compute headroom inside each device's
--max-vrambudget. The range computation lives in a small header (layer_split_ranges.h) so it can be tested on its own.Results
Kaggle T4x2 (two 15 GB T4s), Z-Image Q8_0, 768x768, 4 steps, same VM and inputs for master and this branch:
--max-vram cuda0=5,cuda1=12Every successful run is pixel-identical to the same binary's single-device render. Speed is unchanged (about 2.1 to 2.3 s/it everywhere): the gain is capacity, not throughput, since layer split runs one device at a time.
Not covered: row split, three or more devices, unequal devices (unit-level only), MiniMax-H3 (the run used Z-Image).