feat(inference): add ComponentDeviceMap for multi-GPU placement (PR1) - #1262
feat(inference): add ComponentDeviceMap for multi-GPU placement (PR1)#1262greenstephen wants to merge 5 commits into
Conversation
Introduce per-component device mapping so DiT, VAE, text encoder, and LM can be assigned to distinct CUDA devices while preserving single-GPU behavior when no mapping is provided. Co-authored-by: Cursor <cursoragent@cursor.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughAdds a new ChangesComponent device mapping
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant InitServiceOrchestratorMixin
participant InitServiceSetupMixin
participant device_map
Caller->>InitServiceOrchestratorMixin: initialize_service(gpu_mapping)
InitServiceOrchestratorMixin->>InitServiceSetupMixin: _resolve_component_device_map(resolved_device, gpu_mapping)
InitServiceSetupMixin->>device_map: resolve_component_device_map(...)
device_map->>device_map: discover_gpus() / compute_auto_device_map() or parse_gpu_mapping()
device_map-->>InitServiceSetupMixin: ComponentDeviceMap
InitServiceSetupMixin-->>InitServiceOrchestratorMixin: self.device_map
InitServiceOrchestratorMixin->>device_map: set_active_cuda_device(self.device)
InitServiceOrchestratorMixin->>InitServiceOrchestratorMixin: load model/vae/text_encoder using device_map devices
InitServiceOrchestratorMixin-->>Caller: status text + last_init_params with gpu mapping
Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
acestep/device_map.py (1)
1-279: 📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy liftModule exceeds the 200 LOC hard cap.
device_map.pyspans 279 lines, over the guideline's hard cap for module size. Consider splitting by responsibility (e.g., dataclasses/normalization helpers in one module, mapping-string parsing in another, GPU discovery/logging in a third) while re-exporting the public names fromdevice_map.pyto preserve the stable facade for existing imports.As per coding guidelines, "Module/API Readiness: Module LOC policy is met (<=150 target, <=200 hard cap or justified exception)" and "If a module would exceed 200 LOC, split by responsibility before merging, or add a short justification in PR notes with a concrete follow-up split plan." Based on learnings, module-size concerns should only be raised once a file exceeds 200 LOC, which is the case here.
🤖 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 `@acestep/device_map.py` around lines 1 - 279, Split device_map.py because it exceeds the module hard cap; move related responsibilities into smaller modules such as device dataclasses/normalization helpers, GPU mapping parsing/resolution, and GPU discovery/logging, then re-export the public API from device_map.py so existing imports keep working. Keep the stable symbols like GpuInfo, ComponentDeviceMap, parse_gpu_mapping, resolve_component_device_map, discover_gpus, and log_device_map available from the facade module after the split.Sources: Coding guidelines, Learnings
acestep/core/generation/handler/init_service_orchestrator.py (1)
94-111: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winUse the text encoder’s device when choosing dtype —
_load_text_encoder_and_tokenizer()still casts withself.dtype, so a multi-GPUgpu_mappingcan place it on a device that doesn’t support that precision. Match_load_vae_model()and resolve dtype fromdevicethere too.🤖 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 `@acestep/core/generation/handler/init_service_orchestrator.py` around lines 94 - 111, The dtype selection in initialize_service_orchestrator currently uses resolved_device, but _load_text_encoder_and_tokenizer() still relies on self.dtype, which can mismatch the text encoder’s actual device under multi-GPU gpu_mapping. Update the dtype resolution logic in the service init flow to use the same device-aware approach as _load_vae_model(), resolving dtype from the text encoder’s device rather than a shared self.dtype. Keep the existing ROCm/CUDA/xpu branching, but ensure the final dtype is derived from the device that will load the text encoder and tokenizer so precision matches the target device.
🧹 Nitpick comments (2)
acestep/device_map.py (2)
142-166: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueNarrow the blind
except Exceptionwhen probing compute capability.Ruff flags this as BLE001. Catching
Exceptionhere can mask unrelated failures (e.g., programming errors) silently as "no compute capability".As per coding guidelines, "Error handling: Avoid bare `except:` clauses; catch specific exceptions."🛡️ Narrow the exception type
- try: - capability = torch.cuda.get_device_capability(index) - except Exception: - capability = None + try: + capability = torch.cuda.get_device_capability(index) + except RuntimeError: + capability = None🤖 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 `@acestep/device_map.py` around lines 142 - 166, The compute-capability probe in discover_gpus() is catching Exception too broadly, which can hide unrelated bugs. Update the try/except around torch.cuda.get_device_capability(index) to catch only the specific expected CUDA-related error type(s) for this lookup, and keep the fallback to None only for that case while allowing other exceptions to surface. Use discover_gpus() and the compute_capability assignment in GpuInfo as the target points.Sources: Coding guidelines, Linters/SAST tools
205-252: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winConsider validating mapping indices against discovered GPU count.
parse_gpu_mapping/_format_device_for_backendnever check the requested index againsttorch.cuda.device_count()(ordiscover_gpus()), so an out-of-range index (e.g.dit:5on a 2-GPU box) only surfaces later as a deep runtime error during model loading instead of a clear, earlyDeviceMapError.🤖 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 `@acestep/device_map.py` around lines 205 - 252, parse_gpu_mapping currently accepts GPU indices without checking them against the actual available devices, so invalid entries can fail much later during model load. Add an early validation step in parse_gpu_mapping (or inside _format_device_for_backend) that compares each requested index against the discovered GPU count for the selected backend, and raise a DeviceMapError when an index is out of range. Make sure this covers both the single-device path and the component-pair path for dit, vae, text_encoder, and lm.
🤖 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 `@acestep/core/generation/handler/init_service_orchestrator.py`:
- Around line 85-86: Guard the CUDA device switch inside initialize_service so
it does not mutate process-wide state during re-entrant or overlapping calls.
Update the init_service_orchestrator flow around is_cuda_device and
set_active_cuda_device to keep device selection local to the load path or
serialize that section, so concurrent API/UI entrypoints do not redirect work to
the wrong GPU.
In `@acestep/device_map.py`:
- Around line 169-178: The XPU branch in _format_device_for_backend currently
returns xpu:0, which breaks exact device checks that expect bare xpu. Update
_format_device_for_backend so the xpu backend normalizes to xpu for component
mapping just like the other non-indexed backends, while still keeping the
existing validation for nonzero indices and leaving the cuda path unchanged.
---
Outside diff comments:
In `@acestep/core/generation/handler/init_service_orchestrator.py`:
- Around line 94-111: The dtype selection in initialize_service_orchestrator
currently uses resolved_device, but _load_text_encoder_and_tokenizer() still
relies on self.dtype, which can mismatch the text encoder’s actual device under
multi-GPU gpu_mapping. Update the dtype resolution logic in the service init
flow to use the same device-aware approach as _load_vae_model(), resolving dtype
from the text encoder’s device rather than a shared self.dtype. Keep the
existing ROCm/CUDA/xpu branching, but ensure the final dtype is derived from the
device that will load the text encoder and tokenizer so precision matches the
target device.
In `@acestep/device_map.py`:
- Around line 1-279: Split device_map.py because it exceeds the module hard cap;
move related responsibilities into smaller modules such as device
dataclasses/normalization helpers, GPU mapping parsing/resolution, and GPU
discovery/logging, then re-export the public API from device_map.py so existing
imports keep working. Keep the stable symbols like GpuInfo, ComponentDeviceMap,
parse_gpu_mapping, resolve_component_device_map, discover_gpus, and
log_device_map available from the facade module after the split.
---
Nitpick comments:
In `@acestep/device_map.py`:
- Around line 142-166: The compute-capability probe in discover_gpus() is
catching Exception too broadly, which can hide unrelated bugs. Update the
try/except around torch.cuda.get_device_capability(index) to catch only the
specific expected CUDA-related error type(s) for this lookup, and keep the
fallback to None only for that case while allowing other exceptions to surface.
Use discover_gpus() and the compute_capability assignment in GpuInfo as the
target points.
- Around line 205-252: parse_gpu_mapping currently accepts GPU indices without
checking them against the actual available devices, so invalid entries can fail
much later during model load. Add an early validation step in parse_gpu_mapping
(or inside _format_device_for_backend) that compares each requested index
against the discovered GPU count for the selected backend, and raise a
DeviceMapError when an index is out of range. Make sure this covers both the
single-device path and the component-pair path for dit, vae, text_encoder, and
lm.
🪄 Autofix (Beta)
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
Run ID: a1a67daf-6734-4569-9972-e6e426ca619b
📒 Files selected for processing (8)
acestep/core/generation/handler/init_service_catalog.pyacestep/core/generation/handler/init_service_loader.pyacestep/core/generation/handler/init_service_offload_context.pyacestep/core/generation/handler/init_service_orchestrator.pyacestep/core/generation/handler/init_service_setup.pyacestep/core/generation/handler/init_service_test.pyacestep/device_map.pyacestep/test_device_map.py
Move GPU discovery, layout, parsing, resolution, and status helpers into focused submodules under acestep/device_map/ with a stable public facade. Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
acestep/device_map/parsing.py (1)
99-104: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate raw-mapping resolution logic.
The "explicit value or env fallback, stripped" logic here duplicates lines 60-62 in
parse_gpu_mapping. Extracting a shared helper would prevent the two from drifting apart.♻️ Proposed refactor
+def _resolve_raw_mapping(mapping: Optional[str]) -> str: + raw = (mapping or "").strip() + if not raw: + raw = os.environ.get(GPU_MAPPING_ENV, "").strip() + return raw + + def parse_gpu_mapping( mapping: Optional[str], *, default_device: str, ) -> Optional[ComponentDeviceMap]: ... - raw = (mapping or "").strip() - if not raw: - raw = os.environ.get(GPU_MAPPING_ENV, "").strip() + raw = _resolve_raw_mapping(mapping) if not raw or raw.lower() == "auto": return None ... def raw_gpu_mapping_value(gpu_mapping: Optional[str]) -> str: """Return the effective raw mapping string from args or environment.""" - raw = (gpu_mapping or "").strip() - if not raw: - raw = os.environ.get(GPU_MAPPING_ENV, "").strip() - return raw + return _resolve_raw_mapping(gpu_mapping)🤖 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 `@acestep/device_map/parsing.py` around lines 99 - 104, The raw mapping fallback-and-strip logic is duplicated between parse_gpu_mapping and raw_gpu_mapping_value, so extract that shared behavior into a single helper and have both callers use it. Update the raw_gpu_mapping_value function and the relevant logic in parse_gpu_mapping to delegate to the shared helper so the explicit value-or-environment fallback stays consistent and cannot drift.
🤖 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 `@acestep/device_map/devices.py`:
- Around line 13-20: The cuda_device_index helper currently lets malformed
“cuda:” strings bubble up as raw ValueError from the int conversion instead of
using DeviceMapError. Update cuda_device_index to validate or catch parsing
failures for normalized.startswith("cuda:") inputs, and re-raise them as
DeviceMapError with the existing domain error style so all invalid CUDA device
strings are handled consistently.
In `@acestep/device_map/discovery.py`:
- Around line 21-24: The discovery logic is swallowing too broadly in the device
capability probe; narrow the `except Exception` in `get_device_capability` to
the specific CUDA failure cases expected from
`torch.cuda.get_device_capability`, and add a `loguru.logger` warning/error when
the lookup fails so invalid devices are visible without hiding unrelated bugs.
Keep the fallback to `None` in the same `discovery.py` code path, but make the
failure handling explicit and tied to the `get_device_capability(index)` call.
In `@acestep/device_map/parsing.py`:
- Around line 66-68: The fallback branch in device backend parsing is a no-op
because the unknown-device path always returns the same value already stored in
backend. Update the logic in parsing.py around device_type(default_device) and
is_cuda_device(default_device) so unrecognized device strings are handled
explicitly instead of using a redundant ternary; either remove the dead branch
or replace it with real fallback/validation behavior in the backend selection
flow.
---
Nitpick comments:
In `@acestep/device_map/parsing.py`:
- Around line 99-104: The raw mapping fallback-and-strip logic is duplicated
between parse_gpu_mapping and raw_gpu_mapping_value, so extract that shared
behavior into a single helper and have both callers use it. Update the
raw_gpu_mapping_value function and the relevant logic in parse_gpu_mapping to
delegate to the shared helper so the explicit value-or-environment fallback
stays consistent and cannot drift.
🪄 Autofix (Beta)
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
Run ID: c207a7b9-885b-44b1-ab99-ceb1c5fb492e
📒 Files selected for processing (10)
acestep/device_map/__init__.pyacestep/device_map/constants.pyacestep/device_map/devices.pyacestep/device_map/discovery.pyacestep/device_map/errors.pyacestep/device_map/layout.pyacestep/device_map/parsing.pyacestep/device_map/resolve.pyacestep/device_map/status.pyacestep/device_map/types.py
✅ Files skipped from review due to trivial changes (1)
- acestep/device_map/errors.py
…ack. Also narrow get_device_capability exception handling to RuntimeError. Co-authored-by: Cursor <cursoragent@cursor.com>
…sing. Replace the no-op backend fallback with an explicit DeviceMapError. Co-authored-by: Cursor <cursoragent@cursor.com>
Wrap cuda_device_index int() parsing so strings like cuda:x raise the module's domain error instead of a raw ValueError. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Ready for merge. Stack order: #1262 → #1263 → #1264. Tested locally on 4× RTX 3090 with: All CodeRabbit threads resolved; CI green. Please merge #1262 first, then #1263, then #1264. |
|
@coderabbitai resume |
✅ Action performedReviews resumed. |
|
So, I've been test driving this (using the 3rd PR branch), with the following setup: uv run acestep --gpu_mapping dit:0,vae:0,text_encoder:0,lm:1 Primary GPU: 4060 ti 16GB What should I be able to do that I couldn't do before? Seems like I should be able to run an XL model and the 4B LM, but maybe I'm wrong. I've tried using that combination and the LLM still gets put on the same GPU as everything else. I saw the message:
Based on the code, that should be cuda:1, so perhaps the ID isn't making it there. I did a bit more digging, and it looks like part of the issue is that the initialization function in llm_inference.py uses a device of "cuda" without a number if the device is set to "auto" (the default Gradio setting): So that should probably be tweaked in PR2, as should several similar statements. Also worth noting that some lines later in that function compare the current device to "cuda", and won't work correctly if it has a device number appended. |
|
Thanks for the detailed dig — this is very helpful, and you’re right on both counts. What you should be able to do Why you’re seeing Loading LLM to cuda Two related issues:
So: your expectation (XL + 4B) is right according to the model size (however in my experience the very largest model is bigger than 16gb so you may need to go down a bit or offload); the placement bug is why it didn’t happen. Workaround until fixed uv run acestep [device_map] Active layout: dit:0, vae:0, text_encoder:0, lm:1 Resolve lm_device from device_map after DiT initialize_service in Gradio service_init. |
|
Just to update you, I was able to get the LM to go to the other GPU by:
The 4B LM does fit by itself on a 12GB GPU (maybe a 10GB one as well). I'll see if I can get the newest changes a try soon. |
|
pokepress, Im glad you were able to get it working that way. I found that the vllm implementation essentially hardcoded to gpu0. I fixed that and some other mapping issues and pushed a fix to my repo. I havent submitted a pull request, but will if they are interested in the feature. |
3 participants
Summary
ComponentDeviceMapand GPU mapping parsing (single:N, explicitdit:0,vae:0,text_encoder:0,lm:1) as the foundation for multi-GPU inference (Only using 1 of multiple GPUs available #426).cuda:Ndevice indices and use the correct CUDA device for flash-attention and dtype selection.Scope
acestep/device_map.py, handler init/load/offload paths, unit testsgpu_mapping=auto), cross-GPU inference routing, CLI/API flags, Gradio UI, LM handler wiring (PR2–PR4)Risk and Compatibility
Regression Checks
acestep.test_device_map— 13 tests passacestep.core.generation.handler.init_service_test— 77/78 pass (1 pre-existing failure:_sync_alignment_configmissing on test host)gpu_mappingis unsetReviewer Notes
Test plan
device=cuda:1andgpu_mapping=single:1— all components load on GPU 1gpu_mapping=dit:0,vae:0,text_encoder:0,lm:1— components load on mapped devices; startup log shows layoutMade with Cursor
Summary by CodeRabbit
cuda:1).