diff --git a/acestep/api/startup_llm_init.py b/acestep/api/startup_llm_init.py index dd5199d1c..d0f0c1015 100644 --- a/acestep/api/startup_llm_init.py +++ b/acestep/api/startup_llm_init.py @@ -23,6 +23,7 @@ def initialize_llm_at_startup( get_model_name: Callable[[str], str], ensure_model_downloaded: Callable[[str, str], str], env_bool: Callable[[str, bool], bool], + dit_handler: Any = None, ) -> None: """Initialize LLM model according to GPU config and environment overrides.""" @@ -73,6 +74,9 @@ def initialize_llm_at_startup( lm_backend = resolve_lm_backend(os.getenv("ACESTEP_LM_BACKEND"), gpu_config) lm_device = os.getenv("ACESTEP_LM_DEVICE", device) + device_map = getattr(dit_handler, "device_map", None) if dit_handler is not None else None + if device_map is not None and device_map.lm is not None: + lm_device = device_map.lm lm_offload_env = os.getenv("ACESTEP_LM_OFFLOAD_TO_CPU") lm_offload = env_bool("ACESTEP_LM_OFFLOAD_TO_CPU", False) if lm_offload_env is not None else offload_to_cpu diff --git a/acestep/api/startup_model_init.py b/acestep/api/startup_model_init.py index c8bca850f..563068625 100644 --- a/acestep/api/startup_model_init.py +++ b/acestep/api/startup_model_init.py @@ -85,6 +85,7 @@ def do_model_initialization( compile_model=compile_model, offload_to_cpu=offload_to_cpu, offload_dit_to_cpu=offload_dit_to_cpu, + gpu_mapping=os.getenv("ACESTEP_GPU_MAPPING"), ) if not ok: app.state._init_error = status_msg @@ -157,6 +158,7 @@ def do_model_initialization( get_model_name=get_model_name, ensure_model_downloaded=ensure_model_downloaded, env_bool=env_bool, + dit_handler=handler, ) print("[API Server] All models initialized successfully!") diff --git a/acestep/core/generation/handler/audio_codes.py b/acestep/core/generation/handler/audio_codes.py index 31d6835c9..e0e8664ad 100644 --- a/acestep/core/generation/handler/audio_codes.py +++ b/acestep/core/generation/handler/audio_codes.py @@ -56,7 +56,8 @@ def _decode_audio_codes_to_latents(self, code_str: str) -> Optional[torch.Tensor with self._load_model_context("model"): quantizer = self.model.tokenizer.quantizer detokenizer = self.model.detokenizer - indices = torch.tensor(code_ids, device=self.device, dtype=torch.long) + dit_device = self._get_component_device("model") + indices = torch.tensor(code_ids, device=dit_device, dtype=torch.long) indices = indices.unsqueeze(0).unsqueeze(-1) quantized = quantizer.get_output_from_indices(indices) @@ -83,7 +84,11 @@ def convert_src_audio_to_codes(self, audio_file) -> str: return "❌ Audio file appears to be silent" latents = self._encode_audio_to_latents(processed_audio) - attention_mask = torch.ones(latents.shape[0], dtype=torch.bool, device=self.device) + attention_mask = torch.ones( + latents.shape[0], + dtype=torch.bool, + device=self._get_component_device("model"), + ) with self._load_model_context("model"): hidden_states = latents.unsqueeze(0) _, indices, _ = self.model.tokenize( diff --git a/acestep/core/generation/handler/conditioning_embed.py b/acestep/core/generation/handler/conditioning_embed.py index ef9e4a572..6a38d5db9 100644 --- a/acestep/core/generation/handler/conditioning_embed.py +++ b/acestep/core/generation/handler/conditioning_embed.py @@ -56,7 +56,9 @@ def _ensure_latent_3d(z: torch.Tensor) -> torch.Tensor: refer_audio = _normalize_audio_2d(refer_audio) with torch.inference_mode(): refer_audio_latent = self.tiled_encode(refer_audio, offload_latent_to_cpu=True) - refer_audio_latent = refer_audio_latent.to(self.device).to(self.dtype) + refer_audio_latent = refer_audio_latent.to( + self._get_component_device("dit") + ).to(self.dtype) if refer_audio_latent.dim() == 2: refer_audio_latent = refer_audio_latent.unsqueeze(0) refer_audio_latent = _ensure_latent_3d(refer_audio_latent.transpose(1, 2)) @@ -65,16 +67,28 @@ def _ensure_latent_3d(z: torch.Tensor) -> torch.Tensor: refer_audio_order_mask.append(batch_idx) refer_audio_latents = torch.cat(refer_audio_latents, dim=0) - refer_audio_order_mask = torch.tensor(refer_audio_order_mask, device=self.device, dtype=torch.long) + dit_device = self._get_component_device("dit") + refer_audio_latents = refer_audio_latents.to(dit_device).to(self.dtype) + refer_audio_order_mask = torch.tensor( + refer_audio_order_mask, device=dit_device, dtype=torch.long + ) return refer_audio_latents, refer_audio_order_mask def infer_text_embeddings(self, text_token_idss): - """Infer text-token embeddings via text encoder.""" + """Infer text-token embeddings via text encoder. + + Token ids are often built on the DiT device; move them onto the + text-encoder device so split maps like ``text_encoder:3`` work. + """ + encoder_device = self._get_component_device("text_encoder") + text_token_idss = text_token_idss.to(encoder_device) with torch.inference_mode(): return self.text_encoder(input_ids=text_token_idss, lyric_attention_mask=None).last_hidden_state def infer_lyric_embeddings(self, lyric_token_ids): """Infer lyric-token embeddings via text encoder embedding table.""" + encoder_device = self._get_component_device("text_encoder") + lyric_token_ids = lyric_token_ids.to(encoder_device) with torch.inference_mode(): return self.text_encoder.embed_tokens(lyric_token_ids) @@ -124,6 +138,22 @@ def preprocess_batch(self, batch) -> Tuple: repaint_mask = batch.get("repaint_mask", None) + dit_device = self._get_component_device("dit") + text_hidden_states = text_hidden_states.to(dit_device) + lyric_hidden_states = lyric_hidden_states.to(dit_device) + text_attention_mask = text_attention_mask.to(dit_device) + lyric_attention_mask = lyric_attention_mask.to(dit_device) + src_latents = src_latents.to(dit_device) + target_latents = target_latents.to(dit_device) + chunk_mask = chunk_mask.to(dit_device) + is_covers = is_covers.to(dit_device) if isinstance(is_covers, torch.Tensor) else is_covers + if precomputed_lm_hints_25hz is not None: + precomputed_lm_hints_25hz = precomputed_lm_hints_25hz.to(dit_device) + if non_cover_text_hidden_states is not None: + non_cover_text_hidden_states = non_cover_text_hidden_states.to(dit_device) + if non_cover_text_attention_masks is not None: + non_cover_text_attention_masks = non_cover_text_attention_masks.to(dit_device) + return ( keys, text_inputs, diff --git a/acestep/core/generation/handler/conditioning_embed_test.py b/acestep/core/generation/handler/conditioning_embed_test.py index 8cb85413a..3d7507aec 100644 --- a/acestep/core/generation/handler/conditioning_embed_test.py +++ b/acestep/core/generation/handler/conditioning_embed_test.py @@ -24,12 +24,21 @@ def embed_tokens(self, token_ids): class _Host(ConditioningEmbedMixin): """Minimal host implementing ConditioningEmbedMixin dependencies.""" - def __init__(self): - self.device = "cpu" + def __init__(self, text_encoder_device="cpu", dit_device="cpu"): + self.device = dit_device self.dtype = torch.float32 self.silence_latent = torch.zeros(1, 128, 6, dtype=torch.float32) self.text_encoder = _FakeTextEncoder() self.tiled_encode_calls = 0 + self._devices = { + "dit": dit_device, + "vae": dit_device, + "text_encoder": text_encoder_device, + "lm": dit_device, + } + + def _get_component_device(self, component): + return self._devices[component] def _ensure_silence_latent_on_device(self): return None @@ -111,6 +120,25 @@ def test_preprocess_batch_returns_expected_tuple_shape(self): self.assertEqual(result[0], ["k1", "k2"]) self.assertEqual(result[3].shape, (2, 128, 6)) + def test_infer_text_embeddings_moves_ids_to_text_encoder_device(self): + """Token ids on DiT device must be moved onto a remote text encoder.""" + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + self.skipTest("needs 2+ CUDA devices") + + seen = {} + + class _DeviceCheckingEncoder(_FakeTextEncoder): + def __call__(self, input_ids, lyric_attention_mask=None): + seen["device"] = str(input_ids.device) + return super().__call__(input_ids, lyric_attention_mask) + + host = _Host(text_encoder_device="cuda:1", dit_device="cuda:0") + host.text_encoder = _DeviceCheckingEncoder() + ids = torch.ones(1, 4, dtype=torch.long, device="cuda:0") + out = host.infer_text_embeddings(ids) + self.assertEqual(seen["device"], "cuda:1") + self.assertEqual(out.shape, (1, 4, 6)) + if __name__ == "__main__": unittest.main() diff --git a/acestep/core/generation/handler/generate_music_decode.py b/acestep/core/generation/handler/generate_music_decode.py index b7330503c..93acf5c7b 100644 --- a/acestep/core/generation/handler/generate_music_decode.py +++ b/acestep/core/generation/handler/generate_music_decode.py @@ -9,6 +9,7 @@ from loguru import logger from acestep.gpu_config import get_effective_free_vram_gb +from acestep.device_map import cuda_device_index, is_cuda_device class GenerateMusicDecodeMixin: @@ -129,7 +130,12 @@ def _decode_generate_music_pred_latents( with torch.inference_mode(): with self._load_model_context("vae"): pred_latents_cpu = pred_latents.detach().cpu() - pred_latents_for_decode = pred_latents.transpose(1, 2).contiguous().to(self.vae.dtype) + vae_component_device = self._get_component_device("vae") + pred_latents_for_decode = ( + pred_latents.transpose(1, 2) + .contiguous() + .to(device=vae_component_device, dtype=self.vae.dtype) + ) del pred_latents self._empty_cache() @@ -140,7 +146,7 @@ def _decode_generate_music_pred_latents( ) using_mlx_vae = self.use_mlx_vae and self.mlx_vae is not None vae_cpu = False - vae_device = None + vae_restore_device = None if not using_mlx_vae: vae_cpu = os.environ.get("ACESTEP_VAE_ON_CPU", "0").lower() in ("1", "true", "yes") if not vae_cpu: @@ -150,7 +156,12 @@ def _decode_generate_music_pred_latents( "(unified memory), keeping VAE on MPS" ) else: - effective_free = get_effective_free_vram_gb() + vae_cuda_index = ( + cuda_device_index(vae_component_device) + if is_cuda_device(vae_component_device) + else 0 + ) + effective_free = get_effective_free_vram_gb(vae_cuda_index) logger.info( "[generate_music] Effective free VRAM before VAE decode: " f"{effective_free:.2f} GB" @@ -163,7 +174,7 @@ def _decode_generate_music_pred_latents( vae_cpu = True if vae_cpu: logger.info("[generate_music] Moving VAE to CPU for decode (ACESTEP_VAE_ON_CPU=1)...") - vae_device = next(self.vae.parameters()).device + vae_restore_device = next(self.vae.parameters()).device self.vae = self.vae.cpu() pred_latents_for_decode = pred_latents_for_decode.cpu() self._empty_cache() @@ -186,9 +197,9 @@ def _decode_generate_music_pred_latents( pred_wavs = decoder_output.sample del decoder_output finally: - if vae_cpu and vae_device is not None: + if vae_cpu and vae_restore_device is not None: logger.info("[generate_music] Restoring VAE to original device after CPU decode path...") - self.vae = self.vae.to(vae_device) + self.vae = self.vae.to(vae_restore_device) self._empty_cache() logger.debug( "[generate_music] After VAE decode: " diff --git a/acestep/core/generation/handler/generate_music_decode_basic_test.py b/acestep/core/generation/handler/generate_music_decode_basic_test.py new file mode 100644 index 000000000..a3afb5f3d --- /dev/null +++ b/acestep/core/generation/handler/generate_music_decode_basic_test.py @@ -0,0 +1,42 @@ +"""Tests for the default ``generate_music`` latent decode path.""" + +import unittest +from unittest.mock import patch + +import torch + +from acestep.core.generation.handler.generate_music_decode_test_support import ( + GENERATE_MUSIC_DECODE_MODULE, + DecodeTestHost, +) + + +class GenerateMusicDecodeBasicTests(unittest.TestCase): + """Verify successful decode timing and output behavior.""" + + def test_decode_pred_latents_updates_decode_time_and_returns_cpu_latents(self): + host = DecodeTestHost() + pred_latents = torch.ones(1, 4, 3) + time_costs = {"total_time_cost": 1.0} + + def _progress(value, desc=None): + host.progress_calls.append((value, desc)) + + with patch.object(GENERATE_MUSIC_DECODE_MODULE.time, "time", side_effect=[10.0, 11.5]): + pred_wavs, pred_latents_cpu, updated_costs = host._decode_generate_music_pred_latents( + pred_latents=pred_latents, + progress=_progress, + use_tiled_decode=False, + time_costs=time_costs, + ) + + self.assertEqual(tuple(pred_wavs.shape), (1, 2, 8)) + self.assertEqual(pred_latents_cpu.device.type, "cpu") + self.assertAlmostEqual(updated_costs["vae_decode_time_cost"], 1.5, places=6) + self.assertAlmostEqual(updated_costs["total_time_cost"], 2.5, places=6) + self.assertAlmostEqual(updated_costs["offload_time_cost"], 0.25, places=6) + self.assertEqual(host.progress_calls[0][0], 0.8) + + +if __name__ == "__main__": + unittest.main() diff --git a/acestep/core/generation/handler/generate_music_decode_cpu_offload_test.py b/acestep/core/generation/handler/generate_music_decode_cpu_offload_test.py new file mode 100644 index 000000000..8df421f43 --- /dev/null +++ b/acestep/core/generation/handler/generate_music_decode_cpu_offload_test.py @@ -0,0 +1,100 @@ +"""Tests for CPU-offload and VAE restoration in ``generate_music`` decode.""" + +import unittest +from unittest.mock import patch + +import torch + +from acestep.core.generation.handler.generate_music_decode_test_support import ( + GENERATE_MUSIC_DECODE_MODULE, + DecodeTestHost, + FakeDecodeOutput, + FakeVae, +) + + +class GenerateMusicDecodeCpuOffloadTests(unittest.TestCase): + """Verify CPU-offload decode paths and device restoration.""" + + def test_decode_pred_latents_restores_vae_device_on_decode_error(self): + class FailingVae(FakeVae): + def __init__(self): + super().__init__() + self.cpu_calls = 0 + self.to_calls = [] + + def decode(self, latents: torch.Tensor): + _ = latents + raise RuntimeError("decode failed") + + def cpu(self): + self.cpu_calls += 1 + return self + + def to(self, *args, **kwargs): + self.to_calls.append((args, kwargs)) + return self + + class FailingHost(DecodeTestHost): + def __init__(self): + super().__init__() + self.use_mlx_vae = False + self.mlx_vae = None + self.vae = FailingVae() + self.empty_cache_calls = 0 + + def _empty_cache(self): + self.empty_cache_calls += 1 + + host = FailingHost() + with patch.dict(GENERATE_MUSIC_DECODE_MODULE.os.environ, {"ACESTEP_VAE_ON_CPU": "1"}, clear=False): + with self.assertRaisesRegex(RuntimeError, "decode failed"): + host._decode_generate_music_pred_latents( + pred_latents=torch.ones(1, 4, 3), + progress=None, + use_tiled_decode=False, + time_costs={"total_time_cost": 1.0}, + ) + + self.assertEqual(host.vae.cpu_calls, 1) + self.assertEqual(len(host.vae.to_calls), 1) + self.assertGreaterEqual(host.empty_cache_calls, 2) + + def test_decode_pred_latents_does_not_restore_latents_to_gpu_after_successful_cpu_decode(self): + class SuccessVae(FakeVae): + def __init__(self): + super().__init__() + self.vae_to_calls = [] + + def decode(self, latents: torch.Tensor): + return FakeDecodeOutput(torch.ones(latents.shape[0], 2, 8)) + + def cpu(self): + return self + + def to(self, *args, **kwargs): + self.vae_to_calls.append(args[0] if args else kwargs) + return self + + class SuccessHost(DecodeTestHost): + def __init__(self): + super().__init__() + self.use_mlx_vae = False + self.mlx_vae = None + self.vae = SuccessVae() + + host = SuccessHost() + with patch.dict(GENERATE_MUSIC_DECODE_MODULE.os.environ, {"ACESTEP_VAE_ON_CPU": "1"}, clear=False): + pred_wavs, _cpu_latents, _costs = host._decode_generate_music_pred_latents( + pred_latents=torch.ones(1, 4, 3), + progress=None, + use_tiled_decode=False, + time_costs={"total_time_cost": 1.0}, + ) + + self.assertEqual(len(host.vae.vae_to_calls), 1) + self.assertEqual(tuple(pred_wavs.shape), (1, 2, 8)) + + +if __name__ == "__main__": + unittest.main() diff --git a/acestep/core/generation/handler/generate_music_decode_prepare_test.py b/acestep/core/generation/handler/generate_music_decode_prepare_test.py new file mode 100644 index 000000000..5752631ba --- /dev/null +++ b/acestep/core/generation/handler/generate_music_decode_prepare_test.py @@ -0,0 +1,50 @@ +"""Tests for ``generate_music`` decode-state preparation.""" + +import unittest + +import torch + +from acestep.core.generation.handler.generate_music_decode_test_support import DecodeTestHost + + +class GenerateMusicDecodePrepareTests(unittest.TestCase): + """Verify decode-state preparation helper behavior.""" + + def test_prepare_decode_state_updates_progress_estimates(self): + host = DecodeTestHost() + outputs = { + "target_latents": torch.ones(1, 4, 3), + "time_costs": {"total_time_cost": 1.0, "diffusion_per_step_time_cost": 0.2}, + } + pred_latents, time_costs = host._prepare_generate_music_decode_state( + outputs=outputs, + infer_steps_for_progress=8, + actual_batch_size=1, + audio_duration=12.0, + latent_shift=0.0, + latent_rescale=1.0, + ) + self.assertEqual(tuple(pred_latents.shape), (1, 4, 3)) + self.assertEqual(time_costs["offload_time_cost"], 0.25) + self.assertEqual(host._last_diffusion_per_step_sec, 0.2) + self.assertEqual(host.estimate_calls[0]["infer_steps"], 8) + + def test_prepare_decode_state_raises_for_nan_latents(self): + host = DecodeTestHost() + outputs = { + "target_latents": torch.tensor([[[float("nan")]]]), + "time_costs": {"total_time_cost": 1.0}, + } + with self.assertRaises(RuntimeError): + host._prepare_generate_music_decode_state( + outputs=outputs, + infer_steps_for_progress=8, + actual_batch_size=1, + audio_duration=None, + latent_shift=0.0, + latent_rescale=1.0, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/acestep/core/generation/handler/generate_music_decode_test.py b/acestep/core/generation/handler/generate_music_decode_test.py deleted file mode 100644 index ece612ff1..000000000 --- a/acestep/core/generation/handler/generate_music_decode_test.py +++ /dev/null @@ -1,316 +0,0 @@ -"""Tests for extracted ``generate_music`` decode helper mixin behavior.""" - -import importlib.util -import types -import sys -import unittest -from contextlib import contextmanager -from pathlib import Path -from unittest.mock import patch - -import torch - - -def _load_generate_music_decode_module(): - """Load ``generate_music_decode.py`` from disk and return its module object. - - Raises ``FileNotFoundError`` or ``ImportError`` when loading fails. - """ - repo_root = Path(__file__).resolve().parents[4] - if str(repo_root) not in sys.path: - sys.path.insert(0, str(repo_root)) - package_paths = { - "acestep": repo_root / "acestep", - "acestep.core": repo_root / "acestep" / "core", - "acestep.core.generation": repo_root / "acestep" / "core" / "generation", - "acestep.core.generation.handler": repo_root / "acestep" / "core" / "generation" / "handler", - } - for package_name, package_path in package_paths.items(): - if package_name in sys.modules: - continue - package_module = types.ModuleType(package_name) - package_module.__path__ = [str(package_path)] - sys.modules[package_name] = package_module - module_path = Path(__file__).with_name("generate_music_decode.py") - spec = importlib.util.spec_from_file_location( - "acestep.core.generation.handler.generate_music_decode", - module_path, - ) - module = importlib.util.module_from_spec(spec) - assert spec.loader is not None - spec.loader.exec_module(module) - return module - - -GENERATE_MUSIC_DECODE_MODULE = _load_generate_music_decode_module() -GenerateMusicDecodeMixin = GENERATE_MUSIC_DECODE_MODULE.GenerateMusicDecodeMixin - - -class _FakeDecodeOutput: - """Minimal VAE decode output container exposing ``sample`` attribute.""" - - def __init__(self, sample: torch.Tensor): - """Store decoded sample tensor for mixin decode flow.""" - self.sample = sample - - -class _FakeVae: - """Minimal VAE stand-in with dtype, decode, and parameter iteration hooks.""" - - def __init__(self): - """Initialize deterministic dtype/device state for decode tests.""" - self.dtype = torch.float32 - self._param = torch.nn.Parameter(torch.zeros(1)) - - def decode(self, latents: torch.Tensor): - """Return deterministic decoded waveform output.""" - return _FakeDecodeOutput(torch.ones(latents.shape[0], 2, 8)) - - def parameters(self): - """Yield one parameter so `.device` lookups remain valid.""" - yield self._param - - def cpu(self): - """Return self for test-only CPU transfer calls.""" - return self - - def to(self, *_args, **_kwargs): - """Return self for test-only device transfer calls.""" - return self - - -class _Host(GenerateMusicDecodeMixin): - """Minimal decode-mixin host exposing deterministic state for assertions.""" - - def __init__(self): - """Initialize deterministic runtime state for decode tests.""" - self.current_offload_cost = 0.25 - self.debug_stats = False - self._last_diffusion_per_step_sec = None - self.estimate_calls = [] - self.progress_calls = [] - self.device = "cpu" - self.use_mlx_vae = True - self.mlx_vae = object() - self.vae = _FakeVae() - - def _update_progress_estimate(self, **kwargs): - """Capture estimate updates for assertions.""" - self.estimate_calls.append(kwargs) - - @contextmanager - def _load_model_context(self, _model_name): - """Provide no-op model context manager for decode tests.""" - yield - - def _empty_cache(self): - """Provide no-op cache clear helper for decode tests.""" - return None - - def _memory_allocated(self): - """Return deterministic allocated-memory value for debug logging.""" - return 0.0 - - def _max_memory_allocated(self): - """Return deterministic max-memory value for debug logging.""" - return 0.0 - - def _mlx_vae_decode(self, latents): - """Return deterministic decoded waveform for MLX decode branch.""" - _ = latents - return torch.ones(1, 2, 8) - - def tiled_decode(self, latents): - """Return deterministic decoded waveform for tiled decode branch.""" - _ = latents - return torch.ones(1, 2, 8) - - -class GenerateMusicDecodeMixinTests(unittest.TestCase): - """Verify decode-state preparation and latent decode helper behavior.""" - - def test_prepare_decode_state_updates_progress_estimates(self): - """It updates timing fields and progress estimate metadata for valid latents.""" - host = _Host() - outputs = { - "target_latents": torch.ones(1, 4, 3), - "time_costs": {"total_time_cost": 1.0, "diffusion_per_step_time_cost": 0.2}, - } - pred_latents, time_costs = host._prepare_generate_music_decode_state( - outputs=outputs, - infer_steps_for_progress=8, - actual_batch_size=1, - audio_duration=12.0, - latent_shift=0.0, - latent_rescale=1.0, - ) - self.assertEqual(tuple(pred_latents.shape), (1, 4, 3)) - self.assertEqual(time_costs["offload_time_cost"], 0.25) - self.assertEqual(host._last_diffusion_per_step_sec, 0.2) - self.assertEqual(host.estimate_calls[0]["infer_steps"], 8) - - def test_prepare_decode_state_raises_for_nan_latents(self): - """It raises runtime error when diffusion latents contain NaN values.""" - host = _Host() - outputs = { - "target_latents": torch.tensor([[[float("nan")]]]), - "time_costs": {"total_time_cost": 1.0}, - } - with self.assertRaises(RuntimeError): - host._prepare_generate_music_decode_state( - outputs=outputs, - infer_steps_for_progress=8, - actual_batch_size=1, - audio_duration=None, - latent_shift=0.0, - latent_rescale=1.0, - ) - - def test_decode_pred_latents_updates_decode_time_and_returns_cpu_latents(self): - """It decodes latents and updates decode timing metrics in time_costs.""" - host = _Host() - pred_latents = torch.ones(1, 4, 3) - time_costs = {"total_time_cost": 1.0} - - def _progress(value, desc=None): - """Capture progress updates for assertions.""" - host.progress_calls.append((value, desc)) - - with patch.object(GENERATE_MUSIC_DECODE_MODULE.time, "time", side_effect=[10.0, 11.5]): - pred_wavs, pred_latents_cpu, updated_costs = host._decode_generate_music_pred_latents( - pred_latents=pred_latents, - progress=_progress, - use_tiled_decode=False, - time_costs=time_costs, - ) - - self.assertEqual(tuple(pred_wavs.shape), (1, 2, 8)) - self.assertEqual(pred_latents_cpu.device.type, "cpu") - self.assertAlmostEqual(updated_costs["vae_decode_time_cost"], 1.5, places=6) - self.assertAlmostEqual(updated_costs["total_time_cost"], 2.5, places=6) - self.assertAlmostEqual(updated_costs["offload_time_cost"], 0.25, places=6) - self.assertEqual(host.progress_calls[0][0], 0.8) - - def test_decode_pred_latents_restores_vae_device_on_decode_error(self): - """It restores VAE device in the CPU-offload path even when decode raises.""" - - class _FailingVae(_FakeVae): - """VAE double that raises during decode and records transfer calls.""" - - def __init__(self): - """Initialize transfer call trackers for restoration assertions.""" - super().__init__() - self.cpu_calls = 0 - self.to_calls = [] - - def decode(self, latents: torch.Tensor): - """Raise decode error to exercise restoration in finally branch.""" - _ = latents - raise RuntimeError("decode failed") - - def cpu(self): - """Record explicit CPU transfer and return self.""" - self.cpu_calls += 1 - return self - - def to(self, *args, **kwargs): - """Record restore transfer target and return self.""" - self.to_calls.append((args, kwargs)) - return self - - class _FailingHost(_Host): - """Host variant that forces non-MLX VAE decode and tracks cache clears.""" - - def __init__(self): - """Set non-MLX state so CPU offload path is exercised deterministically.""" - super().__init__() - self.use_mlx_vae = False - self.mlx_vae = None - self.vae = _FailingVae() - self.empty_cache_calls = 0 - - def _empty_cache(self): - """Count cache-clear calls to verify finally cleanup runs.""" - self.empty_cache_calls += 1 - - host = _FailingHost() - pred_latents = torch.ones(1, 4, 3) - time_costs = {"total_time_cost": 1.0} - - with patch.dict(GENERATE_MUSIC_DECODE_MODULE.os.environ, {"ACESTEP_VAE_ON_CPU": "1"}, clear=False): - with self.assertRaisesRegex(RuntimeError, "decode failed"): - host._decode_generate_music_pred_latents( - pred_latents=pred_latents, - progress=None, - use_tiled_decode=False, - time_costs=time_costs, - ) - - self.assertEqual(host.vae.cpu_calls, 1) - self.assertEqual(len(host.vae.to_calls), 1) - self.assertGreaterEqual(host.empty_cache_calls, 2) - - - def test_decode_pred_latents_does_not_restore_latents_to_gpu_after_successful_cpu_decode(self): - """It does not move pred_latents_for_decode back to GPU after a successful CPU decode. - - The removed ``pred_latents_for_decode = pred_latents_for_decode.to(vae_device)`` line - was causing a wasteful re-allocation of the already-decoded input tensor on the - GPU. After the fix, only the VAE itself is restored; the input latent is not. - """ - - class _SuccessVae(_FakeVae): - """VAE double that records transfer calls and succeeds on decode.""" - - def __init__(self): - """Initialize transfer call trackers.""" - super().__init__() - self.vae_to_calls = [] - self._device = "cuda" - - def decode(self, latents: torch.Tensor): - """Return a simple decoded output.""" - return _FakeDecodeOutput(torch.ones(latents.shape[0], 2, 8)) - - def cpu(self): - """Simulate VAE being moved to CPU.""" - self._device = "cpu" - return self - - def to(self, *args, **kwargs): - """Record VAE device-transfer destinations.""" - self.vae_to_calls.append(args[0] if args else kwargs) - return self - - class _SuccessHost(_Host): - """Host that forces non-MLX VAE so the CPU-decode path is exercised.""" - - def __init__(self): - """Configure non-MLX state and a tracking VAE.""" - super().__init__() - self.use_mlx_vae = False - self.mlx_vae = None - self.vae = _SuccessVae() - self.device = "cuda" - - host = _SuccessHost() - pred_latents = torch.ones(1, 4, 3) - time_costs = {"total_time_cost": 1.0} - - with patch.dict(GENERATE_MUSIC_DECODE_MODULE.os.environ, {"ACESTEP_VAE_ON_CPU": "1"}, clear=False): - pred_wavs, _cpu_latents, _costs = host._decode_generate_music_pred_latents( - pred_latents=pred_latents, - progress=None, - use_tiled_decode=False, - time_costs=time_costs, - ) - - # VAE itself must be restored to its original device. - self.assertEqual(len(host.vae.vae_to_calls), 1) - # The decoded waveform must be returned correctly. - self.assertEqual(tuple(pred_wavs.shape), (1, 2, 8)) - - -if __name__ == "__main__": - unittest.main() - diff --git a/acestep/core/generation/handler/generate_music_decode_test_support.py b/acestep/core/generation/handler/generate_music_decode_test_support.py new file mode 100644 index 000000000..222b85192 --- /dev/null +++ b/acestep/core/generation/handler/generate_music_decode_test_support.py @@ -0,0 +1,77 @@ +"""Shared fixtures for ``generate_music_decode`` mixin tests.""" + +from contextlib import contextmanager + +import acestep.core.generation.handler.generate_music_decode as GENERATE_MUSIC_DECODE_MODULE +import torch +from acestep.core.generation.handler.generate_music_decode import GenerateMusicDecodeMixin + + +class FakeDecodeOutput: + """Minimal VAE decode output container exposing ``sample`` attribute.""" + + def __init__(self, sample: torch.Tensor): + self.sample = sample + + +class FakeVae: + """Minimal VAE stand-in with dtype, decode, and parameter iteration hooks.""" + + def __init__(self): + self.dtype = torch.float32 + self._param = torch.nn.Parameter(torch.zeros(1)) + + def decode(self, latents: torch.Tensor): + return FakeDecodeOutput(torch.ones(latents.shape[0], 2, 8)) + + def parameters(self): + yield self._param + + def cpu(self): + return self + + def to(self, *_args, **_kwargs): + return self + + +class DecodeTestHost(GenerateMusicDecodeMixin): + """Minimal decode-mixin host exposing deterministic state for assertions.""" + + def __init__(self): + self.current_offload_cost = 0.25 + self.debug_stats = False + self._last_diffusion_per_step_sec = None + self.estimate_calls = [] + self.progress_calls = [] + self.device = "cpu" + self.use_mlx_vae = True + self.mlx_vae = object() + self.vae = FakeVae() + + def _update_progress_estimate(self, **kwargs): + self.estimate_calls.append(kwargs) + + @contextmanager + def _load_model_context(self, _model_name): + yield + + def _empty_cache(self): + return None + + def _memory_allocated(self): + return 0.0 + + def _max_memory_allocated(self): + return 0.0 + + def _get_component_device(self, component: str) -> str: + _ = component + return self.device + + def _mlx_vae_decode(self, latents): + _ = latents + return torch.ones(1, 2, 8) + + def tiled_decode(self, latents): + _ = latents + return torch.ones(1, 2, 8) diff --git a/acestep/core/generation/handler/generate_music_decode_vram_test.py b/acestep/core/generation/handler/generate_music_decode_vram_test.py new file mode 100644 index 000000000..5721a42ce --- /dev/null +++ b/acestep/core/generation/handler/generate_music_decode_vram_test.py @@ -0,0 +1,53 @@ +"""Tests for multi-GPU VRAM preflight in ``generate_music`` decode.""" + +import unittest +from unittest.mock import MagicMock, patch + +from acestep.core.generation.handler.generate_music_decode_test_support import ( + GENERATE_MUSIC_DECODE_MODULE, + DecodeTestHost, +) + + +class GenerateMusicDecodeVramTests(unittest.TestCase): + """Verify decode VRAM checks honor mapped component devices.""" + + def test_decode_pred_latents_queries_vram_on_mapped_vae_cuda_index(self): + class VramHost(DecodeTestHost): + def __init__(self): + super().__init__() + self.use_mlx_vae = False + self.mlx_vae = None + self.device = "cuda:0" + self.vae_component_device = "cuda:2" + + def _get_component_device(self, component: str) -> str: + if component == "vae": + return self.vae_component_device + return self.device + + host = VramHost() + latent = MagicMock() + latent.detach.return_value = latent + latent.transpose.return_value = latent + latent.contiguous.return_value = latent + latent.to.return_value = latent + + with patch.object( + GENERATE_MUSIC_DECODE_MODULE, + "get_effective_free_vram_gb", + return_value=8.0, + ) as free_mock: + with patch.object(GENERATE_MUSIC_DECODE_MODULE.time, "time", side_effect=[10.0, 11.0]): + host._decode_generate_music_pred_latents( + pred_latents=latent, + progress=None, + use_tiled_decode=False, + time_costs={"total_time_cost": 1.0}, + ) + + free_mock.assert_called_once_with(2) + + +if __name__ == "__main__": + unittest.main() diff --git a/acestep/core/generation/handler/init_service_catalog.py b/acestep/core/generation/handler/init_service_catalog.py index 9fe339cb0..2fddb76df 100644 --- a/acestep/core/generation/handler/init_service_catalog.py +++ b/acestep/core/generation/handler/init_service_catalog.py @@ -6,6 +6,8 @@ import torch from loguru import logger +from acestep.device_map import cuda_device_index, is_cuda_device + class InitServiceCatalogMixin: """Checkpoint discovery and backend capability helpers.""" @@ -47,16 +49,20 @@ def get_available_acestep_v15_models(self) -> List[str]: def is_flash_attention_available(self, device: Optional[str] = None) -> bool: """Check whether flash attention can be used on the target device.""" - target_device = str(device or self.device or "auto").split(":", 1)[0] + target_device = str(device or self.device or "auto") if target_device == "auto": if not torch.cuda.is_available(): return False - else: - if target_device != "cuda" or not torch.cuda.is_available(): + cuda_index = 0 + elif is_cuda_device(target_device): + if not torch.cuda.is_available(): return False + cuda_index = cuda_device_index(target_device) + else: + return False try: - major, _ = torch.cuda.get_device_capability() + major, _ = torch.cuda.get_device_capability(cuda_index) if major < 8: logger.info( f"[is_flash_attention_available] GPU compute capability {major}.x < 8.0 " diff --git a/acestep/core/generation/handler/init_service_loader.py b/acestep/core/generation/handler/init_service_loader.py index 040c5ce11..921328667 100644 --- a/acestep/core/generation/handler/init_service_loader.py +++ b/acestep/core/generation/handler/init_service_loader.py @@ -8,6 +8,7 @@ from loguru import logger from acestep import gpu_config +from acestep.device_map import cuda_device_index, is_cuda_device from .init_service_loader_components import InitServiceLoaderComponentsMixin @@ -143,7 +144,9 @@ def _load_main_model_from_checkpoint( 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 is_cuda_device(device) and not gpu_config.cuda_supports_bfloat16( + cuda_device_index(device) + ): # 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 diff --git a/acestep/core/generation/handler/init_service_memory_basic.py b/acestep/core/generation/handler/init_service_memory_basic.py index b83fa52ea..d0f5a2999 100644 --- a/acestep/core/generation/handler/init_service_memory_basic.py +++ b/acestep/core/generation/handler/init_service_memory_basic.py @@ -11,6 +11,9 @@ import torch from loguru import logger +from acestep.device_map.devices import normalize_component_device +from acestep.device_map.errors import DeviceMapError + # Cached libc handle for mallopt/malloc_trim calls (Linux only). _LIBC = None _MALLOPT_APPLIED = False @@ -101,7 +104,7 @@ def _is_on_target_device(self, tensor, target_device): target_type = target_device.type else: target_type = torch.device(str(target_device)).type - except Exception: + except (RuntimeError, TypeError, ValueError): target_type = str(target_device).strip().lower().split(":", 1)[0] if not target_type: logger.warning( @@ -111,6 +114,16 @@ def _is_on_target_device(self, tensor, target_device): return False return tensor.device.type == target_type + 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 + try: + expected = torch.device(normalize_component_device(str(target_device))) + except (DeviceMapError, RuntimeError, TypeError): + return False + return tensor.device == expected + @staticmethod def _get_affine_quantized_tensor_class(): """Return the AffineQuantizedTensor class from torchao, or None if unavailable.""" @@ -152,10 +165,15 @@ def _has_quantized_params(self, module): return False def _ensure_silence_latent_on_device(self): - """Ensure ``silence_latent`` is on ``self.device``.""" + """Ensure ``silence_latent`` is on the active DiT device.""" if hasattr(self, "silence_latent") and self.silence_latent is not None: - if not self._is_on_target_device(self.silence_latent, self.device): - self.silence_latent = self.silence_latent.to(self.device).to(self.dtype) + target = ( + self._get_component_device("dit") + if getattr(self, "device_map", None) is not None + else self.device + ) + if not self._tensor_on_exact_device(self.silence_latent, target): + self.silence_latent = self.silence_latent.to(target).to(self.dtype) @staticmethod def _get_rss_mb() -> float: diff --git a/acestep/core/generation/handler/init_service_offload_context.py b/acestep/core/generation/handler/init_service_offload_context.py index 92af32fa7..f148d50f8 100644 --- a/acestep/core/generation/handler/init_service_offload_context.py +++ b/acestep/core/generation/handler/init_service_offload_context.py @@ -16,17 +16,21 @@ def _load_model_context(self, model_name: str): yield return + target_device = self._get_component_device(model_name) + if model_name == "model" and not self.offload_dit_to_cpu: model = getattr(self, model_name, None) if model is not None: try: param = next(model.parameters()) if param.device.type == "cpu": - logger.info(f"[_load_model_context] Moving {model_name} to {self.device} (persistent)") - self._recursive_to_device(model, self.device, self.dtype) + logger.info( + f"[_load_model_context] Moving {model_name} to {target_device} (persistent)" + ) + self._recursive_to_device(model, target_device, self.dtype) self._release_system_memory() if hasattr(self, "silence_latent"): - self.silence_latent = self.silence_latent.to(self.device).to(self.dtype) + self.silence_latent = self.silence_latent.to(target_device).to(self.dtype) except StopIteration: pass yield @@ -38,16 +42,18 @@ def _load_model_context(self, model_name: str): return rss_before = self._get_rss_mb() - logger.info(f"[_load_model_context] Loading {model_name} to {self.device} (RSS: {rss_before:.0f} MB)") + logger.info( + f"[_load_model_context] Loading {model_name} to {target_device} (RSS: {rss_before:.0f} MB)" + ) start_time = time.time() if model_name == "vae": - vae_dtype = self._get_vae_dtype() - self._recursive_to_device(model, self.device, vae_dtype) + vae_dtype = self._get_vae_dtype(target_device) + self._recursive_to_device(model, target_device, vae_dtype) else: - self._recursive_to_device(model, self.device, self.dtype) + self._recursive_to_device(model, target_device, self.dtype) if model_name == "model" and hasattr(self, "silence_latent"): - self.silence_latent = self.silence_latent.to(self.device).to(self.dtype) + self.silence_latent = self.silence_latent.to(target_device).to(self.dtype) load_time = time.time() - start_time self.current_offload_cost += load_time @@ -56,7 +62,7 @@ def _load_model_context(self, model_name: str): self._release_system_memory() rss_after = self._get_rss_mb() logger.info( - f"[_load_model_context] Loaded {model_name} to {self.device} in {load_time:.4f}s " + f"[_load_model_context] Loaded {model_name} to {target_device} in {load_time:.4f}s " f"(RSS: {rss_before:.0f} -> {rss_after:.0f} MB, delta: {rss_after - rss_before:+.0f} MB)" ) diff --git a/acestep/core/generation/handler/init_service_orchestrator.py b/acestep/core/generation/handler/init_service_orchestrator.py index 4b0322470..064ba684e 100644 --- a/acestep/core/generation/handler/init_service_orchestrator.py +++ b/acestep/core/generation/handler/init_service_orchestrator.py @@ -9,6 +9,7 @@ from loguru import logger from acestep import gpu_config +from acestep.device_map import cuda_device_index, device_type, is_cuda_device, set_active_cuda_device _ROCM_DTYPE_MAP = { "float32": torch.float32, @@ -58,6 +59,7 @@ def initialize_service( prefer_source: Optional[str] = None, use_mlx_dit: bool = True, vae_checkpoint: Optional[str] = None, + gpu_mapping: Optional[str] = None, ) -> Tuple[str, bool]: """Initialize model artifacts and runtime backends for generation. @@ -72,24 +74,37 @@ def initialize_service( ) resolved_device = self._resolve_initialize_device(device) - self.device = resolved_device + if gpu_mapping is None: + gpu_mapping = os.environ.get("ACESTEP_GPU_MAPPING") + lm_model_path = os.environ.get("ACESTEP_LM_MODEL_PATH") + self.device_map = self._resolve_component_device_map( + resolved_device=resolved_device, + gpu_mapping=gpu_mapping, + config_path=config_path, + lm_model_path=lm_model_path, + ) + self.device = self.device_map.dit self.offload_to_cpu = offload_to_cpu self.offload_dit_to_cpu = offload_dit_to_cpu + if is_cuda_device(self.device): + set_active_cuda_device(self.device) + normalized_compile, normalized_quantization, mlx_compile_requested = self._configure_initialize_runtime( device=resolved_device, compile_model=compile_model, quantization=quantization, ) self.compiled = normalized_compile - if resolved_device == "cuda" and gpu_config.is_rocm_available(): + if is_cuda_device(resolved_device) and gpu_config.is_rocm_available(): self.dtype = _resolve_rocm_dtype() logger.info( f"[initialize_service] ROCm/HIP device detected: using dtype={self.dtype} " "(set ACESTEP_ROCM_DTYPE=bfloat16 or float16 to override)" ) - elif resolved_device == "cuda": - if gpu_config.cuda_supports_bfloat16(): + elif is_cuda_device(resolved_device): + dit_cuda_index = cuda_device_index(self.device) + if gpu_config.cuda_supports_bfloat16(dit_cuda_index): self.dtype = torch.bfloat16 else: self.dtype = torch.float16 @@ -98,7 +113,7 @@ def initialize_service( "using float16 instead of bfloat16." ) else: - self.dtype = torch.bfloat16 if resolved_device == "xpu" else torch.float32 + self.dtype = torch.bfloat16 if device_type(resolved_device) == "xpu" else torch.float32 self.quantization = normalized_quantization try: self._validate_quantization_setup( @@ -155,20 +170,20 @@ def initialize_service( model_path = os.path.join(checkpoint_dir, config_path) self._load_main_model_from_checkpoint( model_checkpoint_path=model_path, - device=resolved_device, + device=self.device_map.dit, use_flash_attention=use_flash_attention, compile_model=normalized_compile, quantization=self.quantization, ) vae_path = self._load_vae_model( checkpoint_dir=checkpoint_dir, - device=resolved_device, + device=self.device_map.vae, compile_model=normalized_compile, vae_variant=resolved_vae_variant, ) text_encoder_path = self._load_text_encoder_and_tokenizer( checkpoint_dir=checkpoint_dir, - device=resolved_device, + device=self.device_map.text_encoder, ) mlx_dit_status, mlx_vae_status = self._initialize_mlx_backends( @@ -178,7 +193,7 @@ def initialize_service( ) status_msg = self._build_initialize_status_message( - device=resolved_device, + device=self.device, model_path=model_path, vae_path=vae_path, text_encoder_path=text_encoder_path, @@ -192,11 +207,14 @@ def initialize_service( mlx_dit_status=mlx_dit_status, mlx_vae_status=mlx_vae_status, ) + if self.device_map.is_multi_device(): + status_msg += f"\nGPU mapping: {self.device_map.summary()}" self.last_init_params = { "project_root": project_root, "config_path": config_path, - "device": resolved_device, + "device": self.device, + "gpu_mapping": self.device_map.summary(), "use_flash_attention": use_flash_attention, "compile_model": normalized_compile, "offload_to_cpu": offload_to_cpu, diff --git a/acestep/core/generation/handler/init_service_setup.py b/acestep/core/generation/handler/init_service_setup.py index c7e7932bf..a0fbc2f9d 100644 --- a/acestep/core/generation/handler/init_service_setup.py +++ b/acestep/core/generation/handler/init_service_setup.py @@ -1,11 +1,18 @@ """Runtime setup helpers for initialization orchestration.""" -from typing import Any, Optional, Tuple +from typing import Any, Dict, Optional, Tuple import torch from loguru import logger from acestep import gpu_config +from acestep.device_map import ( + ComponentDeviceMap, + is_cuda_device, + log_device_map, + normalize_component_device, + resolve_component_device_map, +) class InitServiceSetupMixin: @@ -23,15 +30,23 @@ def _resolve_initialize_device(self, requested_device: str) -> str: return "xpu" return "cpu" - if device == "cuda" and not gpu_config.is_cuda_available(): - if gpu_config.is_mps_available(): - logger.warning("[initialize_service] CUDA requested but unavailable. Falling back to MPS.") - return "mps" - if gpu_config.is_xpu_available(): - logger.warning("[initialize_service] CUDA requested but unavailable. Falling back to XPU.") - return "xpu" - logger.warning("[initialize_service] CUDA requested but unavailable. Falling back to CPU.") - return "cpu" + if is_cuda_device(device): + if not gpu_config.is_cuda_available(): + if gpu_config.is_mps_available(): + logger.warning( + "[initialize_service] CUDA device requested but unavailable. Falling back to MPS." + ) + return "mps" + if gpu_config.is_xpu_available(): + logger.warning( + "[initialize_service] CUDA device requested but unavailable. Falling back to XPU." + ) + return "xpu" + logger.warning( + "[initialize_service] CUDA device requested but unavailable. Falling back to CPU." + ) + return "cpu" + return normalize_component_device(device) if device == "mps" and not gpu_config.is_mps_available(): if gpu_config.is_cuda_available(): @@ -55,6 +70,115 @@ def _resolve_initialize_device(self, requested_device: str) -> str: return device + def _resolve_component_device_map( + self, + *, + resolved_device: str, + gpu_mapping: Optional[str] = None, + config_path: Optional[str] = None, + lm_model_path: Optional[str] = None, + use_lm: bool = True, + batch_size: int = 1, + ) -> ComponentDeviceMap: + """Resolve per-component device placement for initialization.""" + device_map = resolve_component_device_map( + requested_device=resolved_device, + gpu_mapping=gpu_mapping, + config_path=config_path, + lm_model_path=lm_model_path, + use_lm=use_lm, + batch_size=batch_size, + ) + log_device_map(device_map) + return device_map + + def _get_component_device(self, component: str) -> str: + """Return the device string for a model component.""" + device_map = getattr(self, "device_map", None) + if device_map is None: + return self.device + if component == "model": + component = "dit" + return device_map.device_for(component) + + def _to_component_device(self, value: Any, component: str) -> Any: + """Move a tensor to the device assigned to *component* when needed.""" + if value is None or not isinstance(value, torch.Tensor): + return value + target = self._get_component_device(component) + device = torch.device(target) + if value.device != device: + return value.to(device) + return value + + def _route_service_payload_to_dit(self, payload: Dict[str, Any]) -> Dict[str, Any]: + """Move service-generation tensors onto the DiT device for diffusion.""" + device_map = getattr(self, "device_map", None) + if device_map is None or not device_map.is_multi_device(): + return payload + + routed = dict(payload) + for key in ( + "text_hidden_states", + "text_attention_mask", + "lyric_hidden_states", + "lyric_attention_mask", + "refer_audio_acoustic_hidden_states_packed", + "refer_audio_order_mask", + "src_latents", + "chunk_mask", + "is_covers", + "precomputed_lm_hints_25Hz", + "non_cover_text_hidden_states", + "non_cover_text_attention_masks", + "repaint_mask", + "target_latents", + ): + if key in routed: + routed[key] = self._to_component_device(routed[key], "dit") + return routed + + def _route_service_generate_kwargs_to_dit( + self, + generate_kwargs: Dict[str, Any], + payload: Dict[str, Any], + ) -> Dict[str, Any]: + """Move diffusion kwargs tensors onto the DiT device for multi-GPU runs.""" + device_map = getattr(self, "device_map", None) + if device_map is None or not device_map.is_multi_device(): + return generate_kwargs + + routed = dict(generate_kwargs) + payload_field_map = { + "text_hidden_states": "text_hidden_states", + "text_attention_mask": "text_attention_mask", + "lyric_hidden_states": "lyric_hidden_states", + "lyric_attention_mask": "lyric_attention_mask", + "refer_audio_acoustic_hidden_states_packed": "refer_audio_acoustic_hidden_states_packed", + "refer_audio_order_mask": "refer_audio_order_mask", + "src_latents": "src_latents", + "chunk_masks": "chunk_mask", + "is_covers": "is_covers", + "non_cover_text_hidden_states": "non_cover_text_hidden_states", + "non_cover_text_attention_mask": "non_cover_text_attention_masks", + "precomputed_lm_hints_25Hz": "precomputed_lm_hints_25Hz", + "repaint_mask": "repaint_mask", + "clean_src_latents": "target_latents", + } + for kwarg_key, payload_key in payload_field_map.items(): + if payload_key in payload: + routed[kwarg_key] = payload[payload_key] + + for key in ( + "silence_latent", + "timesteps", + "repaint_mask", + "clean_src_latents", + ): + if key in routed: + routed[key] = self._to_component_device(routed[key], "dit") + return routed + def _configure_initialize_runtime( self, *, diff --git a/acestep/core/generation/handler/init_service_test.py b/acestep/core/generation/handler/init_service_test.py index 7f70bdaba..257dd0dc9 100644 --- a/acestep/core/generation/handler/init_service_test.py +++ b/acestep/core/generation/handler/init_service_test.py @@ -120,6 +120,35 @@ def test_is_on_target_device_malformed_target_logs_and_returns_false(self): self.assertFalse(host._is_on_target_device(t, ":0")) warning.assert_called_once() + def test_tensor_on_exact_device_distinguishes_cuda_indices(self): + """It compares full CUDA device strings, not just backend types.""" + host = _Host(project_root="K:/fake_root", device="cuda:0") + tensor = types.SimpleNamespace(device=torch.device("cuda:0")) + self.assertTrue(host._tensor_on_exact_device(tensor, "cuda:0")) + self.assertFalse(host._tensor_on_exact_device(tensor, "cuda:3")) + self.assertTrue(host._is_on_target_device(tensor, "cuda:3")) + + def test_ensure_silence_latent_on_device_moves_across_cuda_indices(self): + """``silence_latent`` must follow ``device_map.dit``, not just the CUDA backend.""" + host = _Host(project_root="K:/fake_root", device="cuda:3") + host.dtype = torch.float32 + host.device_map = host._resolve_component_device_map( + resolved_device="cuda:0", + gpu_mapping="dit:3,vae:0,text_encoder:0,lm:1", + ) + moved_to = [] + + class _Latent: + device = torch.device("cuda:0") + + def to(self, device_or_dtype, *args, **kwargs): + moved_to.append(device_or_dtype) + return self + + host.silence_latent = _Latent() + host._ensure_silence_latent_on_device() + self.assertEqual(moved_to[0], "cuda:3") + self.assertEqual(len(moved_to), 2) def test_get_auto_decode_chunk_size_uses_cuda_device_index(self): """It probes effective VRAM on the selected CUDA device index.""" @@ -261,6 +290,90 @@ def test_resolve_initialize_device_auto_prefers_cuda(self): with patch("torch.xpu", new=types.SimpleNamespace(is_available=lambda: False), create=True): self.assertEqual(host._resolve_initialize_device("auto"), "cuda") + def test_resolve_initialize_device_preserves_cuda_index(self): + """It preserves explicit CUDA device indices such as ``cuda:1``.""" + host = _Host(project_root="K:/fake_root", device="auto") + with patch("torch.cuda.is_available", return_value=True): + self.assertEqual(host._resolve_initialize_device("cuda:1"), "cuda:1") + + def test_resolve_component_device_map_uses_explicit_mapping(self): + """It resolves per-component devices from a GPU mapping string.""" + host = _Host(project_root="K:/fake_root", device="cuda:0") + device_map = host._resolve_component_device_map( + resolved_device="cuda:0", + gpu_mapping="dit:0,vae:0,text_encoder:0,lm:1", + ) + self.assertEqual(device_map.dit, "cuda:0") + self.assertEqual(device_map.lm, "cuda:1") + self.assertTrue(device_map.is_multi_device()) + + def test_get_component_device_returns_component_specific_device(self): + """It returns the mapped device for offload/load helpers.""" + host = _Host(project_root="K:/fake_root", device="cuda:0") + host.device_map = host._resolve_component_device_map( + resolved_device="cuda:0", + gpu_mapping="dit:0,vae:2,text_encoder:0,lm:1", + ) + self.assertEqual(host._get_component_device("model"), "cuda:0") + self.assertEqual(host._get_component_device("vae"), "cuda:2") + + def test_route_service_payload_to_dit_moves_tensors(self): + """It routes diffusion tensors onto the DiT device in multi-GPU mode.""" + host = _Host(project_root="K:/fake_root", device="cuda:0") + host.device_map = host._resolve_component_device_map( + resolved_device="cuda:0", + gpu_mapping="dit:0,vae:0,text_encoder:0,lm:1", + ) + tensor = object() + payload = { + "text_hidden_states": tensor, + "ignored": "keep-me", + } + with patch.object(host, "_to_component_device", side_effect=lambda value, component: value) as move_mock: + routed = host._route_service_payload_to_dit(payload) + move_mock.assert_called_once_with(tensor, "dit") + self.assertIs(routed["text_hidden_states"], tensor) + self.assertEqual(routed["ignored"], "keep-me") + + def test_route_service_generate_kwargs_to_dit_syncs_from_routed_payload(self): + """Diffusion kwargs should mirror routed payload tensors in multi-GPU mode.""" + host = _Host(project_root="K:/fake_root", device="cuda:3") + host.device_map = host._resolve_component_device_map( + resolved_device="cuda:0", + gpu_mapping="dit:3,vae:0,text_encoder:0,lm:1", + ) + src_on_dit = object() + payload = { + "text_hidden_states": object(), + "text_attention_mask": object(), + "lyric_hidden_states": object(), + "lyric_attention_mask": object(), + "refer_audio_acoustic_hidden_states_packed": object(), + "refer_audio_order_mask": object(), + "src_latents": src_on_dit, + "chunk_mask": object(), + "is_covers": object(), + "non_cover_text_hidden_states": None, + "non_cover_text_attention_masks": None, + "precomputed_lm_hints_25Hz": None, + "target_latents": object(), + } + generate_kwargs = { + "text_hidden_states": object(), + "src_latents": object(), + "chunk_masks": object(), + "silence_latent": object(), + "timesteps": object(), + } + with patch.object(host, "_to_component_device", side_effect=lambda value, component: f"moved:{component}") as move_mock: + routed_kwargs = host._route_service_generate_kwargs_to_dit(generate_kwargs, payload) + self.assertIs(routed_kwargs["src_latents"], src_on_dit) + self.assertEqual(routed_kwargs["chunk_masks"], payload["chunk_mask"]) + self.assertEqual(routed_kwargs["clean_src_latents"], "moved:dit") + self.assertEqual(routed_kwargs["silence_latent"], "moved:dit") + self.assertEqual(routed_kwargs["timesteps"], "moved:dit") + self.assertEqual(move_mock.call_count, 3) + def test_configure_initialize_runtime_redirects_compile_on_mps(self): """It converts MPS compile intent to MLX compile and disables quantization.""" host = _Host(project_root="K:/fake_root", device="mps") @@ -464,6 +577,7 @@ def _fake_load_main_model(**_kwargs): "project_root", "config_path", "device", + "gpu_mapping", "use_flash_attention", "compile_model", "offload_to_cpu", @@ -476,6 +590,7 @@ def _fake_load_main_model(**_kwargs): self.assertEqual(set(host.last_init_params.keys()), expected_keys) self.assertEqual(host.last_init_params["config_path"], "acestep-v15-turbo") self.assertEqual(host.last_init_params["device"], "cpu") + self.assertEqual(host.last_init_params["gpu_mapping"], "dit:cpu, vae:cpu, text_encoder:cpu, lm:cpu") self.assertEqual(host.last_init_params["vae_checkpoint"], "official") ensure_models.assert_called_once() sync_code.assert_called_once() @@ -918,6 +1033,44 @@ def test_resolve_rocm_dtype_unknown_value_falls_back_to_float32(self): result = ORCHESTRATOR_MODULE._resolve_rocm_dtype() self.assertEqual(result, torch.float32) + def test_initialize_service_records_multi_gpu_mapping(self): + """It records an explicit GPU mapping in init params and status output.""" + host = _Host(project_root="K:/fake_root", device="cuda:0") + + def _fake_load_main_model(**_kwargs): + host.config = types.SimpleNamespace(_attn_implementation="sdpa") + host.model = object() + + with patch.object(GPU_CONFIG_MODULE, "is_cuda_available", return_value=True), \ + patch.object(GPU_CONFIG_MODULE, "is_rocm_available", return_value=False), \ + patch.object(GPU_CONFIG_MODULE, "cuda_supports_bfloat16", return_value=True): + with patch.object(host, "_ensure_models_present", return_value=None): + with patch.object(host, "_sync_model_code_if_needed"): + with patch.object(host, "_load_main_model_from_checkpoint", side_effect=_fake_load_main_model) as load_main_mock: + with patch.object(host, "_load_vae_model", return_value="vae") as load_vae_mock: + with patch.object( + host, + "_load_text_encoder_and_tokenizer", + return_value="te", + ) as load_text_mock: + with patch.object( + host, + "_initialize_mlx_backends", + return_value=("Disabled", "Disabled"), + ): + status, ok = host.initialize_service( + project_root="K:/fake_root", + config_path="acestep-v15-turbo", + device="cuda:0", + gpu_mapping="dit:0,vae:0,text_encoder:0,lm:1", + ) + self.assertTrue(ok) + self.assertIn("GPU mapping: dit:0, vae:0, text_encoder:0, lm:1", status) + self.assertEqual(host.last_init_params["gpu_mapping"], "dit:0, vae:0, text_encoder:0, lm:1") + self.assertEqual(load_main_mock.call_args.kwargs["device"], "cuda:0") + self.assertEqual(load_vae_mock.call_args.kwargs["device"], "cuda:0") + self.assertEqual(load_text_mock.call_args.kwargs["device"], "cuda:0") + def test_initialize_service_uses_float32_on_rocm(self): """It sets dtype=float32 during initialization when ROCm is active.""" host = self._make_rocm_host() diff --git a/acestep/core/generation/handler/service_generate_execute.py b/acestep/core/generation/handler/service_generate_execute.py index 46c6b0a4f..0500acdfc 100644 --- a/acestep/core/generation/handler/service_generate_execute.py +++ b/acestep/core/generation/handler/service_generate_execute.py @@ -132,7 +132,16 @@ def _build_service_generate_kwargs( "retake_variance": retake_variance, } if timesteps is not None: - kwargs["timesteps"] = torch.tensor(timesteps, dtype=torch.float32, device=self.device) + dit_device = ( + self._get_component_device("dit") + if getattr(self, "device_map", None) is not None + else self.device + ) + kwargs["timesteps"] = torch.tensor( + timesteps, + dtype=torch.float32, + device=dit_device, + ) return kwargs def _execute_service_generate_diffusion( @@ -155,6 +164,10 @@ def _execute_service_generate_diffusion( self, payload=payload, generate_kwargs=generate_kwargs, seed_param=seed_param, flow_edit_ctx=flow_edit_ctx, ) + payload = self._route_service_payload_to_dit(payload) + generate_kwargs = self._route_service_generate_kwargs_to_dit(generate_kwargs, payload) + if hasattr(self, "silence_latent") and self.silence_latent is not None: + self.silence_latent = self._to_component_device(self.silence_latent, "dit") dit_backend = ( "MLX (native)" if (self.use_mlx_dit and self.mlx_decoder is not None) else f"PyTorch ({self.device})" ) diff --git a/acestep/core/generation/handler/service_generate_flow_edit_source.py b/acestep/core/generation/handler/service_generate_flow_edit_source.py index 638e10b1b..9f408af39 100644 --- a/acestep/core/generation/handler/service_generate_flow_edit_source.py +++ b/acestep/core/generation/handler/service_generate_flow_edit_source.py @@ -77,13 +77,9 @@ def embed_source( ) -> Tuple[torch.Tensor, torch.Tensor]: """Run text + lyric encoders on the source tokens. - Tokens come back from the tokenizer on CPU; the regular batch path - moves them to ``handler.device`` before encoding (see - ``preprocess_batch``), so we mirror that here. + Tokens may live on CPU or the DiT device; ``infer_*_embeddings`` + moves them onto the text-encoder component device. """ - device = handler.device - text_token_idss = text_token_idss.to(device=device) - lyric_token_idss = lyric_token_idss.to(device=device) with handler._load_model_context("text_encoder"): text_hs = handler.infer_text_embeddings(text_token_idss) lyric_hs = handler.infer_lyric_embeddings(lyric_token_idss) diff --git a/acestep/device_map/__init__.py b/acestep/device_map/__init__.py new file mode 100644 index 000000000..b992a0c0a --- /dev/null +++ b/acestep/device_map/__init__.py @@ -0,0 +1,64 @@ +""" +Component-level device placement for multi-GPU inference. + +Public API facade — import from ``acestep.device_map`` as before. +""" + +from acestep.device_map.constants import GPU_MAPPING_ENV, LM_DEVICE_ENV +from acestep.device_map.devices import ( + cuda_device_index, + device_index_label, + device_type, + is_cuda_device, + normalize_component_device, + set_active_cuda_device, +) +from acestep.device_map.discovery import discover_gpus, format_gpu_list_text +from acestep.device_map.errors import DeviceMapError +from acestep.device_map.layout import ( + compute_auto_device_map, + estimate_dit_peak_gb, + estimate_lm_total_gb, +) +from acestep.device_map.parsing import parse_gpu_mapping +from acestep.device_map.resolve import log_device_map, resolve_component_device_map +from acestep.device_map.status import ( + collect_gpu_runtime_status, + device_map_to_dict, + gpu_info_to_dict, + log_lm_device_deprecation, +) +from acestep.device_map.types import ( + ComponentDeviceMap, + GpuInfo, + LayoutError, + LayoutRequest, +) + +__all__ = [ + "GPU_MAPPING_ENV", + "LM_DEVICE_ENV", + "ComponentDeviceMap", + "DeviceMapError", + "GpuInfo", + "LayoutError", + "LayoutRequest", + "collect_gpu_runtime_status", + "compute_auto_device_map", + "cuda_device_index", + "device_index_label", + "device_map_to_dict", + "device_type", + "discover_gpus", + "estimate_dit_peak_gb", + "estimate_lm_total_gb", + "format_gpu_list_text", + "gpu_info_to_dict", + "is_cuda_device", + "log_device_map", + "log_lm_device_deprecation", + "normalize_component_device", + "parse_gpu_mapping", + "resolve_component_device_map", + "set_active_cuda_device", +] diff --git a/acestep/device_map/constants.py b/acestep/device_map/constants.py new file mode 100644 index 000000000..c45d91dd9 --- /dev/null +++ b/acestep/device_map/constants.py @@ -0,0 +1,12 @@ +"""Shared constants for multi-GPU device placement.""" + +from __future__ import annotations + +import re + +GPU_MAPPING_ENV = "ACESTEP_GPU_MAPPING" +LM_DEVICE_ENV = "ACESTEP_LM_DEVICE" + +COMPONENT_KEYS = ("dit", "vae", "text_encoder", "lm") +SINGLE_PATTERN = re.compile(r"^single:(\d+)$") +PAIR_PATTERN = re.compile(r"^([a-z_]+):(\d+)$") diff --git a/acestep/device_map/devices.py b/acestep/device_map/devices.py new file mode 100644 index 000000000..430f19163 --- /dev/null +++ b/acestep/device_map/devices.py @@ -0,0 +1,58 @@ +"""Device-string normalization and CUDA helpers.""" + +from __future__ import annotations + +from acestep.device_map.errors import DeviceMapError + + +def device_type(device: str) -> str: + """Return the backend type token from a device string.""" + return str(device).split(":", 1)[0] + + +def cuda_device_index(device: str) -> int: + """Return the CUDA device index for a device string.""" + normalized = str(device) + if normalized == "cuda": + return 0 + if normalized.startswith("cuda:"): + suffix = normalized.split(":", 1)[1] + try: + return int(suffix) + except ValueError as exc: + raise DeviceMapError(f"Not a CUDA device string: {device!r}") from exc + raise DeviceMapError(f"Not a CUDA device string: {device!r}") + + +def is_cuda_device(device: str) -> bool: + """Return whether *device* refers to a CUDA backend.""" + return device_type(device) == "cuda" + + +def normalize_component_device(device: str) -> str: + """Normalize device strings used for component placement.""" + value = str(device).strip() + if not value: + raise DeviceMapError("Device string cannot be empty") + if value == "cuda": + return "cuda:0" + return value + + +def device_index_label(device: str) -> str: + """Return a short index label for logs/UI.""" + if is_cuda_device(device): + return str(cuda_device_index(device)) + return device_type(device) + + +def set_active_cuda_device(device: str) -> None: + """Set the active CUDA device when loading or running on a specific GPU.""" + if not is_cuda_device(device): + return + import torch + + if not torch.cuda.is_available(): + return + index = cuda_device_index(device) + torch.cuda.set_device(index) diff --git a/acestep/device_map/discovery.py b/acestep/device_map/discovery.py new file mode 100644 index 000000000..84d3e15fa --- /dev/null +++ b/acestep/device_map/discovery.py @@ -0,0 +1,60 @@ +"""CUDA device discovery helpers.""" + +from __future__ import annotations + +from typing import List, Optional + +from acestep.device_map.types import GpuInfo + + +def discover_gpus() -> List[GpuInfo]: + """Enumerate visible CUDA devices and their current VRAM stats.""" + import torch + + if not torch.cuda.is_available(): + return [] + + gpus: List[GpuInfo] = [] + for index in range(torch.cuda.device_count()): + props = torch.cuda.get_device_properties(index) + free_bytes, total_bytes = torch.cuda.mem_get_info(index) + try: + capability = torch.cuda.get_device_capability(index) + except RuntimeError: + capability = None + gpus.append( + GpuInfo( + logical_index=index, + name=props.name, + total_vram_gb=round(total_bytes / (1024**3), 2), + free_vram_gb=round(free_bytes / (1024**3), 2), + compute_capability=capability, + ) + ) + return gpus + + +def format_gpu_list_text(gpus: Optional[List[GpuInfo]] = None) -> str: + """Return a human-readable table of visible CUDA devices.""" + visible = gpus if gpus is not None else discover_gpus() + if not visible: + return "No CUDA devices detected." + + lines = [ + "Visible CUDA devices:", + " idx name total(GB) free(GB) compute", + ] + for gpu in visible: + capability = "" + if gpu.compute_capability is not None: + capability = f"{gpu.compute_capability[0]}.{gpu.compute_capability[1]}" + lines.append( + " {idx:>3} {name:<28} {total:>8.2f} {free:>7.2f} {capability}".format( + idx=gpu.logical_index, + name=gpu.name[:28], + total=gpu.total_vram_gb, + free=gpu.free_vram_gb, + capability=capability or "-", + ) + ) + return "\n".join(lines) diff --git a/acestep/device_map/errors.py b/acestep/device_map/errors.py new file mode 100644 index 000000000..12a671972 --- /dev/null +++ b/acestep/device_map/errors.py @@ -0,0 +1,5 @@ +"""Exceptions raised by device-map parsing and resolution.""" + + +class DeviceMapError(ValueError): + """Raised when a GPU mapping string cannot be parsed or validated.""" diff --git a/acestep/device_map/layout.py b/acestep/device_map/layout.py new file mode 100644 index 000000000..039086abb --- /dev/null +++ b/acestep/device_map/layout.py @@ -0,0 +1,114 @@ +"""VRAM-aware automatic multi-GPU layout.""" + +from __future__ import annotations + +from typing import List, Optional, Union + +from acestep.gpu_config import ( + DIT_INFERENCE_VRAM_PER_BATCH, + LM_VRAM, + MODEL_VRAM, + VRAM_SAFETY_MARGIN_GB, + get_lm_model_size, +) + +from acestep.device_map.types import ComponentDeviceMap, GpuInfo, LayoutError, LayoutRequest + + +def _dit_model_vram_key(dit_type: str) -> str: + if dit_type.startswith("dit_"): + return dit_type + return f"dit_{dit_type}" + + +def estimate_dit_peak_gb(dit_type: str, batch_size: int) -> float: + """Estimate DiT GPU memory including co-located VAE and text encoder.""" + model_key = _dit_model_vram_key(dit_type) + weights = MODEL_VRAM.get(model_key, MODEL_VRAM["dit_turbo"]) + per_batch = DIT_INFERENCE_VRAM_PER_BATCH.get( + dit_type, + DIT_INFERENCE_VRAM_PER_BATCH["turbo"], + ) + aux = ( + MODEL_VRAM["vae"] + + MODEL_VRAM["text_encoder"] + + MODEL_VRAM["cuda_context"] + + VRAM_SAFETY_MARGIN_GB + ) + return weights + (per_batch * max(1, batch_size)) + aux + + +def estimate_lm_total_gb(lm_model_path: Optional[str]) -> float: + """Estimate LM weights plus KV cache on the target GPU.""" + model_size = get_lm_model_size(lm_model_path or "acestep-5Hz-lm-1.7B") + lm_info = LM_VRAM.get(model_size, LM_VRAM["0.6B"]) + return lm_info["weights"] + lm_info["kv_cache_4k"] + 0.3 + + +def _sort_gpus_for_layout(gpus: List[GpuInfo]) -> List[GpuInfo]: + return sorted( + gpus, + key=lambda gpu: (-gpu.free_vram_gb, -gpu.total_vram_gb, gpu.logical_index), + ) + + +def _lm_free_vram_gb(gpu: GpuInfo, dit_gpu: GpuInfo, dit_need_gb: float) -> float: + """Return free VRAM on *gpu* available for LM after reserving the DiT stack.""" + if gpu.logical_index != dit_gpu.logical_index: + return gpu.free_vram_gb + return max(0.0, gpu.free_vram_gb - dit_need_gb) + + +def compute_auto_device_map(request: LayoutRequest) -> Union[ComponentDeviceMap, LayoutError]: + """Place components across visible CUDA devices based on free VRAM.""" + gpus = _sort_gpus_for_layout(request.gpus) + if not gpus: + return LayoutError( + "No CUDA devices detected for auto layout", + suggestions=("Install CUDA drivers or pass an explicit --gpu-mapping",), + ) + + dit_need_gb = estimate_dit_peak_gb(request.dit_type, request.batch_size) + dit_gpu = next((gpu for gpu in gpus if gpu.free_vram_gb >= dit_need_gb), None) + if dit_gpu is None: + return LayoutError( + f"No GPU has {dit_need_gb:.1f}GB free for DiT ({request.dit_type})", + suggestions=( + "Reduce batch size", + "Use a smaller DiT checkpoint", + "Enable CPU offload", + "Pass an explicit gpu mapping", + ), + ) + + dit_device = f"cuda:{dit_gpu.logical_index}" + lm_device = None + if request.use_lm: + lm_need_gb = estimate_lm_total_gb(request.lm_model_path) + lm_candidates = [gpu for gpu in gpus if gpu.logical_index != dit_gpu.logical_index] + lm_candidates.extend(gpu for gpu in gpus if gpu.logical_index == dit_gpu.logical_index) + lm_gpu = next( + ( + gpu + for gpu in lm_candidates + if _lm_free_vram_gb(gpu, dit_gpu, dit_need_gb) >= lm_need_gb + ), + None, + ) + if lm_gpu is None: + return LayoutError( + f"No GPU has {lm_need_gb:.1f}GB free for LM ({request.lm_model_path or 'default'})", + suggestions=( + "Use a smaller LM model", + "Pass gpu_mapping with lm on a specific GPU", + "Disable LM initialization", + ), + ) + lm_device = f"cuda:{lm_gpu.logical_index}" + + return ComponentDeviceMap( + dit=dit_device, + vae=dit_device, + text_encoder=dit_device, + lm=lm_device, + ) diff --git a/acestep/device_map/parsing.py b/acestep/device_map/parsing.py new file mode 100644 index 000000000..74fde9200 --- /dev/null +++ b/acestep/device_map/parsing.py @@ -0,0 +1,116 @@ +"""GPU mapping string parsing.""" + +from __future__ import annotations + +import os +from typing import Dict, Optional + +from acestep.device_map.constants import COMPONENT_KEYS, GPU_MAPPING_ENV, PAIR_PATTERN, SINGLE_PATTERN +from acestep.device_map.devices import device_type +from acestep.device_map.errors import DeviceMapError +from acestep.device_map.types import ComponentDeviceMap + + +def _format_device_for_backend(backend: str, index: int) -> str: + if backend == "cuda": + return f"cuda:{index}" + if backend in {"mps", "xpu", "cpu"}: + if index != 0: + raise DeviceMapError( + f"Component index {index} is invalid for backend '{backend}'" + ) + return backend if backend != "xpu" else "xpu:0" + raise DeviceMapError(f"Unsupported backend for component mapping: {backend!r}") + + +def _parse_mapping_pairs(mapping: str) -> Dict[str, int]: + pairs: Dict[str, int] = {} + for raw_part in mapping.split(","): + part = raw_part.strip() + if not part: + continue + match = PAIR_PATTERN.fullmatch(part) + if match is None: + raise DeviceMapError( + f"Invalid mapping segment {part!r}; expected 'component:index'" + ) + component, index_text = match.group(1), match.group(2) + if component not in COMPONENT_KEYS: + raise DeviceMapError( + f"Unknown component {component!r}; expected one of {COMPONENT_KEYS}" + ) + if component in pairs: + raise DeviceMapError(f"Duplicate component entry: {component!r}") + pairs[component] = int(index_text) + if not pairs: + raise DeviceMapError("GPU mapping must include at least one component") + return pairs + + +_SUPPORTED_MAPPING_BACKENDS = frozenset({"cuda", "mps", "xpu", "cpu"}) + + +def _resolve_mapping_backend(default_device: str) -> str: + """Return the backend token used to interpret explicit GPU indices.""" + backend = device_type(default_device) + if backend in _SUPPORTED_MAPPING_BACKENDS: + return backend + raise DeviceMapError( + f"Unsupported default device {default_device!r} for GPU mapping; " + "expected cuda, mps, xpu, or cpu" + ) + + +def parse_gpu_mapping( + mapping: Optional[str], + *, + default_device: str, +) -> Optional[ComponentDeviceMap]: + """ + Parse a GPU mapping string into a component device map. + + Returns None when mapping is unset or explicitly set to 'auto'. + """ + raw = (mapping or "").strip() + if not raw: + raw = os.environ.get(GPU_MAPPING_ENV, "").strip() + if not raw or raw.lower() == "auto": + return None + + backend = _resolve_mapping_backend(default_device) + + single_match = SINGLE_PATTERN.fullmatch(raw) + if single_match is not None: + index = int(single_match.group(1)) + device = _format_device_for_backend(backend, index) + return ComponentDeviceMap.from_single_device(device) + + pairs = _parse_mapping_pairs(raw) + dit_index = pairs.get("dit") + if dit_index is None: + raise DeviceMapError("Explicit GPU mapping must include a 'dit' component") + + dit_device = _format_device_for_backend(backend, dit_index) + vae_device = _format_device_for_backend(backend, pairs.get("vae", dit_index)) + text_encoder_device = _format_device_for_backend( + backend, + pairs.get("text_encoder", dit_index), + ) + lm_device = None + if "lm" in pairs: + lm_device = _format_device_for_backend(backend, pairs["lm"]) + + return ComponentDeviceMap( + dit=dit_device, + vae=vae_device, + text_encoder=text_encoder_device, + lm=lm_device, + ) + + +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 diff --git a/acestep/device_map/resolve.py b/acestep/device_map/resolve.py new file mode 100644 index 000000000..be3011506 --- /dev/null +++ b/acestep/device_map/resolve.py @@ -0,0 +1,70 @@ +"""Resolve the effective component layout for service initialization.""" + +from __future__ import annotations + +import os +from typing import Optional + +from loguru import logger + +from acestep.gpu_config import get_dit_type_from_path + +from acestep.device_map.devices import is_cuda_device, normalize_component_device +from acestep.device_map.discovery import discover_gpus +from acestep.device_map.layout import compute_auto_device_map +from acestep.device_map.parsing import parse_gpu_mapping, raw_gpu_mapping_value +from acestep.device_map.types import ComponentDeviceMap, LayoutRequest + + +def resolve_component_device_map( + *, + requested_device: str, + gpu_mapping: Optional[str] = None, + config_path: Optional[str] = None, + lm_model_path: Optional[str] = None, + use_lm: bool = True, + batch_size: int = 1, +) -> ComponentDeviceMap: + """ + Resolve the effective component device map for service initialization. + + When no mapping is provided, all components share ``requested_device``. + When mapping is ``auto`` and multiple CUDA devices are visible, compute a + VRAM-aware layout; otherwise fall back to single-device placement. + """ + normalized_device = normalize_component_device(requested_device) + raw_mapping = raw_gpu_mapping_value(gpu_mapping) + + if raw_mapping.lower() == "auto" and is_cuda_device(normalized_device): + gpus = discover_gpus() + if len(gpus) >= 2: + dit_type = get_dit_type_from_path(config_path or "") + layout = compute_auto_device_map( + LayoutRequest( + gpus=gpus, + dit_type=dit_type, + lm_model_path=lm_model_path, + use_lm=use_lm, + batch_size=batch_size, + ) + ) + if isinstance(layout, ComponentDeviceMap): + logger.info("[device_map] Auto layout selected from {} GPU(s)", len(gpus)) + return layout + logger.warning( + "[device_map] Auto layout failed: {}. Suggestions: {}", + layout.message, + "; ".join(layout.suggestions), + ) + + parsed = parse_gpu_mapping(gpu_mapping, default_device=normalized_device) + if parsed is None: + return ComponentDeviceMap.from_single_device(normalized_device) + return parsed + + +def log_device_map(device_map: ComponentDeviceMap) -> None: + """Emit a startup log line describing the active component layout.""" + logger.info("[device_map] Active layout: {}", device_map.summary()) + if device_map.is_multi_device(): + logger.info("[device_map] Multi-device placement enabled with cross-GPU routing") diff --git a/acestep/device_map/status.py b/acestep/device_map/status.py new file mode 100644 index 000000000..841b385f9 --- /dev/null +++ b/acestep/device_map/status.py @@ -0,0 +1,81 @@ +"""CLI/API status serialization and deprecation helpers.""" + +from __future__ import annotations + +import os +from typing import Any, Dict, Optional + +from loguru import logger + +from acestep.device_map.constants import GPU_MAPPING_ENV, LM_DEVICE_ENV +from acestep.device_map.discovery import discover_gpus +from acestep.device_map.types import ComponentDeviceMap, GpuInfo + + +def gpu_info_to_dict(gpu: GpuInfo) -> Dict[str, object]: + """Serialize a :class:`GpuInfo` record for CLI and API responses.""" + payload: Dict[str, object] = { + "index": gpu.logical_index, + "name": gpu.name, + "total_vram_gb": gpu.total_vram_gb, + "free_vram_gb": gpu.free_vram_gb, + } + if gpu.compute_capability is not None: + payload["compute_capability"] = list(gpu.compute_capability) + return payload + + +def device_map_to_dict(device_map: ComponentDeviceMap) -> Dict[str, object]: + """Serialize an active component layout for status endpoints.""" + return { + "dit": device_map.dit, + "vae": device_map.vae, + "text_encoder": device_map.text_encoder, + "lm": device_map.lm, + "lm_tensor_parallel": device_map.lm_tensor_parallel, + "summary": device_map.summary(), + "multi_device": device_map.is_multi_device(), + } + + +def collect_gpu_runtime_status(handler: Any = None) -> Dict[str, object]: + """Collect GPU inventory and the active component layout for API status.""" + gpus = [gpu_info_to_dict(gpu) for gpu in discover_gpus()] + device_map = getattr(handler, "device_map", None) if handler is not None else None + mapping_env = os.environ.get(GPU_MAPPING_ENV, "").strip() or None + return { + "gpus": gpus, + "gpu_mapping": mapping_env, + "device_map": device_map_to_dict(device_map) if device_map is not None else None, + } + + +def log_lm_device_deprecation( + *, + explicit_lm_device: Optional[str] = None, + gpu_mapping_env: Optional[str] = None, + using_device_map_lm: bool = False, +) -> None: + """Warn when legacy ``ACESTEP_LM_DEVICE`` should be replaced by gpu mapping.""" + explicit = (explicit_lm_device or os.getenv(LM_DEVICE_ENV, "")).strip() + if not explicit: + return + + mapping = (gpu_mapping_env or os.environ.get(GPU_MAPPING_ENV, "")).strip() + if mapping and using_device_map_lm: + logger.warning( + "[device_map] ACESTEP_LM_DEVICE={} is ignored because {}={} " + "assigns the LM device. Use 'lm:N' in the mapping instead.", + explicit, + GPU_MAPPING_ENV, + mapping, + ) + return + + if not mapping: + logger.warning( + "[device_map] ACESTEP_LM_DEVICE is deprecated; use {} instead " + "(e.g. 'single:1' or 'dit:0,vae:0,text_encoder:0,lm:1'). Current value: {}", + GPU_MAPPING_ENV, + explicit, + ) diff --git a/acestep/device_map/types.py b/acestep/device_map/types.py new file mode 100644 index 000000000..0e8437dc4 --- /dev/null +++ b/acestep/device_map/types.py @@ -0,0 +1,93 @@ +"""Datatypes describing GPUs and per-component device layouts.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import List, Optional, Tuple + +from acestep.device_map.devices import device_index_label, normalize_component_device +from acestep.device_map.constants import COMPONENT_KEYS + + +@dataclass(frozen=True) +class GpuInfo: + """Describes one visible CUDA device using logical indices.""" + + logical_index: int + name: str + total_vram_gb: float + free_vram_gb: float + compute_capability: Optional[Tuple[int, int]] = None + + +@dataclass(frozen=True) +class ComponentDeviceMap: + """Per-component runtime device strings for inference.""" + + dit: str + vae: str + text_encoder: str + lm: Optional[str] = None + lm_tensor_parallel: int = 1 + + 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}") + value = getattr(self, key) + if value is None: + raise ValueError(f"Component '{component}' has no assigned device") + return value + + def is_multi_device(self) -> bool: + """Return True when components target more than one distinct device.""" + devices = {self.dit, self.vae, self.text_encoder} + if self.lm is not None: + devices.add(self.lm) + return len(devices) > 1 + + def summary(self) -> str: + """Return a compact human-readable layout description.""" + parts = [ + f"dit:{device_index_label(self.dit)}", + f"vae:{device_index_label(self.vae)}", + f"text_encoder:{device_index_label(self.text_encoder)}", + ] + if self.lm is not None: + parts.append(f"lm:{device_index_label(self.lm)}") + if self.lm_tensor_parallel > 1: + parts.append(f"lm_tp={self.lm_tensor_parallel}") + return ", ".join(parts) + + @classmethod + def from_single_device(cls, device: str) -> "ComponentDeviceMap": + """Place all inference components on one device.""" + normalized = normalize_component_device(device) + return cls( + dit=normalized, + vae=normalized, + text_encoder=normalized, + lm=normalized, + ) + + +@dataclass(frozen=True) +class LayoutRequest: + """Inputs used to compute an automatic multi-GPU layout.""" + + gpus: List[GpuInfo] + dit_type: str = "turbo" + lm_model_path: Optional[str] = None + use_lm: bool = True + batch_size: int = 1 + + +@dataclass(frozen=True) +class LayoutError: + """Returned when automatic layout cannot place all requested components.""" + + message: str + suggestions: Tuple[str, ...] = field(default_factory=tuple) diff --git a/acestep/gpu_config.py b/acestep/gpu_config.py index 8b1b4b5e5..f835c03ad 100644 --- a/acestep/gpu_config.py +++ b/acestep/gpu_config.py @@ -974,7 +974,11 @@ def find_best_lm_model_on_disk( def get_lm_gpu_memory_ratio( - model_path: str, total_gpu_memory_gb: float + model_path: str, + total_gpu_memory_gb: float, + device_index: int = 0, + *, + reserve_dit_inference_gb: float = 1.5, ) -> Tuple[float, float]: """ Calculate GPU memory utilization ratio for LM model. @@ -987,6 +991,9 @@ def get_lm_gpu_memory_ratio( Args: model_path: LM model path (e.g., "acestep-5Hz-lm-0.6B") total_gpu_memory_gb: Total GPU memory in GB (used as fallback) + device_index: CUDA device index to query for free VRAM + reserve_dit_inference_gb: Headroom to leave on the LM GPU for DiT activations + when DiT is co-located on the same device (set to 0 when split across GPUs) Returns: Tuple of (gpu_memory_utilization_ratio, target_memory_gb) @@ -1008,7 +1015,7 @@ def get_lm_gpu_memory_ratio( import torch if torch.cuda.is_available(): - free_bytes, total_bytes = torch.cuda.mem_get_info() + free_bytes, total_bytes = torch.cuda.mem_get_info(device_index) free_gb = free_bytes / (1024**3) actual_total_gb = total_bytes / (1024**3) @@ -1031,8 +1038,7 @@ def get_lm_gpu_memory_ratio( # The ratio is relative to total GPU memory (nano-vllm convention), # but we compute it so that the LM only claims what's actually free # minus a safety margin for DiT inference activations. - # Reserve at least 1.5 GB for DiT inference activations - dit_reserve_gb = 1.5 + dit_reserve_gb = max(0.0, reserve_dit_inference_gb) usable_for_lm = max(0, free_gb - dit_reserve_gb - VRAM_SAFETY_MARGIN_GB) # Cap to what the LM actually needs diff --git a/acestep/llm_backend_compat.py b/acestep/llm_backend_compat.py index 7bd6a2751..b2c4f9001 100644 --- a/acestep/llm_backend_compat.py +++ b/acestep/llm_backend_compat.py @@ -3,6 +3,8 @@ import importlib import sys +from acestep.device_map import is_cuda_device + def _has_working_triton_installation() -> bool: """Return whether the Triton modules required by nano-vllm import cleanly.""" @@ -25,7 +27,7 @@ def get_vllm_preflight_warning(*, device: str, platform: str | None = None) -> s A warning string when vLLM should fall back to PyTorch, otherwise ``None``. """ active_platform = sys.platform if platform is None else platform - if device != "cuda" or active_platform != "win32": + if not is_cuda_device(device) or active_platform != "win32": return None if _has_working_triton_installation(): return None diff --git a/acestep/llm_backend_compat_test.py b/acestep/llm_backend_compat_test.py index 4dab0887a..f1faf5b85 100644 --- a/acestep/llm_backend_compat_test.py +++ b/acestep/llm_backend_compat_test.py @@ -46,9 +46,22 @@ def test_get_vllm_preflight_warning_returns_none_outside_windows_cuda(self) -> N ): warning = get_vllm_preflight_warning(device="cpu", platform="win32") linux_warning = get_vllm_preflight_warning(device="cuda", platform="linux") + indexed_linux = get_vllm_preflight_warning(device="cuda:1", platform="linux") self.assertIsNone(warning) self.assertIsNone(linux_warning) + self.assertIsNone(indexed_linux) + + def test_get_vllm_preflight_warning_accepts_indexed_cuda_on_windows(self) -> None: + """Indexed CUDA devices must still trigger the Windows Triton preflight.""" + with patch( + "acestep.llm_backend_compat._has_working_triton_installation", + return_value=False, + ): + warning = get_vllm_preflight_warning(device="cuda:1", platform="win32") + + self.assertIsNotNone(warning) + self.assertIn("Windows", warning) @unittest.skipIf(LLMHandler is None, f"llm_inference import unavailable: {_IMPORT_ERROR}") diff --git a/acestep/llm_inference.py b/acestep/llm_inference.py index 3690af56d..7cdb12933 100644 --- a/acestep/llm_inference.py +++ b/acestep/llm_inference.py @@ -26,6 +26,12 @@ from acestep.constrained_logits_processor import MetadataConstrainedLogitsProcessor from acestep.constants import DEFAULT_LM_INSTRUCTION, DEFAULT_LM_UNDERSTAND_INSTRUCTION, DEFAULT_LM_INSPIRED_INSTRUCTION, DEFAULT_LM_REWRITE_INSTRUCTION, DURATION_MIN, DURATION_MAX from acestep.gpu_config import get_lm_gpu_memory_ratio, get_gpu_memory_gb, get_lm_model_size, get_global_gpu_config +from acestep.device_map import ( + cuda_device_index, + is_cuda_device, + normalize_component_device, + set_active_cuda_device, +) # Minimum free VRAM (GB) required to attempt vLLM initialization. # vLLM's KV cache allocator adapts to available memory, so we only need a @@ -199,15 +205,23 @@ def get_gpu_memory_utilization(self, model_path: str = None, minimal_gpu: float Tuple of (gpu_memory_utilization_ratio, low_gpu_memory_mode) """ try: - device = torch.device("cuda:0") - total_gpu_mem_bytes = torch.cuda.get_device_properties(device).total_memory + if not is_cuda_device(self.device) or not torch.cuda.is_available(): + raise RuntimeError("CUDA device required for vLLM memory budgeting") + + device_index = cuda_device_index(normalize_component_device(self.device)) + total_gpu_mem_bytes = torch.cuda.get_device_properties(device_index).total_memory total_gpu = total_gpu_mem_bytes / 1024**3 low_gpu_memory_mode = False # Use adaptive GPU memory ratio based on model size if model_path: - ratio, target_memory_gb = get_lm_gpu_memory_ratio(model_path, total_gpu) + ratio, target_memory_gb = get_lm_gpu_memory_ratio( + model_path, + total_gpu, + device_index=device_index, + reserve_dit_inference_gb=0.0, + ) logger.info(f"Adaptive LM memory allocation: model={model_path}, target={target_memory_gb}GB, ratio={ratio:.3f}, total_gpu={total_gpu:.1f}GB") # Enable low memory mode for small GPUs @@ -217,7 +231,8 @@ def get_gpu_memory_utilization(self, model_path: str = None, minimal_gpu: float return ratio, low_gpu_memory_mode # Fallback to original logic if no model_path provided - reserved_mem_bytes = torch.cuda.memory_reserved(device) + device_index = cuda_device_index(normalize_component_device(self.device)) + reserved_mem_bytes = torch.cuda.memory_reserved(device_index) reserved_gpu = reserved_mem_bytes / 1024**3 available_gpu = total_gpu - reserved_gpu @@ -529,7 +544,7 @@ def initialize( device = "xpu" else: device = "cpu" - elif device == "cuda" and not torch.cuda.is_available(): + elif is_cuda_device(device) and not torch.cuda.is_available(): if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): logger.warning("[initialize] CUDA requested but unavailable. Falling back to MPS.") device = "mps" @@ -560,6 +575,10 @@ def initialize( logger.warning("[initialize] XPU requested but unavailable. Falling back to CPU.") device = "cpu" + # Normalize bare "cuda" → "cuda:0" and preserve mapped "cuda:N". + if is_cuda_device(device): + device = normalize_component_device(device) + self.device = device self.offload_to_cpu = offload_to_cpu @@ -569,7 +588,7 @@ def initialize( # produce NaN/inf when naively converted to float16 (different exponent range). # The DiT and VAE use float16 on MPS where it actually helps throughput. if dtype is None: - if device in ["cuda", "xpu"]: + if is_cuda_device(device) or device.startswith("xpu"): self.dtype = torch.bfloat16 else: self.dtype = torch.float32 @@ -601,7 +620,8 @@ def initialize( } # Proactive CUDA cleanup before LM load to reduce fragmentation on mode/model switch - if device == "cuda" and torch.cuda.is_available(): + if is_cuda_device(device) and torch.cuda.is_available(): + set_active_cuda_device(device) gc.collect() torch.cuda.empty_cache() torch.cuda.synchronize() @@ -641,9 +661,10 @@ def initialize( # RuntimeError: Offset increment outside graph capture encountered unexpectedly is_rocm = hasattr(torch.version, 'hip') and torch.version.hip is not None is_jetson = False - if device == "cuda" and torch.cuda.is_available(): + if is_cuda_device(device) and torch.cuda.is_available(): try: - dev_name = torch.cuda.get_device_name(0).lower() + dit_idx = cuda_device_index(device) + dev_name = torch.cuda.get_device_name(dit_idx).lower() is_jetson = any(k in dev_name for k in ("orin", "xavier", "tegra")) if is_jetson: logger.info(f"Jetson GPU detected ({dev_name}): disabling CUDA graph capture for nano-vllm") @@ -702,7 +723,7 @@ def initialize( status_msg = f"✅ 5Hz LM initialized (PyTorch fallback, MLX not available)\nModel: {full_lm_model_path}\nBackend: PyTorch" return status_msg, True - if backend == "vllm" and device != "cuda": + if backend == "vllm" and not is_cuda_device(device): logger.info( f"[initialize] vllm backend requires CUDA, using PyTorch backend for device={device}." ) @@ -720,19 +741,23 @@ def initialize( # Initialize based on user-selected backend if backend == "vllm": _warn_if_prerelease_python() - total_gb = get_gpu_memory_gb() if device == "cuda" else 0.0 + total_gb = 0.0 free_gb = 0.0 - if device == "cuda" and torch.cuda.is_available(): + if is_cuda_device(device) and torch.cuda.is_available(): try: + device_index = cuda_device_index(device) + total_gb = get_gpu_memory_gb() if hasattr(torch.cuda, "mem_get_info"): - free_bytes, _ = torch.cuda.mem_get_info() + free_bytes, _ = torch.cuda.mem_get_info(device_index) free_gb = free_bytes / (1024**3) else: - total_bytes = torch.cuda.get_device_properties(0).total_memory - free_gb = (total_bytes - torch.cuda.memory_reserved(0)) / (1024**3) + total_bytes = torch.cuda.get_device_properties(device_index).total_memory + free_gb = ( + total_bytes - torch.cuda.memory_reserved(device_index) + ) / (1024**3) except Exception: free_gb = 0.0 - if device == "cuda" and free_gb < VRAM_SAFE_FREE_GB: + if is_cuda_device(device) and free_gb < VRAM_SAFE_FREE_GB: logger.warning( f"vLLM disabled due to insufficient free VRAM (total={total_gb:.2f}GB, free={free_gb:.2f}GB, need>={VRAM_SAFE_FREE_GB}GB free) — falling back to PyTorch backend" ) @@ -801,6 +826,8 @@ def _initialize_5hz_lm_vllm(self, model_path: str, enforce_eager: bool = False, return "❌ nano-vllm is not installed. Please install it using 'cd acestep/third_parts/nano-vllm && pip install .'" try: + if is_cuda_device(self.device): + set_active_cuda_device(self.device) current_device = torch.cuda.current_device() device_name = torch.cuda.get_device_name(current_device) diff --git a/acestep/llm_inference_cuda_index_test.py b/acestep/llm_inference_cuda_index_test.py new file mode 100644 index 000000000..d355614f6 --- /dev/null +++ b/acestep/llm_inference_cuda_index_test.py @@ -0,0 +1,63 @@ +"""Unit tests for indexed CUDA device resolution in LLMHandler.initialize.""" + +from __future__ import annotations + +import unittest +from unittest.mock import patch + +try: + from acestep.llm_inference import LLMHandler + from acestep.device_map import is_cuda_device + _IMPORT_ERROR = None +except ImportError as exc: # pragma: no cover + LLMHandler = None + is_cuda_device = None + _IMPORT_ERROR = exc + + +@unittest.skipIf(LLMHandler is None, f"llm_inference import unavailable: {_IMPORT_ERROR}") +class TestLmInitializeCudaIndex(unittest.TestCase): + """Verify auto/cuda:N normalization for LM multi-GPU placement.""" + + def _resolved_device(self, device: str) -> str: + """Run initialize far enough to resolve self.device, then abort cleanly.""" + handler = LLMHandler() + with patch("acestep.llm_inference.os.path.exists", return_value=False): + status, ok = handler.initialize( + checkpoint_dir="/tmp/checkpoints", + lm_model_path="acestep-5Hz-lm-1.7B", + backend="pt", + device=device, + offload_to_cpu=False, + dtype=None, + ) + self.assertFalse(ok) + self.assertIn("not found", status) + return handler.device + + @patch("acestep.llm_inference.torch.cuda.is_available", return_value=True) + def test_auto_normalizes_to_cuda0(self, _mock_cuda): + """auto on CUDA must become cuda:0, not bare cuda.""" + self.assertEqual(self._resolved_device("auto"), "cuda:0") + + @patch("acestep.llm_inference.torch.cuda.is_available", return_value=True) + def test_bare_cuda_normalizes_to_cuda0(self, _mock_cuda): + """Bare cuda must normalize to cuda:0.""" + self.assertEqual(self._resolved_device("cuda"), "cuda:0") + + @patch("acestep.llm_inference.torch.cuda.is_available", return_value=True) + def test_mapped_cuda_index_preserved(self, _mock_cuda): + """Explicit cuda:N from device_map must be preserved.""" + self.assertEqual(self._resolved_device("cuda:1"), "cuda:1") + + def test_vllm_gate_accepts_indexed_cuda(self): + """vLLM CUDA gate must treat cuda:N as CUDA (not force PT).""" + device = "cuda:1" + backend = "vllm" + if backend == "vllm" and not is_cuda_device(device): + backend = "pt" + self.assertEqual(backend, "vllm") + + +if __name__ == "__main__": + unittest.main() diff --git a/acestep/models/common/apg_guidance.py b/acestep/models/common/apg_guidance.py index 3f114a281..0889420d2 100644 --- a/acestep/models/common/apg_guidance.py +++ b/acestep/models/common/apg_guidance.py @@ -19,15 +19,18 @@ def project( dims=[-1], ): dtype = v0.dtype - device_type = v0.device.type - if device_type == "mps": + target_device = v0.device + if target_device.type == "mps": v0, v1 = v0.cpu(), v1.cpu() v0, v1 = v0.double(), v1.double() v1 = torch.nn.functional.normalize(v1, dim=dims) v0_parallel = (v0 * v1).sum(dim=dims, keepdim=True) * v1 v0_orthogonal = v0 - v0_parallel - return v0_parallel.to(dtype).to(device_type), v0_orthogonal.to(dtype).to(device_type) + return ( + v0_parallel.to(device=target_device, dtype=dtype), + v0_orthogonal.to(device=target_device, dtype=dtype), + ) def apg_forward( diff --git a/acestep/models/common/apg_guidance_test.py b/acestep/models/common/apg_guidance_test.py new file mode 100644 index 000000000..4c38d3bb8 --- /dev/null +++ b/acestep/models/common/apg_guidance_test.py @@ -0,0 +1,72 @@ +"""Unit tests for APG guidance helpers.""" + +from __future__ import annotations + +import unittest + +import torch + +from acestep.models.common.apg_guidance import MomentumBuffer, apg_forward, project + + +class ApgGuidanceDeviceTests(unittest.TestCase): + """Verify guidance math preserves tensor device placement.""" + + def test_project_preserves_cpu_device(self): + v0 = torch.randn(2, 4, 8) + v1 = torch.randn(2, 4, 8) + parallel, orthogonal = project(v0, v1, dims=[1]) + self.assertEqual(parallel.device, v0.device) + self.assertEqual(orthogonal.device, v0.device) + + def test_apg_forward_preserves_cpu_device(self): + pred_cond = torch.randn(2, 4, 8) + pred_uncond = pred_cond + 0.25 + guided = apg_forward( + pred_cond=pred_cond, + pred_uncond=pred_uncond, + guidance_scale=7.0, + momentum_buffer=MomentumBuffer(), + dims=[1], + ) + self.assertEqual(guided.device, pred_cond.device) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") + def test_project_preserves_non_default_cuda_index(self): + if torch.cuda.device_count() < 2: + self.skipTest("Need at least 2 CUDA devices") + device = torch.device("cuda:1") + v0 = torch.randn(2, 4, 8, device=device) + v1 = torch.randn(2, 4, 8, device=device) + parallel, orthogonal = project(v0, v1, dims=[1]) + self.assertEqual(parallel.device, device) + self.assertEqual(orthogonal.device, device) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") + def test_apg_forward_preserves_non_default_cuda_index(self): + if torch.cuda.device_count() < 2: + self.skipTest("Need at least 2 CUDA devices") + device = torch.device("cuda:1") + pred_cond = torch.randn(2, 4, 8, device=device) + pred_uncond = pred_cond + 0.25 + guided = apg_forward( + pred_cond=pred_cond, + pred_uncond=pred_uncond, + guidance_scale=7.0, + momentum_buffer=MomentumBuffer(), + dims=[1], + ) + self.assertEqual(guided.device, device) + + @unittest.skipUnless(hasattr(torch.backends, "mps") and torch.backends.mps.is_available(), "MPS required") + def test_project_returns_mps_tensors_after_cpu_math(self): + device = torch.device("mps") + v0 = torch.randn(2, 4, 8, device=device) + v1 = torch.randn(2, 4, 8, device=device) + parallel, orthogonal = project(v0, v1, dims=[1]) + self.assertEqual(parallel.device.type, "mps") + self.assertEqual(orthogonal.device.type, "mps") + + +if __name__ == "__main__": + unittest.main() diff --git a/acestep/test_device_map.py b/acestep/test_device_map.py new file mode 100644 index 000000000..cfe4a4f05 --- /dev/null +++ b/acestep/test_device_map.py @@ -0,0 +1,264 @@ +"""Unit tests for multi-GPU device map parsing and resolution.""" + +import os +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from acestep.device_map import ( + ComponentDeviceMap, + DeviceMapError, + GpuInfo, + LayoutError, + LayoutRequest, + collect_gpu_runtime_status, + compute_auto_device_map, + cuda_device_index, + device_map_to_dict, + device_type, + estimate_dit_peak_gb, + estimate_lm_total_gb, + format_gpu_list_text, + gpu_info_to_dict, + is_cuda_device, + log_lm_device_deprecation, + normalize_component_device, + parse_gpu_mapping, + resolve_component_device_map, +) + + +class DeviceMapParsingTests(unittest.TestCase): + """Tests for mapping string parsing and validation.""" + + def test_normalize_component_device_expands_bare_cuda(self): + self.assertEqual(normalize_component_device("cuda"), "cuda:0") + + def test_normalize_component_device_preserves_cuda_index(self): + self.assertEqual(normalize_component_device("cuda:2"), "cuda:2") + + def test_parse_gpu_mapping_returns_none_for_empty_and_auto(self): + self.assertIsNone(parse_gpu_mapping(None, default_device="cuda:0")) + self.assertIsNone(parse_gpu_mapping("", default_device="cuda:0")) + self.assertIsNone(parse_gpu_mapping("auto", default_device="cuda:0")) + + def test_parse_gpu_mapping_single_places_all_components_on_index(self): + device_map = parse_gpu_mapping("single:1", default_device="cuda:0") + self.assertIsNotNone(device_map) + assert device_map is not None + self.assertEqual(device_map.dit, "cuda:1") + self.assertEqual(device_map.vae, "cuda:1") + self.assertEqual(device_map.text_encoder, "cuda:1") + self.assertEqual(device_map.lm, "cuda:1") + + def test_parse_gpu_mapping_explicit_components(self): + device_map = parse_gpu_mapping( + "dit:0,vae:0,text_encoder:0,lm:1", + default_device="cuda:0", + ) + self.assertIsNotNone(device_map) + assert device_map is not None + self.assertEqual(device_map.dit, "cuda:0") + self.assertEqual(device_map.lm, "cuda:1") + self.assertTrue(device_map.is_multi_device()) + + def test_parse_gpu_mapping_defaults_aux_components_to_dit(self): + device_map = parse_gpu_mapping("dit:2", default_device="cuda:0") + self.assertIsNotNone(device_map) + assert device_map is not None + self.assertEqual(device_map.dit, "cuda:2") + self.assertEqual(device_map.vae, "cuda:2") + self.assertEqual(device_map.text_encoder, "cuda:2") + self.assertIsNone(device_map.lm) + + def test_parse_gpu_mapping_rejects_unknown_component(self): + with self.assertRaises(DeviceMapError): + parse_gpu_mapping("decoder:0", default_device="cuda:0") + + def test_parse_gpu_mapping_rejects_missing_dit(self): + with self.assertRaises(DeviceMapError): + parse_gpu_mapping("vae:0,lm:1", default_device="cuda:0") + + def test_parse_gpu_mapping_rejects_unsupported_default_device(self): + with self.assertRaises(DeviceMapError): + parse_gpu_mapping("single:0", default_device="unknown-backend") + + def test_parse_gpu_mapping_reads_env_when_argument_missing(self): + with patch.dict(os.environ, {"ACESTEP_GPU_MAPPING": "single:3"}, clear=False): + device_map = parse_gpu_mapping(None, default_device="cuda:0") + self.assertIsNotNone(device_map) + assert device_map is not None + self.assertEqual(device_map.dit, "cuda:3") + + +class DeviceMapResolutionTests(unittest.TestCase): + """Tests for legacy and explicit device map resolution.""" + + def test_resolve_component_device_map_legacy_single_device(self): + device_map = resolve_component_device_map( + requested_device="cuda", + gpu_mapping=None, + ) + self.assertEqual(device_map.dit, "cuda:0") + self.assertFalse(device_map.is_multi_device()) + + def test_resolve_component_device_map_preserves_requested_index(self): + device_map = resolve_component_device_map( + requested_device="cuda:1", + gpu_mapping=None, + ) + self.assertEqual(device_map.dit, "cuda:1") + self.assertEqual(device_map.lm, "cuda:1") + + def test_resolve_component_device_map_auto_uses_multi_gpu_layout(self): + gpus = [ + GpuInfo(0, "GPU0", 24.0, 22.0), + GpuInfo(1, "GPU1", 24.0, 23.0), + ] + with patch("acestep.device_map.resolve.discover_gpus", return_value=gpus): + device_map = resolve_component_device_map( + requested_device="cuda:0", + gpu_mapping="auto", + config_path="acestep-v15-xl-sft", + lm_model_path="acestep-5Hz-lm-4B", + ) + self.assertEqual(device_map.dit, "cuda:1") + self.assertEqual(device_map.lm, "cuda:0") + self.assertTrue(device_map.is_multi_device()) + + def test_resolve_component_device_map_auto_falls_back_on_single_gpu(self): + with patch("acestep.device_map.resolve.discover_gpus", return_value=[GpuInfo(0, "GPU0", 24.0, 22.0)]): + device_map = resolve_component_device_map( + requested_device="cuda:0", + gpu_mapping="auto", + ) + self.assertEqual(device_map.dit, "cuda:0") + self.assertFalse(device_map.is_multi_device()) + + +class AutoLayoutTests(unittest.TestCase): + """Tests for VRAM-aware automatic layout selection.""" + + def test_compute_auto_device_map_splits_dit_and_lm(self): + gpus = [ + GpuInfo(0, "GPU0", 24.0, 10.0), + GpuInfo(1, "GPU1", 24.0, 23.0), + ] + layout = compute_auto_device_map( + LayoutRequest( + gpus=gpus, + dit_type="xl_base", + lm_model_path="acestep-5Hz-lm-4B", + ) + ) + self.assertIsInstance(layout, ComponentDeviceMap) + self.assertEqual(layout.dit, "cuda:1") + self.assertEqual(layout.lm, "cuda:0") + + def test_compute_auto_device_map_returns_error_when_dit_does_not_fit(self): + gpus = [GpuInfo(0, "GPU0", 8.0, 2.0)] + layout = compute_auto_device_map( + LayoutRequest(gpus=gpus, dit_type="xl_base") + ) + self.assertIsInstance(layout, LayoutError) + self.assertIn("No GPU has", layout.message) + + def test_compute_auto_device_map_rejects_same_gpu_lm_without_headroom(self): + """Co-located LM fallback must reserve the DiT footprint on the shared GPU.""" + gpus = [GpuInfo(0, "GPU0", 24.0, 22.0)] + layout = compute_auto_device_map( + LayoutRequest( + gpus=gpus, + dit_type="xl_base", + lm_model_path="acestep-5Hz-lm-4B", + ) + ) + self.assertIsInstance(layout, LayoutError) + self.assertIn("No GPU has", layout.message) + + def test_estimate_helpers_use_config_profiles(self): + self.assertGreater(estimate_dit_peak_gb("xl_base", batch_size=2), estimate_dit_peak_gb("turbo", 1)) + self.assertGreater( + estimate_lm_total_gb("acestep-5Hz-lm-4B"), + estimate_lm_total_gb("acestep-5Hz-lm-0.6B"), + ) + + +class DeviceAliasTests(unittest.TestCase): + """Tests for CUDA alias helpers.""" + + def test_device_for_maps_model_alias_to_dit(self): + device_map = ComponentDeviceMap.from_single_device("cuda:0") + self.assertEqual(device_map.device_for("model"), "cuda:0") + + def test_cuda_device_index_helpers(self): + self.assertTrue(is_cuda_device("cuda:2")) + self.assertEqual(device_type("cuda:2"), "cuda") + self.assertEqual(cuda_device_index("cuda"), 0) + self.assertEqual(cuda_device_index("cuda:2"), 2) + + def test_cuda_device_index_raises_device_map_error_for_malformed_index(self): + with self.assertRaises(DeviceMapError): + cuda_device_index("cuda:x") + + +class GpuStatusHelpersTests(unittest.TestCase): + """Tests for CLI/API GPU listing and status helpers.""" + + def test_format_gpu_list_text_reports_no_devices(self): + self.assertEqual(format_gpu_list_text([]), "No CUDA devices detected.") + + def test_format_gpu_list_text_includes_device_rows(self): + text = format_gpu_list_text( + [GpuInfo(0, "RTX 3090", 24.0, 22.5, (8, 6))] + ) + self.assertIn("RTX 3090", text) + self.assertIn("24.00", text) + + def test_gpu_info_and_device_map_serialization(self): + gpu_payload = gpu_info_to_dict(GpuInfo(1, "GPU1", 24.0, 20.0, (8, 6))) + self.assertEqual(gpu_payload["index"], 1) + self.assertEqual(gpu_payload["compute_capability"], [8, 6]) + + device_map = ComponentDeviceMap( + dit="cuda:0", + vae="cuda:0", + text_encoder="cuda:0", + lm="cuda:1", + ) + map_payload = device_map_to_dict(device_map) + self.assertTrue(map_payload["multi_device"]) + self.assertIn("lm:1", map_payload["summary"]) + + def test_collect_gpu_runtime_status_includes_handler_layout(self): + handler = SimpleNamespace( + device_map=ComponentDeviceMap.from_single_device("cuda:2") + ) + with patch("acestep.device_map.status.discover_gpus", return_value=[]): + status = collect_gpu_runtime_status(handler=handler) + self.assertEqual(status["device_map"]["dit"], "cuda:2") + self.assertEqual(status["gpus"], []) + + def test_log_lm_device_deprecation_without_mapping(self): + with patch.dict(os.environ, {"ACESTEP_LM_DEVICE": "cuda:1"}, clear=False), patch( + "acestep.device_map.status.logger" + ) as mock_logger: + log_lm_device_deprecation() + mock_logger.warning.assert_called_once() + + def test_log_lm_device_deprecation_ignored_when_mapping_assigns_lm(self): + with patch.dict( + os.environ, + { + "ACESTEP_LM_DEVICE": "cuda:0", + "ACESTEP_GPU_MAPPING": "dit:0,vae:0,text_encoder:0,lm:1", + }, + clear=False, + ), patch("acestep.device_map.status.logger") as mock_logger: + log_lm_device_deprecation(using_device_map_lm=True) + mock_logger.warning.assert_called_once() + self.assertIn("ignored", mock_logger.warning.call_args[0][0]) + + +if __name__ == "__main__": + unittest.main() diff --git a/acestep/ui/gradio/events/generation/service_init.py b/acestep/ui/gradio/events/generation/service_init.py index 545aaf1e7..ad0f88598 100644 --- a/acestep/ui/gradio/events/generation/service_init.py +++ b/acestep/ui/gradio/events/generation/service_init.py @@ -86,19 +86,8 @@ def init_service_wrapper( quantization = False quant_value = None - # Compute lm_device only when initializing the LLM to avoid overwriting a - # previously-resolved device (e.g. "cuda") with the raw UI value ("auto"). - # "auto" is resolved to the concrete device inside llm_handler.initialize(). - if init_llm: - if not gpu_config.available_lm_models: - logger.warning( - f"⚠️ GPU tier {gpu_config.tier} ({gpu_config.gpu_memory_gb:.1f}GB) does not support LM on GPU. " - "Falling back to CPU for LM initialization." - ) - lm_device = "cpu" - else: - lm_device = device - + # Tier / backend checks can run before DiT init; LM device must be resolved + # after initialize_service so device_map.lm from ACESTEP_GPU_MAPPING is set. if init_llm and lm_model_path and gpu_config.available_lm_models: if not is_lm_model_size_allowed(lm_model_path, gpu_config.available_lm_models): logger.warning( @@ -130,8 +119,24 @@ def init_service_wrapper( offload_to_cpu=offload_to_cpu, offload_dit_to_cpu=offload_dit_to_cpu, quantization=quant_value, use_mlx_dit=mlx_dit, vae_checkpoint=vae_checkpoint, + gpu_mapping=os.environ.get("ACESTEP_GPU_MAPPING"), ) + lm_device = device + if init_llm: + if not gpu_config.available_lm_models: + logger.warning( + f"⚠️ GPU tier {gpu_config.tier} ({gpu_config.gpu_memory_gb:.1f}GB) does not support LM on GPU. " + "Falling back to CPU for LM initialization." + ) + lm_device = "cpu" + else: + device_map = getattr(dit_handler, "device_map", None) + if device_map is not None and device_map.lm is not None: + lm_device = device_map.lm + else: + lm_device = device + if init_llm: checkpoint_dir = os.path.join(project_root, "checkpoints") diff --git a/acestep/ui/gradio/events/generation/service_init_test.py b/acestep/ui/gradio/events/generation/service_init_test.py index 427de34b0..33c7d5a3a 100644 --- a/acestep/ui/gradio/events/generation/service_init_test.py +++ b/acestep/ui/gradio/events/generation/service_init_test.py @@ -222,6 +222,7 @@ def test_init_llm_with_auto_device_calls_initialize(self, mock_gpu_config): dit_handler.initialize_service.return_value = ("ok", True) dit_handler.model = MagicMock() dit_handler.is_turbo_model.return_value = True + dit_handler.device_map = None llm_handler = MagicMock() llm_handler.llm_initialized = False @@ -251,6 +252,62 @@ def test_init_llm_with_auto_device_calls_initialize(self, mock_gpu_config): "initialize() must receive 'auto' so it can resolve to the best device", ) + @patch("acestep.ui.gradio.events.generation.service_init.get_global_gpu_config") + def test_init_llm_uses_device_map_lm_after_dit_init(self, mock_gpu_config): + """LM device must come from device_map after DiT initialize_service. + + Regression: resolving lm_device before initialize_service left device_map + empty on first Gradio init, so UI device='auto' collapsed the LM onto + bare cuda:0 even when ACESTEP_GPU_MAPPING set lm:1. + """ + module = self._import_module() + + mock_gpu_config.return_value = MagicMock( + available_lm_models=["acestep-5Hz-lm-1.7B"], + lm_backend_restriction=None, + tier="tier6", + gpu_memory_gb=24.0, + max_duration_with_lm=600, + max_duration_without_lm=600, + max_batch_size_with_lm=4, + max_batch_size_without_lm=8, + ) + + dit_handler = MagicMock() + dit_handler.device_map = None + + def _init_service(*_args, **_kwargs): + dit_handler.device_map = MagicMock(lm="cuda:1") + return ("ok", True) + + dit_handler.initialize_service.side_effect = _init_service + dit_handler.model = MagicMock() + dit_handler.is_turbo_model.return_value = True + + llm_handler = MagicMock() + llm_handler.llm_initialized = False + llm_handler.initialize.return_value = ("[OK] LLM initialized", True) + + module.init_service_wrapper( + dit_handler, + llm_handler, + "/some/project/checkpoints", + "acestep-v15-turbo", + "auto", + True, + "acestep-5Hz-lm-1.7B", + "pt", + use_flash_attention=False, + offload_to_cpu=False, + offload_dit_to_cpu=False, + compile_model=False, + quantization=False, + ) + + llm_handler.initialize.assert_called_once() + _, call_kwargs = llm_handler.initialize.call_args + self.assertEqual(call_kwargs.get("device"), "cuda:1") + @patch("acestep.ui.gradio.events.generation.service_init.get_global_gpu_config") def test_legacy_cuda_config_forces_pt_backend(self, mock_gpu_config): """Legacy CUDA restrictions should override a requested vllm backend."""