feat: multi-GPU component placement (DiT / 5Hz LM on separate GPUs) for Gradio + CLI - #1269
Open
Doud-FR wants to merge 5 commits into
Open
feat: multi-GPU component placement (DiT / 5Hz LM on separate GPUs) for Gradio + CLI#1269Doud-FR wants to merge 5 commits into
Doud-FR wants to merge 5 commits into
Conversation
Contributor
|
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:
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 (1)
📝 WalkthroughWalkthroughThe changes add per-component device mapping with environment overrides and free-VRAM ranking, apply DiT/VAE placement during service initialization, and route LM initialization through mapped devices with adjusted offload and indexed-device dtype handling. ChangesComponent device placement and LM initialization
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant ServiceInit
participant DeviceMapping
participant InitServiceOrchestrator
participant LLMHandler
ServiceInit->>DeviceMapping: resolve LM device and offload setting
InitServiceOrchestrator->>DeviceMapping: resolve and validate component devices
InitServiceOrchestrator->>InitServiceOrchestrator: load DiT, VAE, and text encoder on mapped devices
ServiceInit->>LLMHandler: initialize LM with mapped device
LLMHandler-->>ServiceInit: complete LM initialization
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 |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR extends ACE-Step’s multi-GPU “component placement” support into the Gradio and CLI initialization paths so the DiT and 5Hz LM can be pinned to different GPUs (via env overrides and/or auto-ranking), reducing VRAM contention and enabling larger combined configurations on dual-GPU rigs.
Changes:
- Adds a new component→device mapping utility with env overrides for deterministic placement.
- Routes DiT/VAE/text-encoder and the 5Hz LM to potentially different devices in handler + Gradio/CLI init paths.
- Adjusts LM initialization to better handle
cuda:Ndevice strings and (intended) direct bf16 loading to avoid transient VRAM spikes.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| acestep/ui/gradio/events/generation/service_init.py | Uses component device map for LM placement during Gradio init and adjusts LM CPU offload behavior. |
| acestep/llm_inference.py | Updates PyTorch LM load path and dtype selection for device strings like cuda:N. |
| acestep/core/generation/handler/init_service_orchestrator.py | Adds per-component device placement (DiT/VAE) during service initialization. |
| acestep/core/generation/device_mapping.py | New utility module: ranks GPUs by free VRAM, supports env overrides, and validates cuda:N indices. |
| acestep/acestep_v15_pipeline.py | CLI path: disables single-GPU LM offload/downgrade heuristics in multi-GPU LM setups and initializes LM on mapped device. |
…rides, add device_mapping tests
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/ui/gradio/events/generation/service_init.py`:
- Around line 144-150: Cache the initial device map before the init_llm block
and reuse it after dit_handler.initialize_service() instead of calling
resolve_component_device_map() again. In
acestep/ui/gradio/events/generation/service_init.py lines 92-106, initialize
_dev_map = None before the if init_llm block; in lines 144-150, compute
_multi_gpu_lm from _dev_map with a None guard, preserving the existing
_lm_offload behavior.
🪄 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: 6fa7b41d-bb9b-4d3a-b17a-dc728ab0b432
📒 Files selected for processing (4)
acestep/acestep_v15_pipeline.pyacestep/core/generation/device_mapping_test.pyacestep/core/generation/handler/init_service_orchestrator.pyacestep/ui/gradio/events/generation/service_init.py
🚧 Files skipped from review as they are similar to previous changes (2)
- acestep/acestep_v15_pipeline.py
- acestep/core/generation/handler/init_service_orchestrator.py
…oid post-DiT-load re-ranking)
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 4 comments.
Comments suppressed due to low confidence (2)
acestep/core/generation/handler/init_service_orchestrator.py:84
resolve_component_device_map()is applied even when the caller explicitly requested a non-CUDA device (e.g.device="cpu"). On CUDA hosts this will still produce acuda:Nmapping and override the user’s choice, forcing GPU initialization unexpectedly.
resolved_device = self._resolve_initialize_device(device)
if component_device_map is None:
component_device_map = resolve_component_device_map()
validate_component_device_map(component_device_map)
dit_device = component_device_map.dit or resolved_device
acestep/acestep_v15_pipeline.py:435
- The CLI computes
_multi_gpu_lmfrom the CUDA component map unconditionally. On CUDA hosts where the user selected a non-CUDA--device(e.g.cpu/mps/xpu),resolve_component_device_map()will still returncuda:Nentries and can incorrectly disable the CPU offload/downgrade heuristics.
from acestep.core.generation.device_mapping import resolve_component_device_map as _rcdm
_cmap = _rcdm()
_multi_gpu_lm = bool(_cmap.dit and _cmap.lm and _cmap.lm != _cmap.dit)
…es, normalize cuda:N device checks
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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_loader.py`:
- Line 146: Update the CUDA branch in the device-selection logic to parse the
index from the mapped device string when it uses the cuda:N form, then pass that
index to gpu_config.cuda_supports_bfloat16(). Preserve the existing behavior for
the unindexed cuda device.
🪄 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: a528d601-1d7e-4273-9d76-1cf3c2734f52
📒 Files selected for processing (4)
acestep/acestep_v15_pipeline.pyacestep/core/generation/handler/init_service_loader.pyacestep/core/generation/handler/init_service_orchestrator.pyacestep/ui/gradio/events/generation/service_init.py
🚧 Files skipped from review as they are similar to previous changes (3)
- acestep/ui/gradio/events/generation/service_init.py
- acestep/core/generation/handler/init_service_orchestrator.py
- acestep/acestep_v15_pipeline.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Enables multi-GPU component placement for the Gradio UI and command-line generation path, so the DiT and the 5Hz LM can run on separate GPUs.
On dual mid-range cards (e.g. 2×16 GB) this makes it possible to run the XL DiT + 4B LM together at full diffusion speed — a setup that otherwise needs a single ≥24 GB card, or heavy CPU offload of the DiT (which cripples diffusion throughput and times out on long tracks).
Builds on the great groundwork in #1149 by @imsarang (the
ComponentDeviceMapconcept andinitialize_servicewiring, which covered the API-server path). This PR extends it to the Gradio/CLI generation path — where the LM was still initialized on the DiT's device — and adds the pieces needed for real asymmetric multi-GPU rigs.What's included
device_mapping.py— component→device mapping with explicit env overrides (ACESTEP_DIT_DEVICE,ACESTEP_VAE_DEVICE,ACESTEP_LM_DEVICE) for deterministic pinning, on top of the free-VRAM auto-ranking.init_service_orchestrator.py— place DiT / VAE / text_encoder on the mapped devices.acestep_v15_pipeline.py&service_init.py— route the 5Hz LM to its mapped device in the CLI and Gradio init paths, and keep it resident on its dedicated card (no CPU offload) even when the DiT handler offloads VAE/text_encoder to free the DiT's VRAM.llm_inference.py— recognizecuda:N(not just"cuda") when picking the LM dtype, and load the PyTorch LM directly in bf16 instead of loading fp32 then converting — this removes a transient 2× VRAM peak that OOMs a dedicated 16 GB card for the 4B LM.Example config (2×16 GB: RTX 5080 + RTX 4060 Ti)
ACESTEP_DIT_DEVICE=cuda:0 # DiT + VAE + text_encoder on the fast card
ACESTEP_VAE_DEVICE=cuda:0
ACESTEP_LM_DEVICE=cuda:1 # 4B LM resident on the second card
MAX_CUDA_VRAM=24 # unlock the 4B LM in the tier logic
launch flags: --offload_to_cpu true --offload_dit_to_cpu false --backend pt
Tested
RTX 5080 (16 GB) + RTX 4060 Ti (16 GB), Windows, torch 2.7.1+cu128:
initialize_serviceplaces DiT on cuda:0, 4B LM (bf16) on cuda:1 — no OOM.offload_dit_to_cpu, which slowed diffusion to ~35–52 s/step and timed out on long tracks.Notes
ptbackend on its dedicated card (reliable device placement). Making nano-vLLM honor a non-default CUDA device would further speed the LM phase — left for a follow-up.device_mapping.py/ orchestrator bits and keep only the Gradio-path + dtype fixes.Summary by CodeRabbit
cuda:1).