feat(inference): multi-GPU auto-layout and cross-GPU routing (PR2) - #1263
feat(inference): multi-GPU auto-layout and cross-GPU routing (PR2)#1263greenstephen wants to merge 14 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>
Compute VRAM-aware device placement when gpu_mapping=auto on multi-GPU systems, route conditioning tensors to the DiT GPU at inference boundaries, and initialize the LM on its mapped GPU with correct memory budgeting. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR adds a multi-GPU device-mapping package and updates initialization, generation, LLM startup, and UI wiring to resolve and use per-component devices. It also adjusts APG guidance tensor placement to preserve input devices. ChangesMulti-device routing and startup wiring
Estimated code review effort: 4 (Complex) | ~75 minutes Sequence Diagram(s)sequenceDiagram
participant UI
participant initialize_service
participant resolve_component_device_map
participant initialize_llm_at_startup
UI->>initialize_service: gpu_mapping
initialize_service->>resolve_component_device_map: requested_device, gpu_mapping
resolve_component_device_map-->>initialize_service: ComponentDeviceMap
initialize_service->>initialize_llm_at_startup: dit_handler
initialize_llm_at_startup-->>UI: lm_device / startup status
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 (1)
acestep/core/generation/handler/generate_music_decode.py (1)
133-164: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winUse a separate VAE device variable here (
acestep/core/generation/handler/generate_music_decode.py:133-164)
vae_deviceis reset toNonebefore the new VRAM check runs, so the CUDA branch can no longer use the actual VAE target device. Rename the first variable (for example,vae_target_device) and keep the decode target separate from the restore/offload device.🤖 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/generate_music_decode.py` around lines 133 - 164, The VAE target device is being overwritten before the VRAM check, so the CUDA path loses the real decode destination. Update the logic in generate_music_decode’s VAE decode flow to keep the initial device from _get_component_device("vae") in a separate variable (for example, a distinct target device name) and use that for pred_latents_for_decode and the CUDA free-memory check, while reserving the later variable for restore/offload handling. Make sure the using_mlx_vae, vae_cpu, and cuda_device_index/is_cuda_device branches all reference the correct device variable so the decode target and offload target stay separate.
🧹 Nitpick comments (5)
acestep/test_device_map.py (2)
76-81: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueHardcoded env var name instead of imported constant.
"ACESTEP_GPU_MAPPING"is hardcoded here, whileparse_gpu_mapping(per the upstream contract snippet) reads via aGPU_MAPPING_ENVconstant. Importing and using that constant would keep the test coupled to the source of truth instead of a duplicated literal.🤖 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/test_device_map.py` around lines 76 - 81, The test in test_parse_gpu_mapping should stop hardcoding the environment variable name and use the same source-of-truth constant as parse_gpu_mapping. Update the patch.dict setup to reference GPU_MAPPING_ENV, so the test stays aligned with the parser contract and avoids duplicating the literal environment key.
28-176: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTest methods lack mandatory docstrings.
None of the individual test methods in this file have docstrings (only the class docstrings are present), unlike the sibling
init_service_test.pyfile where nearly every test method includes one.As per coding guidelines, "Docstrings are mandatory for all new or modified Python modules, classes, and functions." Concise one-line docstrings (as used throughout
init_service_test.py) would bring this file into compliance.🤖 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/test_device_map.py` around lines 28 - 176, Add concise one-line docstrings to every test method in this file so the modified tests comply with the docstring requirement; update each method in the test classes that uses normalize_component_device, parse_gpu_mapping, resolve_component_device_map, compute_auto_device_map, estimate helpers, and CUDA alias helpers, following the style used in init_service_test.py. Keep the docstrings short and descriptive, placed directly under each test_* method definition.Source: Coding guidelines
acestep/device_map.py (1)
54-64: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
DeviceMapErrorfor consistency with the rest of the module.
device_forraises bareKeyError/ValueErrorfor invalid component names or unassigned devices, while every other validation path in this module (_format_device_for_backend,_parse_mapping_pairs,normalize_component_device, etc.) raisesDeviceMapError. Downstream code that specifically catchesDeviceMapError(the module's documented domain error) won't catch these.♻️ Proposed fix
def device_for(self, component: str) -> str: """Return the resolved device string for a component key.""" key = component.strip().lower() if key == "model": key = "dit" if key not in _COMPONENT_KEYS: - raise KeyError(f"Unknown component: {component}") + raise DeviceMapError(f"Unknown component: {component}") value = getattr(self, key) if value is None: - raise ValueError(f"Component '{component}' has no assigned device") + raise DeviceMapError(f"Component '{component}' has no assigned device") return value🤖 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 54 - 64, `device_for` should use the module’s domain exception instead of raising bare `KeyError` and `ValueError`. Update the validation paths in `DeviceMap.device_for` so invalid component names and missing assignments both raise `DeviceMapError`, matching the behavior used by `_format_device_for_backend`, `_parse_mapping_pairs`, and `normalize_component_device` so callers can consistently catch one exception type.acestep/core/generation/handler/conditioning_embed.py (1)
9-16: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocstring dependency list omits
_get_component_device.Similar to
audio_codes.py, the "Depends on host members" contract doesn't list_get_component_device, now used at Lines 59-61, 70, and 133.As per coding guidelines: "Docstrings must be concise and include purpose plus key inputs/outputs and raised exceptions when relevant."
🤖 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/conditioning_embed.py` around lines 9 - 16, The docstring for ConditioningEmbedMixin has an incomplete “Depends on host members” contract because it omits _get_component_device, which is used by the mixin methods. Update the class docstring in ConditioningEmbedMixin to include _get_component_device alongside the other required methods, keeping the description concise and aligned with the existing dependency list.acestep/core/generation/handler/audio_codes.py (1)
11-18: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocstring dependency list omits
_get_component_device.The mixin's "Depends on host members" contract lists methods it relies on but doesn't mention
_get_component_device, now used at Lines 59-60 and 87-91.As per coding guidelines: "Docstrings must be concise and include purpose plus key inputs/outputs and raised exceptions when relevant."
🤖 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/audio_codes.py` around lines 11 - 18, The AudioCodesMixin docstring’s host-dependency contract is missing the `_get_component_device` method even though `AudioCodesMixin` uses it in its audio parsing and latent conversion flow. Update the class docstring in `AudioCodesMixin` so the “Depends on host members” list includes `_get_component_device` alongside the existing required methods, keeping the description concise and aligned with the mixin’s actual dependencies.
🤖 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.py`:
- Around line 1-425: This module is over the 200-LOC hard cap and should be
split by responsibility. Move the GPU/data types and small helpers like GpuInfo,
ComponentDeviceMap, DeviceMapError, and normalization/formatting utilities into
a types/parsing module, and move auto-layout logic such as LayoutRequest,
LayoutError, estimate_dit_peak_gb, estimate_lm_total_gb,
compute_auto_device_map, and discover_gpus into a layout module. Keep
resolve_component_device_map and log_device_map in a thin device_map facade that
re-exports the public API so callers of ComponentDeviceMap, parse_gpu_mapping,
and resolve_component_device_map do not need to change.
- Around line 158-203: The fallback selection in compute_auto_device_map
currently lets LM co-locate on dit_gpu based only on lm_need_gb, which can OOM
when the DiT already consumes VRAM. Update compute_auto_device_map so the LM
candidate search in the request.use_lm branch first prefers a different GPU, and
only falls back to dit_gpu if that device has enough free_vram_gb for both
estimate_dit_peak_gb(request.dit_type, request.batch_size) and
estimate_lm_total_gb(request.lm_model_path) combined; keep the
ComponentDeviceMap and LayoutError behavior unchanged otherwise.
---
Outside diff comments:
In `@acestep/core/generation/handler/generate_music_decode.py`:
- Around line 133-164: The VAE target device is being overwritten before the
VRAM check, so the CUDA path loses the real decode destination. Update the logic
in generate_music_decode’s VAE decode flow to keep the initial device from
_get_component_device("vae") in a separate variable (for example, a distinct
target device name) and use that for pred_latents_for_decode and the CUDA
free-memory check, while reserving the later variable for restore/offload
handling. Make sure the using_mlx_vae, vae_cpu, and
cuda_device_index/is_cuda_device branches all reference the correct device
variable so the decode target and offload target stay separate.
---
Nitpick comments:
In `@acestep/core/generation/handler/audio_codes.py`:
- Around line 11-18: The AudioCodesMixin docstring’s host-dependency contract is
missing the `_get_component_device` method even though `AudioCodesMixin` uses it
in its audio parsing and latent conversion flow. Update the class docstring in
`AudioCodesMixin` so the “Depends on host members” list includes
`_get_component_device` alongside the existing required methods, keeping the
description concise and aligned with the mixin’s actual dependencies.
In `@acestep/core/generation/handler/conditioning_embed.py`:
- Around line 9-16: The docstring for ConditioningEmbedMixin has an incomplete
“Depends on host members” contract because it omits _get_component_device, which
is used by the mixin methods. Update the class docstring in
ConditioningEmbedMixin to include _get_component_device alongside the other
required methods, keeping the description concise and aligned with the existing
dependency list.
In `@acestep/device_map.py`:
- Around line 54-64: `device_for` should use the module’s domain exception
instead of raising bare `KeyError` and `ValueError`. Update the validation paths
in `DeviceMap.device_for` so invalid component names and missing assignments
both raise `DeviceMapError`, matching the behavior used by
`_format_device_for_backend`, `_parse_mapping_pairs`, and
`normalize_component_device` so callers can consistently catch one exception
type.
In `@acestep/test_device_map.py`:
- Around line 76-81: The test in test_parse_gpu_mapping should stop hardcoding
the environment variable name and use the same source-of-truth constant as
parse_gpu_mapping. Update the patch.dict setup to reference GPU_MAPPING_ENV, so
the test stays aligned with the parser contract and avoids duplicating the
literal environment key.
- Around line 28-176: Add concise one-line docstrings to every test method in
this file so the modified tests comply with the docstring requirement; update
each method in the test classes that uses normalize_component_device,
parse_gpu_mapping, resolve_component_device_map, compute_auto_device_map,
estimate helpers, and CUDA alias helpers, following the style used in
init_service_test.py. Keep the docstrings short and descriptive, placed directly
under each test_* method definition.
🪄 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: d997cfaf-8d50-442e-92fa-345fa8fee757
📒 Files selected for processing (17)
acestep/api/startup_llm_init.pyacestep/api/startup_model_init.pyacestep/core/generation/handler/audio_codes.pyacestep/core/generation/handler/conditioning_embed.pyacestep/core/generation/handler/generate_music_decode.pyacestep/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/core/generation/handler/service_generate_execute.pyacestep/device_map.pyacestep/gpu_config.pyacestep/llm_inference.pyacestep/test_device_map.pyacestep/ui/gradio/events/generation/service_init.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>
…ndex. Route generate_kwargs tensors to the DiT device, move latents during preprocess, and preserve CUDA device indices in APG project(). Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
acestep/device_map/discovery.py (1)
21-24: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNarrow the exception type flagged by Ruff (BLE001).
except Exceptionis a blind catch-all; aRuntimeError(or narrower torch-specific exception) would satisfy the "catch specific exceptions" guideline while preserving the safeNonefallback.♻️ Proposed narrowing
try: capability = torch.cuda.get_device_capability(index) - except Exception: + 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/discovery.py` around lines 21 - 24, The exception handling in the device capability lookup is too broad and triggers Ruff BLE001. In the function that calls torch.cuda.get_device_capability, replace the blanket except Exception with a narrower exception type such as RuntimeError or the most specific torch-related exception that can be raised here, while keeping the existing fallback of setting capability to None.Source: Linters/SAST tools
acestep/models/common/apg_guidance_test.py (1)
12-59: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the MPS device-restoration branch.
Tests cover CPU and non-default CUDA-index device preservation, but not the
device.type == "mps"branch inproject()— the exact path with the device-restoration bug flagged inapg_guidance.py. Since real MPS hardware likely isn't available in CI, consider mockingtorch.device.type/monkeypatching to exercise that branch and assert the returned tensors land back on the original (mocked) MPS device.🤖 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/models/common/apg_guidance_test.py` around lines 12 - 59, The device-preservation tests in ApgGuidanceDeviceTests currently miss the mps-specific restoration path inside project(), which is the branch tied to the reported bug. Add a test that exercises the device.type == "mps" logic by mocking or monkeypatching the device check in project() so the branch runs without real MPS hardware. Verify that both returned tensors from project() are restored to the original mocked MPS device, similar to the existing CPU and cuda:1 assertions.
🤖 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_memory_basic.py`:
- Around line 154-163: The `_ensure_silence_latent_on_device` method is using
`_is_on_target_device`, which only checks device type and can miss GPU index
mismatches when `device_map.dit` points to a different CUDA device. Update the
device validation at this call site to compare the full normalized target device
for `silence_latent` against the DiT device (or `self.device` when no
`device_map` is set), and only move it when the exact device index does not
match.
In `@acestep/device_map/layout.py`:
- Around line 79-93: The LM placement logic in request.use_lm can still choose
dit_gpu even when its remaining VRAM after DiT allocation is insufficient.
Update the lm_gpu selection in layout.py to account for the DiT reservation when
evaluating colocated placement, using the LM budget helper or an equivalent
remaining-capacity check with reserve_dit_inference_gb set appropriately. Keep
the fix localized around estimate_lm_total_gb, lm_candidates, and the lm_gpu
selection so the fallback to dit_gpu only succeeds when the post-DiT headroom is
actually enough.
In `@acestep/models/common/apg_guidance.py`:
- Around line 21-34: The tensor projection helper in apg_guidance currently
overwrites the original MPS device after the CPU fallback, so the return tensors
stay on CPU instead of being moved back. Update the device handling in the
projection function that computes v0_parallel and v0_orthogonal so it preserves
the original device before any .cpu() fallback, then uses that saved device in
the final .to(...) calls to return tensors to MPS correctly.
---
Nitpick comments:
In `@acestep/device_map/discovery.py`:
- Around line 21-24: The exception handling in the device capability lookup is
too broad and triggers Ruff BLE001. In the function that calls
torch.cuda.get_device_capability, replace the blanket except Exception with a
narrower exception type such as RuntimeError or the most specific torch-related
exception that can be raised here, while keeping the existing fallback of
setting capability to None.
In `@acestep/models/common/apg_guidance_test.py`:
- Around line 12-59: The device-preservation tests in ApgGuidanceDeviceTests
currently miss the mps-specific restoration path inside project(), which is the
branch tied to the reported bug. Add a test that exercises the device.type ==
"mps" logic by mocking or monkeypatching the device check in project() so the
branch runs without real MPS hardware. Verify that both returned tensors from
project() are restored to the original mocked MPS device, similar to the
existing CPU and cuda:1 assertions.
🪄 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: 6b980c49-a248-448c-b2c6-b125501c0592
📒 Files selected for processing (18)
acestep/core/generation/handler/conditioning_embed.pyacestep/core/generation/handler/init_service_memory_basic.pyacestep/core/generation/handler/init_service_setup.pyacestep/core/generation/handler/init_service_test.pyacestep/core/generation/handler/service_generate_execute.pyacestep/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.pyacestep/models/common/apg_guidance.pyacestep/models/common/apg_guidance_test.pyacestep/test_device_map.py
✅ Files skipped from review due to trivial changes (2)
- acestep/device_map/errors.py
- acestep/device_map/init.py
🚧 Files skipped from review as they are similar to previous changes (4)
- acestep/core/generation/handler/service_generate_execute.py
- acestep/core/generation/handler/conditioning_embed.py
- acestep/core/generation/handler/init_service_setup.py
- acestep/test_device_map.py
…ack. Also narrow get_device_capability exception handling to RuntimeError. Co-authored-by: Cursor <cursoragent@cursor.com>
Reject unsupported default_device backends explicitly in GPU mapping parsing, and preserve the original MPS device when project() uses CPU math fallback. Co-authored-by: Cursor <cursoragent@cursor.com>
_is_on_target_device only compares backend types, so cuda:0 and cuda:3 look equivalent. Use exact device matching when placing silence_latent on the DiT device from device_map. Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
🧹 Nitpick comments (1)
acestep/core/generation/handler/init_service_memory_basic.py (1)
114-125: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winNarrow the blind
except Exception.Ruff flags this as BLE001.
normalize_component_deviceonly raisesDeviceMapError(on empty string), andtorch.device(...)raisesRuntimeError/TypeErroron invalid strings — catching those specifically avoids silently swallowing unrelated bugs.As per coding guidelines, "Error handling: Avoid bare
except:clauses; catch specific exceptions."♻️ Proposed fix
+from acestep.device_map.errors import DeviceMapError + def _tensor_on_exact_device(self, tensor, target_device: str) -> bool: """Return whether *tensor* is on the exact device string (including CUDA index).""" if tensor is None: return True from acestep.device_map.devices import normalize_component_device try: expected = torch.device(normalize_component_device(str(target_device))) - except Exception: + except (DeviceMapError, RuntimeError, TypeError): return False return tensor.device == expected🤖 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_memory_basic.py` around lines 114 - 125, The `_tensor_on_exact_device` helper in `init_service_memory_basic.py` is catching `Exception` too broadly, which triggers BLE001. Update the `try`/`except` around `normalize_component_device` and `torch.device(...)` to catch only the expected failures from those calls, using the specific exception types they can raise, and leave other errors unhandled so real bugs are not hidden.Sources: Coding guidelines, Linters/SAST tools
🤖 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.
Nitpick comments:
In `@acestep/core/generation/handler/init_service_memory_basic.py`:
- Around line 114-125: The `_tensor_on_exact_device` helper in
`init_service_memory_basic.py` is catching `Exception` too broadly, which
triggers BLE001. Update the `try`/`except` around `normalize_component_device`
and `torch.device(...)` to catch only the expected failures from those calls,
using the specific exception types they can raise, and leave other errors
unhandled so real bugs are not hidden.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: b5b2f4bf-cedd-4eac-bc44-25b891fab657
📒 Files selected for processing (2)
acestep/core/generation/handler/init_service_memory_basic.pyacestep/core/generation/handler/init_service_test.py
Module-level import avoids per-call overhead; catch DeviceMapError and RuntimeError instead of a broad Exception in _tensor_on_exact_device. Co-authored-by: Cursor <cursoragent@cursor.com>
Replace bare Exception handlers with DeviceMapError/RuntimeError/TypeError for exact device checks and RuntimeError/TypeError/ValueError for backend alias parsing in _is_on_target_device. Co-authored-by: Cursor <cursoragent@cursor.com>
vae_component_device was overwritten with None before the CUDA free-VRAM check, so multi-GPU layouts always queried cuda:0. Keep a separate vae_restore_device for the CPU offload restore path. 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>
Move shared fixtures into generate_music_decode_test_support.py and split scenario tests into prepare, basic, CPU-offload, and VRAM modules. Co-authored-by: Cursor <cursoragent@cursor.com>
Drop the custom sys.modules package stubbing loader; import the mixin through the public package path like other handler tests. Co-authored-by: Cursor <cursoragent@cursor.com>
|
So, just to put some comments specific to this PR here, the code in llm_inference.py only respects the LM device setting if the device setting is "cuda" as opposed to "auto". Additionally, it seems like the backend needs to be set to PT since VLLM doesn't seem to be respecting the setting and uses the same GPU for all models. Let me know if you have other things you want me to try. |
|
Thanks — good catches, and useful to keep them on this PR. auto vs cuda: Agreed. In llm_inference.initialize, device == "auto" is collapsed to bare "cuda" (no index), so a mapped cuda:1 never sticks on that path. Related: several later checks use device == "cuda" / device != "cuda", which also break indexed devices. One concrete footgun is the vLLM gate: if backend == "vllm" and device != "cuda": vLLM same-GPU behavior: That matches what we’re seeing. Until indexed CUDA is handled cleanly in the vLLM init path, --backend pt is the reliable way to keep the LM on lm:N. We’ll treat proper vLLM + cuda:N as part of the same fix. Things worth trying: --backend pt with --gpu-mapping "dit:0,vae:0,text_encoder:0,lm:1" and --init-service true |
Resolve Gradio lm_device after DiT initialize_service so device_map.lm is applied, normalize/preserve indexed CUDA devices in llm_inference, and move text-encoder token ids onto the text-encoder component device. Co-authored-by: Cursor <cursoragent@cursor.com>
|
@pokepress Thanks again — you were right about the LM landing on bare I’ve pushed a fix on this branch (
Please retry (after pulling this branch): uv run acestep \
--init-service true \
--config-path acestep-v15-xl-sft \
--lm_model_path acestep-5Hz-lm-4B \
--init-llm true \
--gpu-mapping "dit:0,vae:0,text_encoder:0,lm:1" |
3 participants
Summary
gpu_mapping=autolayout for multi-GPU CUDA systems (DiT stack on one GPU, LM on another).device_map.lmwith per-GPU vLLM memory budgeting.Stacked on
a38527f).Scope
device_map.pyauto-layout, cross-GPU tensor routing, LM device wiring, tests--gpu-mappingflag, Gradio UI panel, LM tensor parallelism (PR3–PR4)Risk and Compatibility
ACESTEP_GPU_MAPPING=autoor explicit multi-component mapsRegression Checks
acestep.test_device_map— 18 tests passacestep.core.generation.handler.init_service_test— 83/84 pass (1 pre-existing failure)Test plan
ACESTEP_GPU_MAPPING=autoon 2×24GB — XL DiT + 4B LM init without OOMACESTEP_GPU_MAPPING=dit:0,vae:0,text_encoder:0,lm:1— generate a short trackMade with Cursor
Summary by CodeRabbit