feat: multi-GPU component placement (DiT / 5Hz LM on separate GPUs) for Gradio + CLI

acestep/acestep_v15_pipeline.py

@@ -426,9 +426,28 @@ def main():
     if args.service_mode:
         print(f"  Backend: {args.backend}")
 
+    # Multi-GPU component placement: when the LM is mapped to a different GPU than
+    # the DiT (via ACESTEP_*_DEVICE overrides OR free-VRAM auto-ranking), the LM no
+    # longer competes for the DiT's VRAM, so the single-GPU offload/downgrade
+    # heuristics below do not apply.
+    # Resolve the component->device map ONCE here (before the DiT is loaded) and
+    # reuse it for the LM placement below so both stay consistent. Only consult the
+    # CUDA map when a CUDA/auto device is requested: an explicit cpu/mps/xpu device
+    # must not be silently overridden onto a cuda:N card. Validate env overrides
+    # (e.g. ACESTEP_LM_DEVICE=cuda:9) up front so they fail with a clear message.
+    from acestep.core.generation.device_mapping import (
+        resolve_component_device_map as _rcdm,
+        validate_component_device_map as _vcdm,
+    )
+    _device_kind = str(args.device).split(":", 1)[0]
+    _cmap = _rcdm() if _device_kind in {"auto", "cuda"} else None
+    if _cmap is not None:
+        _vcdm(_cmap)
+    _multi_gpu_lm = bool(_cmap and _cmap.dit and _cmap.lm and _cmap.lm != _cmap.dit)
+
     # Auto-enable CPU offload for tier6 GPUs (16-24GB) when using the 4B LM model
     # The 4B LM (~8GB) + DiT (~4.7GB) + VAE + text encoder exceeds 16-20GB with activations
-    if not args.offload_to_cpu and args.lm_model_path and "4B" in args.lm_model_path:
+    if not args.offload_to_cpu and args.lm_model_path and "4B" in args.lm_model_path and not _multi_gpu_lm:
         if 0 < gpu_memory_gb <= 24:
             args.offload_to_cpu = True
             print(
@@ -438,7 +457,7 @@ def main():
     # Safety: on 16GB GPUs, prevent selecting LM models that are too large.
     # Even with offloading, a 4B LM (8 GB weights + KV cache) leaves almost no
     # headroom for DiT activations on a 16 GB card.
-    if args.lm_model_path and 0 < gpu_memory_gb < VRAM_AUTO_OFFLOAD_THRESHOLD_GB:
+    if args.lm_model_path and 0 < gpu_memory_gb < VRAM_AUTO_OFFLOAD_THRESHOLD_GB and not _multi_gpu_lm:
         if "4B" in args.lm_model_path:
             # Downgrade to 1.7B if available
             fallback = args.lm_model_path.replace("4B", "1.7B")
@@ -566,15 +585,23 @@ def main():
                             file=sys.stderr,
                         )
 
+                    # Reuse the map resolved above (None when a non-CUDA device was
+                    # requested); re-resolving here would re-rank GPUs by free VRAM
+                    # after the DiT load and could disagree with the placement above.
+                    _lm_device = (_cmap.lm if _cmap is not None else None) or args.device
+                    # In multi-GPU mode the LM has its own dedicated card, so it must
+                    # stay resident there (do NOT offload it to CPU even if the DiT
+                    # handler offloads VAE/text-encoder to free the DiT's VRAM).
+                    _lm_offload = False if _multi_gpu_lm else args.offload_to_cpu
                     print(
-                        f"Initializing 5Hz LM: {args.lm_model_path} on {args.device}..."
+                        f"Initializing 5Hz LM: {args.lm_model_path} on {_lm_device} (offload_to_cpu={_lm_offload})..."
                     )
                     lm_status, lm_success = llm_handler.initialize(
                         checkpoint_dir=checkpoint_dir,
                         lm_model_path=args.lm_model_path,
                         backend=args.backend,
-                        device=args.device,
-                        offload_to_cpu=args.offload_to_cpu,
+                        device=_lm_device,
+                        offload_to_cpu=_lm_offload,
                         dtype=None,
                     )
 

acestep/core/generation/device_mapping.py

@@ -0,0 +1,165 @@
+"""Utilities for dynamic multi-device component mapping.
+
+Based on PR #1149 (imsarang) — extended with explicit environment-variable
+overrides (ACESTEP_DIT_DEVICE / ACESTEP_VAE_DEVICE / ACESTEP_LM_DEVICE) so that
+users of asymmetric multi-GPU rigs can pin each component deterministically
+(e.g. DiT on the fast card, LM on the second card).
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import Optional
+
+from loguru import logger  # type: ignore[reportMissingImports]
+
+try:
+    import torch  # type: ignore[reportMissingImports]
+except ImportError:  # pragma: no cover - fallback for minimal environments
+    torch = None  # type: ignore[assignment]
+
+
+@dataclass(frozen=True)
+class ComponentDeviceMap:
+    """Normalized component-to-device mapping."""
+
+    dit: Optional[str] = None
+    vae: Optional[str] = None
+    lm: Optional[str] = None
+
+
+def _rank_cuda_devices_by_free_vram() -> list[int]:
+    """Return CUDA device indices sorted by descending free VRAM."""
+    if torch is None or not torch.cuda.is_available():
+        return []
+
+    device_count = torch.cuda.device_count()
+    free_by_device: list[tuple[int, int]] = []
+    for idx in range(device_count):
+        free_bytes = 0
+        try:
+            free_bytes = int(torch.cuda.mem_get_info(idx)[0])
+        except RuntimeError as exc:
+            logger.debug(
+                "[device_mapping] mem_get_info({}) failed ({}); retrying with device context.",
+                idx,
+                exc,
+            )
+            try:
+                with torch.cuda.device(idx):
+                    free_bytes = int(torch.cuda.mem_get_info()[0])
+            except RuntimeError as exc2:
+                logger.warning(
+                    "[device_mapping] Unable to query free VRAM for cuda:{} ({}); ranking it last.",
+                    idx,
+                    exc2,
+                )
+                free_bytes = 0
+        free_by_device.append((idx, free_bytes))
+    free_by_device.sort(key=lambda pair: pair[1], reverse=True)
+    return [idx for idx, _free in free_by_device]
+
+
+def _env_device_overrides() -> tuple[Optional[str], Optional[str], Optional[str]]:
+    """Read explicit per-component device overrides from the environment."""
+    def _norm(val: Optional[str]) -> Optional[str]:
+        if val is None:
+            return None
+        val = val.strip()
+        return val or None
+
+    return (
+        _norm(os.getenv("ACESTEP_DIT_DEVICE")),
+        _norm(os.getenv("ACESTEP_VAE_DEVICE")),
+        _norm(os.getenv("ACESTEP_LM_DEVICE")),
+    )
+
+
+def resolve_component_device_map() -> ComponentDeviceMap:
+    """Auto-populate component mapping from available CUDA devices.
+
+    Explicit env overrides (ACESTEP_DIT_DEVICE / ACESTEP_VAE_DEVICE /
+    ACESTEP_LM_DEVICE) take precedence over the free-VRAM auto-ranking, allowing
+    deterministic pinning on asymmetric multi-GPU setups.
+    """
+    env_dit, env_vae, env_lm = _env_device_overrides()
+
+    if torch is None or not torch.cuda.is_available():
+        if env_dit or env_vae or env_lm:
+            return ComponentDeviceMap(dit=env_dit, vae=env_vae, lm=env_lm)
+        return ComponentDeviceMap()
+
+    ranked = _rank_cuda_devices_by_free_vram()
+    if not ranked:
+        if env_dit or env_vae or env_lm:
+            return ComponentDeviceMap(dit=env_dit, vae=env_vae, lm=env_lm)
+        return ComponentDeviceMap()
+
+    dit_idx = ranked[0]
+    vae_idx = ranked[1] if len(ranked) > 1 else ranked[0]
+    lm_idx = ranked[2] if len(ranked) > 2 else ranked[-1]
+
+    return ComponentDeviceMap(
+        dit=env_dit or f"cuda:{dit_idx}",
+        vae=env_vae or f"cuda:{vae_idx}",
+        lm=env_lm or f"cuda:{lm_idx}",
+    )
+
+
+def format_component_gpu_hint_text(
+    *,
+    default_device: str = "auto",
+    label: str = "Component GPU hint",
+) -> str:
+    """Format component placement hint for UI display.
+
+    Returns an empty string when all components resolve to the same final device,
+    which avoids noisy UI on non-CUDA and single-device hosts.
+    """
+    mapping = resolve_component_device_map()
+    resolved_devices = [
+        mapping.dit or default_device,
+        mapping.vae or default_device,
+        mapping.lm or default_device,
+    ]
+    if len(set(resolved_devices)) == 1:
+        return ""
+
+    return (
+        f"{label}: "
+        f"DiT={resolved_devices[0]}, "
+        f"VAE={resolved_devices[1]}, "
+        f"LM={resolved_devices[2]}"
+    )
+
+
+def validate_component_device_map(mapping: ComponentDeviceMap) -> None:
+    """Validate that mapped CUDA indices exist on the current host.
+
+    Args:
+        mapping: Component-to-device mapping to validate.
+
+    Raises:
+        ValueError: If any ``cuda:<idx>`` entry references an index out of range
+            for visible CUDA devices.
+    """
+    cuda_count = torch.cuda.device_count() if torch is not None and torch.cuda.is_available() else 0
+    for component, device in (
+        ("dit", mapping.dit),
+        ("vae", mapping.vae),
+        ("lm", mapping.lm),
+    ):
+        if not device or not device.startswith("cuda:"):
+            continue
+        try:
+            idx = int(device.split(":", 1)[1])
+        except (TypeError, ValueError) as exc:
+            raise ValueError(
+                f"Invalid {component} device '{device}': expected 'cuda:<int>' with "
+                f"0 <= idx < {cuda_count}."
+            ) from exc
+        if idx < 0 or idx >= cuda_count:
+            raise ValueError(
+                f"Invalid {component} device '{device}': only {cuda_count} CUDA device(s) available."
+            )

acestep/core/generation/device_mapping_test.py

@@ -0,0 +1,105 @@
+"""Unit tests for multi-GPU component device mapping."""
+
+import os
+import unittest
+from unittest import mock
+
+from acestep.core.generation import device_mapping as dm
+from acestep.core.generation.device_mapping import ComponentDeviceMap
+
+_ENV_KEYS = ("ACESTEP_DIT_DEVICE", "ACESTEP_VAE_DEVICE", "ACESTEP_LM_DEVICE")
+
+
+class _FakeCuda:
+    def __init__(self, count, free_by_idx):
+        self._count = count
+        self._free = free_by_idx
+
+    def is_available(self):
+        return True
+
+    def device_count(self):
+        return self._count
+
+    def mem_get_info(self, idx=0):
+        return (self._free[idx], 16 * 1024 ** 3)
+
+
+class _FakeTorch:
+    def __init__(self, count, free_by_idx):
+        self.cuda = _FakeCuda(count, free_by_idx)
+
+
+def _clear_env():
+    for k in _ENV_KEYS:
+        os.environ.pop(k, None)
+
+
+class DeviceMappingTest(unittest.TestCase):
+    def setUp(self):
+        _clear_env()
+
+    def tearDown(self):
+        _clear_env()
+
+    def test_env_overrides_take_precedence_over_ranking(self):
+        fake = _FakeTorch(2, {0: 5 * 1024 ** 3, 1: 10 * 1024 ** 3})
+        with mock.patch.object(dm, "torch", fake):
+            os.environ["ACESTEP_DIT_DEVICE"] = "cuda:0"
+            os.environ["ACESTEP_VAE_DEVICE"] = "cuda:0"
+            os.environ["ACESTEP_LM_DEVICE"] = "cuda:1"
+            m = dm.resolve_component_device_map()
+        self.assertEqual(m.dit, "cuda:0")
+        self.assertEqual(m.vae, "cuda:0")
+        self.assertEqual(m.lm, "cuda:1")
+
+    def test_auto_ranking_by_free_vram(self):
+        # idx 1 has more free VRAM -> DiT; idx 0 -> VAE and LM (2 GPUs).
+        fake = _FakeTorch(2, {0: 5 * 1024 ** 3, 1: 10 * 1024 ** 3})
+        with mock.patch.object(dm, "torch", fake):
+            m = dm.resolve_component_device_map()
+        self.assertEqual(m.dit, "cuda:1")
+        self.assertEqual(m.vae, "cuda:0")
+        self.assertEqual(m.lm, "cuda:0")
+
+    def test_single_gpu_maps_all_to_same_device(self):
+        fake = _FakeTorch(1, {0: 8 * 1024 ** 3})
+        with mock.patch.object(dm, "torch", fake):
+            m = dm.resolve_component_device_map()
+        self.assertEqual(m.dit, "cuda:0")
+        self.assertEqual(m.vae, "cuda:0")
+        self.assertEqual(m.lm, "cuda:0")
+
+    def test_no_cuda_returns_empty_map_without_env(self):
+        with mock.patch.object(dm, "torch", None):
+            m = dm.resolve_component_device_map()
+        self.assertEqual(m, ComponentDeviceMap())
+
+    def test_validate_raises_on_out_of_range_index(self):
+        fake = _FakeTorch(2, {0: 1, 1: 1})
+        with mock.patch.object(dm, "torch", fake):
+            with self.assertRaises(ValueError):
+                dm.validate_component_device_map(ComponentDeviceMap(dit="cuda:5"))
+
+    def test_validate_passes_on_valid_indices(self):
+        fake = _FakeTorch(2, {0: 1, 1: 1})
+        with mock.patch.object(dm, "torch", fake):
+            dm.validate_component_device_map(
+                ComponentDeviceMap(dit="cuda:0", vae="cuda:0", lm="cuda:1")
+            )
+
+    def test_hint_is_empty_when_single_device(self):
+        fake = _FakeTorch(1, {0: 8 * 1024 ** 3})
+        with mock.patch.object(dm, "torch", fake):
+            self.assertEqual(dm.format_component_gpu_hint_text(), "")
+
+    def test_hint_non_empty_when_components_differ(self):
+        fake = _FakeTorch(2, {0: 5 * 1024 ** 3, 1: 10 * 1024 ** 3})
+        with mock.patch.object(dm, "torch", fake):
+            hint = dm.format_component_gpu_hint_text(label="GPU map")
+        self.assertIn("DiT=cuda:1", hint)
+        self.assertIn("LM=cuda:0", hint)
+
+
+if __name__ == "__main__":
+    unittest.main()

acestep/core/generation/handler/init_service_loader.py

@@ -141,9 +141,14 @@ def _load_main_model_from_checkpoint(
                     exc,
                 )
 
+        # Query bfloat16 support on the *mapped* CUDA device (e.g. cuda:1), not the
+        # current/default device, so mixed-generation multi-GPU rigs select the right
+        # attention backend. A bare "cuda" parses to index None (= current device).
+        _device_kind, _, _device_index = str(device).partition(":")
+        _cuda_index = int(_device_index) if _device_index.isdigit() else None
         if use_flash_attention and self.is_flash_attention_available(device):
             attn_implementation = "flash_attention_2"
-        elif device == "cuda" and not gpu_config.cuda_supports_bfloat16():
+        elif _device_kind == "cuda" and not gpu_config.cuda_supports_bfloat16(_cuda_index):
             # Pre-Ampere GPUs (compute capability < 8.0) run in float16 which
             # can overflow in SDPA's fused softmax with longer sequences,
             # producing NaN/Inf latents (see issues #924, #927).  Eager