diff --git a/recipes/LTX/LTX-2.md b/recipes/LTX/LTX-2.md index ca3201dc1..fa9ec6755 100644 --- a/recipes/LTX/LTX-2.md +++ b/recipes/LTX/LTX-2.md @@ -165,6 +165,40 @@ Pass these fields in online `extra_params` or offline `--extra-body`: --extra-body '{"video_cfg_scale":3.0,"audio_cfg_scale":7.0}' ``` +### Guidance Parallelism + +For LTX, `cfg_parallel_size` is the number of ranks used to execute the +complete guidance plan in parallel. Despite the legacy `cfg` name, it covers +text CFG, STG, and cross-modality guidance passes; guidance rescaling is +applied after all pass predictions are gathered. + +The default one-stage recipe has four Transformer passes per denoise step: +`cond`, `uncond`, `ptb` (STG), and `mod` (cross-modality). The useful balanced +configurations are therefore: + +| `--cfg-parallel-size` | Passes per rank | Guidance-slot utilization | Notes | +|---:|---:|---:|---| +| `1` | 4 | 100% | Single-rank fused guidance batch | +| `2` | 2 | 100% | Recommended two-rank configuration | +| `4` | 1 | 100% | One guidance pass per rank | + +Other positive sizes are accepted. When the pass count is not divisible by +`cfg_parallel_size`, ranks are padded to an equal number of execution slots +and LTX emits a warning with the expected utilization. For example, four +passes on three ranks use six slots and have 66.7% guidance-slot utilization. + +Start an LTX-2.3 server with two-way guidance parallelism: + +```bash +vllm serve diffusers/LTX-2.3-Diffusers --omni \ + --cfg-parallel-size 2 --stage-init-timeout 600 +``` + +The total device count is the product of `cfg_parallel_size` and the other +configured parallel dimensions. Distilled pipelines use positive-only +guidance, so `cfg_parallel_size > 1` does not provide guidance speedup for +them. + ### Python API Normal per-request values belong in `OmniDiffusionSamplingParams`. Pass the @@ -273,8 +307,8 @@ bundled offline CLI do not currently expose `sigmas`. - Fused one-stage requests must resolve to identical LTX guidance and sigma schedules. Keep concurrent requests guidance-homogeneous or use `--max-num-seqs 1`. -- `--cfg-parallel-size 2` supports only the CFG-only plan: disable STG and - modality guidance and set both rescale fields to `0.0`. +- `--cfg-parallel-size` shards the complete LTX guidance plan, including STG, + modality guidance, and rescale-compatible prediction gathering. ## Operational Notes diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index febaf4b0b..d477e05c1 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -577,9 +577,11 @@ def test_from_pipeline_config_derives_sequence_parallel_size_from_allgather_degr assert stage.parallel_config.world_size == 2 -def test_diffusion_parallel_config_rejects_cfg_parallel_size_outside_current_bound(): - with pytest.raises(ValidationError): - OmniStageDiffusionParallelConfig(cfg_parallel_size=4) +def test_diffusion_parallel_config_accepts_four_way_guidance_parallelism(): + cfg = OmniStageDiffusionParallelConfig(cfg_parallel_size=4) + + assert cfg.cfg_parallel_size == 4 + assert cfg.world_size == 4 def test_diffusion_parallel_config_rejects_allgather_with_ulysses_or_ring(): diff --git a/tests/diffusion/models/ltx2/test_ltx2_conditioning.py b/tests/diffusion/models/ltx2/test_ltx2_conditioning.py index 6b4ec70ce..e309dbf47 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_conditioning.py +++ b/tests/diffusion/models/ltx2/test_ltx2_conditioning.py @@ -110,7 +110,7 @@ class TestLTXImageToVideoForwardStages: kwargs = pipe._denoise_timestep_kwargs( ts, - SimpleNamespace(cfg_parallel_ready=False), + SimpleNamespace(guidance_parallel_ready=False), denoise_ctx, video_token_count=2, audio_token_count=1, diff --git a/tests/diffusion/models/ltx2/test_ltx2_parallel.py b/tests/diffusion/models/ltx2/test_ltx2_parallel.py index 92077a229..8d8c22e79 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_parallel.py +++ b/tests/diffusion/models/ltx2/test_ltx2_parallel.py @@ -3,6 +3,7 @@ """Unit tests for LTX guidance and forward-parallel behavior.""" +from contextlib import nullcontext from types import SimpleNamespace import pytest @@ -50,7 +51,8 @@ class TestCFGParallelHelpers: assert torch.allclose(video_combined, expected_video) assert torch.allclose(audio_combined, expected_audio) - def test_cfg_parallel_rejects_non_cfg_guidance(self): + @pytest.mark.parametrize("guidance_world_size", [1, 2, 3, 4]) + def test_guidance_parallel_accepts_official_guidance(self, guidance_world_size): from vllm_omni.diffusion.models.ltx2.ltx2_guidance import ( LTX_GUIDANCE_EXECUTOR, LTXGuidancePlan, @@ -59,26 +61,180 @@ class TestCFGParallelHelpers: plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance) - with pytest.raises(ValueError, match="CFG-only guidance"): - LTX_GUIDANCE_EXECUTOR.validate_cfg_world_size(plan, 2) + LTX_GUIDANCE_EXECUTOR.validate_guidance_world_size(plan, guidance_world_size) - def test_cfg_parallel_dummy_warms_supported_cfg_only_plan(self, monkeypatch): + def test_guidance_parallel_dummy_warms_full_guidance_plan(self, monkeypatch): from vllm_omni.diffusion.models.ltx2 import ltx2_runtime from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline - monkeypatch.setattr(ltx2_runtime, "get_classifier_free_guidance_world_size", lambda: 2) + monkeypatch.setattr(ltx2_runtime, "get_guidance_parallel_world_size", lambda: 2) pipe = object.__new__(LTX2Pipeline) req = SimpleNamespace(is_dummy_run=lambda: True) request_inputs = SimpleNamespace(guidance=LTX23_ONE_STAGE_RECIPE.request_guidance) - cfg_parallel_ready = pipe._setup_forward_runtime(req, request_inputs, attention_kwargs=None) + guidance_parallel_ready = pipe._setup_forward_runtime(req, request_inputs, attention_kwargs=None) - assert cfg_parallel_ready is True - assert pipe._guidance_plan.names == ("cond", "uncond") + assert guidance_parallel_ready is True + assert pipe._guidance_plan.names == ("cond", "uncond", "ptb", "mod") assert pipe._guidance_plan.spec.video.cfg_scale == 3.0 assert pipe._guidance_plan.spec.audio.cfg_scale == 7.0 + def test_guidance_parallel_warns_for_imbalanced_passes(self, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import ltx2_guidance + + spec = ltx2_guidance.LTXGuidanceSpec( + video=ltx2_guidance.LTXModalityGuidance( + cfg_scale=3.0, + stg_scale=1.0, + stg_blocks=(1,), + ) + ) + plan = ltx2_guidance.LTXGuidancePlan.build(spec) + captured = [] + monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_rank", lambda: 0) + monkeypatch.setattr(ltx2_guidance.logger, "warning_once", lambda *args: captured.append(args)) + + ltx2_guidance.LTX_GUIDANCE_EXECUTOR.warn_if_imbalanced(plan, 4, "hq_generate") + + assert plan.names == ("cond", "uncond", "ptb") + assert len(captured) == 1 + assert captured[0][1:] == ("hq_generate", 3, 4, 1, 75.0) + + @pytest.mark.parametrize("guidance_rank", [0, 1]) + def test_default_guidance_is_sharded_and_combined(self, monkeypatch, guidance_rank): + from vllm_omni.diffusion.models.ltx2 import ltx2_guidance + from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseContext + from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState + from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE + + executor = ltx2_guidance.LTX_GUIDANCE_EXECUTOR + plan = ltx2_guidance.LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance) + video_predictions = { + "cond": torch.tensor([[[0.2, -0.3]]]), + "uncond": torch.tensor([[[-0.4, 0.1]]]), + "ptb": torch.tensor([[[0.5, -0.2]]]), + "mod": torch.tensor([[[-0.1, 0.6]]]), + } + audio_predictions = { + "cond": torch.tensor([[[0.7, -0.2]]]), + "uncond": torch.tensor([[[0.1, 0.4]]]), + "ptb": torch.tensor([[[-0.3, 0.2]]]), + "mod": torch.tensor([[[0.4, -0.5]]]), + } + state = LTXAVState( + video=torch.tensor([[[1.0, -2.0]]]), + audio=torch.tensor([[[0.5, 3.0]]]), + ) + denoise_ctx = LTXDenoiseContext( + latents=state.video, + audio_latents=state.audio, + video_coords=torch.zeros(1, 1, 3), + audio_coords=torch.zeros(1, 1, 1), + ) + executor.prepare_denoise_context(plan, True, 2, denoise_ctx) + + class FakeGuidanceGroup: + def __init__(self): + self.call_index = 0 + + def all_gather(self, tensor, separate_tensors=True): + assert separate_tensors + gathered = ( + [video_predictions["cond"], video_predictions["uncond"]], + [video_predictions["ptb"], video_predictions["mod"]], + [audio_predictions["cond"], audio_predictions["uncond"]], + [audio_predictions["ptb"], audio_predictions["mod"]], + )[self.call_index] + torch.testing.assert_close(tensor, gathered[guidance_rank]) + self.call_index += 1 + return gathered + + group = FakeGuidanceGroup() + monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_world_size", lambda: 2) + monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_rank", lambda: guidance_rank) + monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_group", lambda: group) + + local_names = ("cond", "ptb") if guidance_rank == 0 else ("uncond", "mod") + captured = {} + + class FakeTransformer: + def __call__(self, **kwargs): + captured.update(kwargs) + return ( + torch.cat([video_predictions[name] for name in local_names]), + torch.cat([audio_predictions[name] for name in local_names]), + ) + + video_sigma = torch.tensor(0.25) + audio_sigma = torch.tensor(0.5) + + class FakePipeline: + scheduler = SimpleNamespace(sigmas=torch.stack([video_sigma])) + transformer = FakeTransformer() + + @staticmethod + def _transformer_cache_context(name): + assert name == "guided" + return nullcontext() + + @staticmethod + def _build_transformer_kwargs(forward_ctx, denoise_ctx, **kwargs): + del forward_ctx, denoise_ctx + return kwargs + + @staticmethod + def _video_guidance_model_sigma(sigma, denoise_ctx): + del denoise_ctx + return sigma + + positive_context = torch.ones(1, 1, 1) + negative_context = torch.zeros(1, 1, 1) + forward_ctx = SimpleNamespace( + prompt_context=SimpleNamespace( + positive_connector_prompt_embeds=positive_context, + positive_connector_audio_prompt_embeds=positive_context, + negative_connector_prompt_embeds=negative_context, + negative_connector_audio_prompt_embeds=negative_context, + ), + attention_kwargs=None, + audio_scheduler=SimpleNamespace(sigmas=torch.stack([audio_sigma])), + ) + + actual_video, actual_audio = executor.predict_parallel_guidance( + FakePipeline(), + plan, + index=0, + timestep=torch.tensor(1.0), + state=state, + forward_ctx=forward_ctx, + denoise_ctx=denoise_ctx, + ) + + expected_video = executor._guide_modality( + state.video, + video_predictions, + video_sigma, + plan.spec.video, + ) + expected_audio = executor._guide_modality( + state.audio, + audio_predictions, + audio_sigma, + plan.spec.audio, + ) + torch.testing.assert_close(actual_video, expected_video) + torch.testing.assert_close(actual_audio, expected_audio) + assert group.call_index == 4 + assert captured["hidden_states"].shape[0] == 2 + perturbations = captured["attention_kwargs"]["ltx_perturbation_kwargs"] + if guidance_rank == 0: + assert "video_self_attention_mask" in perturbations + assert "a2v_cross_attention_mask" not in perturbations + else: + assert "video_self_attention_mask" not in perturbations + assert "a2v_cross_attention_mask" in perturbations + class TestCFGParallelForwardPath: """Test the LTX-2.3 CFG-parallel denoising path without loading model weights.""" @@ -96,7 +252,7 @@ class TestCFGParallelForwardPath: pipe.tokenizer_max_length = 4 pipe.vae_spatial_compression_ratio = 32 pipe.vae_temporal_compression_ratio = 8 - monkeypatch.setattr(ltx2_runtime, "get_classifier_free_guidance_world_size", lambda: 1) + monkeypatch.setattr(ltx2_runtime, "get_guidance_parallel_world_size", lambda: 1) class StopAtEncodePromptError(Exception): pass @@ -238,10 +394,10 @@ class TestCFGParallelForwardPath: return [audio_pos, audio_neg] raise AssertionError(f"Unexpected gathered tensor: {tensor}") - monkeypatch.setattr(ltx2_runtime, "get_classifier_free_guidance_world_size", lambda: 2) - monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_world_size", lambda: 2) - monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_rank", lambda: cfg_rank) - monkeypatch.setattr(ltx2_guidance, "get_cfg_group", lambda: FakeCfgGroup()) + monkeypatch.setattr(ltx2_runtime, "get_guidance_parallel_world_size", lambda: 2) + monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_world_size", lambda: 2) + monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_rank", lambda: cfg_rank) + monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_group", lambda: FakeCfgGroup()) def fake_retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas=None, mu=None): scheduler.sigmas = torch.tensor([0.25, 0.25], device=device) diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index cbe69302f..2606723f8 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -46,6 +46,16 @@ def test_serve_parser_accepts_strategy_config() -> None: assert args.get_explicit_kwargs_dict()["strategy_config"] == "/tmp/strategy.yaml" +def test_serve_parser_accepts_four_way_cfg_parallelism() -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args(["serve", "fake-model", "--omni", "--cfg-parallel-size", "4"]) + + assert args.cfg_parallel_size == 4 + + def _make_headless_args(**kwargs) -> TrackingNamespace: defaults = { "model": "fake-model", diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index dad16d184..422e8d784 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -368,7 +368,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): ring_degree: int = Field(default=1, ge=1) allgather_degree: int = Field(default=1, ge=1) ulysses_mode: str = "strict" - cfg_parallel_size: int = Field(default=1, ge=1, le=3) + cfg_parallel_size: int = Field(default=1, ge=1) vae_patch_parallel_size: int = Field(default=1, ge=1) vae_parallel_mode: str = "tile" use_hsdp: bool = False diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index ca9d48d96..b97ffc311 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -185,7 +185,7 @@ class DiffusionParallelConfig: """ cfg_parallel_size: int = 1 - """Number of Classifier Free Guidance (CFG) parallel groups.""" + """Number of ranks used to execute guidance passes in parallel.""" vae_patch_parallel_size: int = 1 """Number of ranks used for VAE patch/tile parallelism (decode/encode).""" @@ -233,9 +233,6 @@ class DiffusionParallelConfig: assert self.ring_degree > 0, "Ring degree must be > 0" assert self.allgather_degree > 0, "AllGather degree must be > 0" assert self.cfg_parallel_size > 0, "CFG parallel size must be > 0" - assert self.cfg_parallel_size in [1, 2, 3], ( - f"CFG parallel size must be 1, 2, or 3, but got {self.cfg_parallel_size}" - ) assert self.vae_patch_parallel_size > 0, "VAE patch parallel size must be > 0" assert self.vae_parallel_mode in {"tile", "spatial_shard_height", "spatial_shard_width"}, ( "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, " diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py b/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py index 6bd759057..f4d2f26f6 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py @@ -705,12 +705,12 @@ class LTXI2VConditioningMixin: latents=request_inputs.latents, ) - def _prepare_denoise_context_for_cfg( + def _prepare_denoise_context_for_guidance( self, forward_ctx: LTXForwardContext, denoise_ctx: LTXDenoiseContext, ) -> LTXDenoiseContext: - denoise_ctx = super()._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx) + denoise_ctx = super()._prepare_denoise_context_for_guidance(forward_ctx, denoise_ctx) if denoise_ctx.conditioning_mask is None: if self.unified_text_image_entry: return denoise_ctx @@ -755,9 +755,7 @@ class LTXI2VConditioningMixin: video_token_count=video_token_count, audio_token_count=audio_token_count, ) - conditioning_mask = ( - denoise_ctx.conditioning_mask if forward_ctx.cfg_parallel_ready else denoise_ctx.conditioning_mask_for_model - ) + conditioning_mask = denoise_ctx.conditioning_mask_for_model if conditioning_mask is None: raise ValueError("LTX I2V denoising requires a conditioning mask.") kwargs.update( diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py index 262c9334c..4a149be85 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py @@ -33,7 +33,7 @@ class LTXForwardContext: request_inputs: LTXRequestInputs prompt_context: LTXPromptContext device: torch.device - cfg_parallel_ready: bool + guidance_parallel_ready: bool attention_kwargs: dict[str, Any] | None latent_num_frames: int latent_height: int @@ -358,9 +358,9 @@ def step_denoised_latents( do_true_cfg=pipeline.do_classifier_free_guidance, per_request_scheduler=forward_ctx.video_audio_step_adapter, ) - return pipeline._synchronize_cfg_parallel_step_output( + return pipeline._synchronize_guidance_parallel_step_output( latents, - do_true_cfg=pipeline.do_classifier_free_guidance, + guidance_parallel_ready=forward_ctx.guidance_parallel_ready, ) @@ -382,7 +382,7 @@ class LTXPhaseExecutor: prompt_context: LTXPromptContext | None = None, ) -> LTXPhaseResult: pipeline._check_forward_inputs(request_inputs, image=image) - cfg_parallel_ready = pipeline._setup_forward_runtime(req, request_inputs, attention_kwargs) + guidance_parallel_ready = pipeline._setup_forward_runtime(req, request_inputs, attention_kwargs) device = pipeline.device if prompt_context is None: prompt_context = pipeline._prepare_prompt_context( @@ -429,7 +429,7 @@ class LTXPhaseExecutor: request_inputs=request_inputs, prompt_context=prompt_context, device=device, - cfg_parallel_ready=cfg_parallel_ready, + guidance_parallel_ready=guidance_parallel_ready, attention_kwargs=attention_kwargs, latent_num_frames=latent_num_frames, latent_height=latent_height, @@ -449,7 +449,7 @@ class LTXPhaseExecutor: audio_coords=audio_coords, conditioning_mask=conditioning_mask, ) - denoise_ctx = pipeline._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx) + denoise_ctx = pipeline._prepare_denoise_context_for_guidance(forward_ctx, denoise_ctx) state = LTXDenoiseExecutor.run( pipeline, LTXAVState(video=denoise_ctx.latents, audio=denoise_ctx.audio_latents), diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py b/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py index 0b3d2dc2b..f35f5b998 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py @@ -10,12 +10,19 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any import torch +from vllm.logger import init_logger from vllm_omni.diffusion.distributed.parallel_state import ( - get_cfg_group, - get_classifier_free_guidance_rank, - get_classifier_free_guidance_world_size, + get_cfg_group as get_guidance_parallel_group, ) +from vllm_omni.diffusion.distributed.parallel_state import ( + get_classifier_free_guidance_rank as get_guidance_parallel_rank, +) +from vllm_omni.diffusion.distributed.parallel_state import ( + get_classifier_free_guidance_world_size as get_guidance_parallel_world_size, +) + +logger = init_logger(__name__) if TYPE_CHECKING: from .ltx2_denoise import LTXDenoiseContext, LTXForwardContext @@ -110,10 +117,6 @@ class LTXGuidancePlan: def names(self) -> tuple[str, ...]: return tuple(item.name for item in self.passes) - @property - def cfg_parallel_compatible(self) -> bool: - return self.names == ("cond", "uncond") and not self.spec.do_rescale - def x0_from_velocity(sample: torch.Tensor, velocity: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor: """Convert velocity to x0 using the official fp32 arithmetic order.""" @@ -212,23 +215,69 @@ class LTXGuidanceExecutor: """Build and execute official LTX multi-guidance Transformer passes.""" @staticmethod - def validate_cfg_world_size(plan: LTXGuidancePlan, cfg_world_size: int) -> None: - if cfg_world_size == 1: + def validate_guidance_world_size(plan: LTXGuidancePlan, guidance_world_size: int) -> None: + del plan + if guidance_world_size < 1: + raise ValueError(f"LTX guidance parallelism requires a positive parallel size, got {guidance_world_size}.") + + @staticmethod + def _parallel_assignments(pass_count: int, world_size: int) -> list[list[int]]: + assignments: list[list[int]] = [[] for _ in range(world_size)] + for pass_index in range(pass_count): + assignments[pass_index % world_size].append(pass_index) + return assignments + + @classmethod + def _model_pass_count( + cls, + plan: LTXGuidancePlan, + guidance_parallel_ready: bool, + guidance_world_size: int, + ) -> int: + if not guidance_parallel_ready: + return len(plan.passes) + assignments = cls._parallel_assignments(len(plan.passes), guidance_world_size) + return max(len(indices) for indices in assignments) + + @classmethod + def warn_if_imbalanced( + cls, + plan: LTXGuidancePlan, + guidance_world_size: int, + phase_name: str, + ) -> None: + pass_count = len(plan.passes) + if guidance_world_size <= 1 or pass_count % guidance_world_size == 0 or get_guidance_parallel_rank() != 0: return - if cfg_world_size != 2: - raise ValueError(f"LTX CFG parallelism supports cfg_parallel_size 1 or 2, got {cfg_world_size}.") - if not plan.cfg_parallel_compatible: - raise ValueError("LTX CFG parallelism only supports CFG-only guidance without rescale.") + slots_per_rank = math.ceil(pass_count / guidance_world_size) + total_slots = slots_per_rank * guidance_world_size + wasted_slots = total_slots - pass_count + utilization = 100.0 * pass_count / total_slots + logger.warning_once( + "LTX guidance parallelism is imbalanced for phase %r: %d guidance passes across %d ranks. " + "%d padded/discarded rank slot(s); expected guidance-slot utilization is %.1f%%.", + phase_name, + pass_count, + guidance_world_size, + wasted_slots, + utilization, + ) @staticmethod def prepare_denoise_context( plan: LTXGuidancePlan, - cfg_parallel_ready: bool, + guidance_parallel_ready: bool, + guidance_world_size: int, denoise_ctx: LTXDenoiseContext, ) -> LTXDenoiseContext: - if len(plan.passes) > 1 and not cfg_parallel_ready: - denoise_ctx.video_coords = _repeat_batch(denoise_ctx.video_coords, len(plan.passes)) - denoise_ctx.audio_coords = _repeat_batch(denoise_ctx.audio_coords, len(plan.passes)) + model_pass_count = LTXGuidanceExecutor._model_pass_count( + plan, + guidance_parallel_ready, + guidance_world_size, + ) + if model_pass_count > 1: + denoise_ctx.video_coords = _repeat_batch(denoise_ctx.video_coords, model_pass_count) + denoise_ctx.audio_coords = _repeat_batch(denoise_ctx.audio_coords, model_pass_count) return denoise_ctx @staticmethod @@ -265,7 +314,30 @@ class LTXGuidanceExecutor: ) return velocity_from_x0(sample, guided_x0, sigma) - def predict_parallel_cfg( + @staticmethod + def _guide_modality( + sample: torch.Tensor, + splits: dict[str, torch.Tensor], + sigma: torch.Tensor, + guidance: LTXModalityGuidance, + *, + model_sigma: torch.Tensor | None = None, + ) -> torch.Tensor: + model_sigma = sigma if model_sigma is None else model_sigma + # Official guidance reduces contiguous BSC tensors. The fused + # transformer output may be channel-major after splitting, which + # changes fp32 reduction order and can cross bf16 rounding bounds. + x0 = {name: x0_from_velocity(sample, value, model_sigma).contiguous() for name, value in splits.items()} + guided = combine_guided_x0( + cond=x0["cond"], + uncond_text=x0.get("uncond", 0.0), + uncond_perturbed=x0.get("ptb", 0.0), + uncond_modality=x0.get("mod", 0.0), + guidance=guidance, + ) + return velocity_from_x0(sample, guided, sigma) + + def predict_parallel_guidance( self, pipeline: Any, plan: LTXGuidancePlan, @@ -275,54 +347,93 @@ class LTXGuidanceExecutor: forward_ctx: LTXForwardContext, denoise_ctx: LTXDenoiseContext, ) -> tuple[torch.Tensor, torch.Tensor]: - self.validate_cfg_world_size(plan, get_classifier_free_guidance_world_size()) + guidance_world_size = get_guidance_parallel_world_size() + self.validate_guidance_world_size(plan, guidance_world_size) + guidance_rank = get_guidance_parallel_rank() + assignments = self._parallel_assignments(len(plan.passes), guidance_world_size) + local_pass_indices = assignments[guidance_rank] + model_pass_count = max(len(indices) for indices in assignments) + + # Every rank executes and gathers the same number of slots. Shorter + # assignments use a conditional pass whose output is discarded. + padded_pass_indices: list[int | None] = local_pass_indices + [None] * ( + model_pass_count - len(local_pass_indices) + ) + local_passes = tuple( + plan.passes[0] if pass_index is None else plan.passes[pass_index] for pass_index in padded_pass_indices + ) + local_plan = LTXGuidancePlan(spec=plan.spec, passes=local_passes) + prompt = forward_ctx.prompt_context - negative = get_classifier_free_guidance_rank() == 1 - video_context = ( - prompt.negative_connector_prompt_embeds - if negative and plan.spec.video.do_cfg - else prompt.positive_connector_prompt_embeds + video_contexts: list[torch.Tensor] = [] + audio_contexts: list[torch.Tensor] = [] + for denoise_pass in local_passes: + video_context = ( + prompt.negative_connector_prompt_embeds + if denoise_pass.negative_video_context + else prompt.positive_connector_prompt_embeds + ) + audio_context = ( + prompt.negative_connector_audio_prompt_embeds + if denoise_pass.negative_audio_context + else prompt.positive_connector_audio_prompt_embeds + ) + if video_context is None or audio_context is None: + raise ValueError("Negative prompt context is required when LTX CFG is enabled.") + video_contexts.append(video_context) + audio_contexts.append(audio_context) + + video_input = _repeat_batch(state.video, model_pass_count).to(prompt.positive_connector_prompt_embeds.dtype) + audio_input = _repeat_batch(state.audio, model_pass_count).to( + prompt.positive_connector_audio_prompt_embeds.dtype ) - audio_context = ( - prompt.negative_connector_audio_prompt_embeds - if negative and plan.spec.audio.do_cfg - else prompt.positive_connector_audio_prompt_embeds - ) - if video_context is None or audio_context is None: - raise ValueError("Negative prompt context is required when LTX CFG is enabled.") - video_input = state.video.to(prompt.positive_connector_prompt_embeds.dtype) - audio_input = state.audio.to(prompt.positive_connector_audio_prompt_embeds.dtype) + attention_kwargs = dict(forward_ctx.attention_kwargs or {}) + perturbations = build_perturbation_kwargs(local_plan, state.video.shape[0], video_input) + if perturbations: + attention_kwargs["ltx_perturbation_kwargs"] = perturbations kwargs = pipeline._build_transformer_kwargs( forward_ctx, denoise_ctx, hidden_states=video_input, audio_hidden_states=audio_input, - encoder_hidden_states=video_context, - audio_encoder_hidden_states=audio_context, + encoder_hidden_states=torch.cat(video_contexts), + audio_encoder_hidden_states=torch.cat(audio_contexts), encoder_attention_mask=None, audio_encoder_attention_mask=None, ts=timestep.expand(video_input.shape[0]), + attention_kwargs=attention_kwargs, ) - local_video, local_audio = pipeline.predict_noise(**kwargs) - group = get_cfg_group() - video = group.all_gather(local_video, separate_tensors=True) - audio = group.all_gather(local_audio, separate_tensors=True) + with pipeline._transformer_cache_context("guided"): + local_video, local_audio = pipeline.transformer(**kwargs) + local_video_slots = local_video.chunk(model_pass_count) + local_audio_slots = local_audio.chunk(model_pass_count) + + group = get_guidance_parallel_group() + gathered_video_slots = [group.all_gather(value, separate_tensors=True) for value in local_video_slots] + gathered_audio_slots = [group.all_gather(value, separate_tensors=True) for value in local_audio_slots] + + video_splits: dict[str, torch.Tensor] = {} + audio_splits: dict[str, torch.Tensor] = {} + for pass_index, denoise_pass in enumerate(plan.passes): + owner_rank = pass_index % guidance_world_size + slot_index = pass_index // guidance_world_size + video_splits[denoise_pass.name] = gathered_video_slots[slot_index][owner_rank] + audio_splits[denoise_pass.name] = gathered_audio_slots[slot_index][owner_rank] + video_sigma = pipeline.scheduler.sigmas[index] return ( - self.combine_cfg_velocity( + self._guide_modality( state.video, - video[0], - video[1], + video_splits, video_sigma, - plan.spec.video.cfg_scale, + plan.spec.video, model_sigma=pipeline._video_guidance_model_sigma(video_sigma, denoise_ctx), ), - self.combine_cfg_velocity( + self._guide_modality( state.audio, - audio[0], - audio[1], + audio_splits, forward_ctx.audio_scheduler.sigmas[index], - plan.spec.audio.cfg_scale, + plan.spec.audio, ), ) @@ -336,8 +447,8 @@ class LTXGuidanceExecutor: forward_ctx: LTXForwardContext, denoise_ctx: LTXDenoiseContext, ) -> tuple[torch.Tensor, torch.Tensor]: - if forward_ctx.cfg_parallel_ready: - return self.predict_parallel_cfg(pipeline, plan, index, timestep, state, forward_ctx, denoise_ctx) + if forward_ctx.guidance_parallel_ready: + return self.predict_parallel_guidance(pipeline, plan, index, timestep, state, forward_ctx, denoise_ctx) prompt = forward_ctx.prompt_context video_contexts: list[torch.Tensor] = [] @@ -382,30 +493,8 @@ class LTXGuidanceExecutor: video_splits = dict(zip(plan.names, video_velocity.chunk(pass_count), strict=True)) audio_splits = dict(zip(plan.names, audio_velocity.chunk(pass_count), strict=True)) - def guide_modality( - sample: torch.Tensor, - splits: dict[str, torch.Tensor], - sigma: torch.Tensor, - guidance: LTXModalityGuidance, - *, - model_sigma: torch.Tensor | None = None, - ) -> torch.Tensor: - model_sigma = sigma if model_sigma is None else model_sigma - # Official guidance reduces contiguous BSC tensors. The fused - # transformer output may be channel-major after splitting, which - # changes fp32 reduction order and can cross bf16 rounding bounds. - x0 = {name: x0_from_velocity(sample, value, model_sigma).contiguous() for name, value in splits.items()} - guided = combine_guided_x0( - cond=x0["cond"], - uncond_text=x0.get("uncond", 0.0), - uncond_perturbed=x0.get("ptb", 0.0), - uncond_modality=x0.get("mod", 0.0), - guidance=guidance, - ) - return velocity_from_x0(sample, guided, sigma) - return ( - guide_modality( + self._guide_modality( state.video, video_splits, pipeline.scheduler.sigmas[index], @@ -415,7 +504,7 @@ class LTXGuidanceExecutor: denoise_ctx, ), ), - guide_modality( + self._guide_modality( state.audio, audio_splits, forward_ctx.audio_scheduler.sigmas[index], diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py index b5bb18a9b..4b74ebcb8 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py @@ -17,7 +17,9 @@ from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin -from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size +from vllm_omni.diffusion.distributed.parallel_state import ( + get_classifier_free_guidance_world_size as get_guidance_parallel_world_size, +) from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin @@ -44,8 +46,6 @@ from .ltx2_guidance import ( LTX_GUIDANCE_EXECUTOR, LTXGuidanceExecutor, LTXGuidancePlan, - LTXGuidanceSpec, - LTXModalityGuidance, ) from .ltx2_recipes import ( LTXPhaseRecipe, @@ -259,7 +259,7 @@ class LTXRuntime( def _enter_phase(self, phase: LTXPhaseRecipe) -> None: """Hook for a future phase-weight strategy.""" - del phase + self._active_phase_name = phase.name def prepare_latents( self, @@ -323,6 +323,10 @@ class LTXRuntime( def do_classifier_free_guidance(self): return self._guidance_plan.spec.do_cfg + @property + def do_guidance(self): + return len(self._guidance_plan.passes) > 1 + @property def interrupt(self): return self._interrupt @@ -372,16 +376,16 @@ class LTXRuntime( ), ) - def _synchronize_cfg_parallel_step_output( + def _synchronize_guidance_parallel_step_output( self, latents: tuple[torch.Tensor, torch.Tensor], - do_true_cfg: bool, + guidance_parallel_ready: bool, ) -> tuple[torch.Tensor, torch.Tensor]: - if not (do_true_cfg and get_classifier_free_guidance_world_size() > 1): + if not (guidance_parallel_ready and get_guidance_parallel_world_size() > 1): return latents # CUDA async execution otherwise permits numerical drift to accumulate - # across CFG-parallel denoise steps. + # across guidance-parallel denoise steps. latents = tuple(tensor.contiguous() for tensor in latents) current_omni_platform.synchronize() return latents @@ -393,26 +397,16 @@ class LTXRuntime( attention_kwargs: dict[str, Any] | None, ) -> bool: self._guidance_plan = LTXGuidancePlan.build(request_inputs.guidance) - del attention_kwargs + del req, attention_kwargs self._interrupt = False - cfg_world_size = get_classifier_free_guidance_world_size() - if ( - req.is_dummy_run() - and cfg_world_size == 2 - and self.do_classifier_free_guidance - and not self._guidance_plan.cfg_parallel_compatible - ): - # The single-device warmup follows the full recipe. CFG2 executes - # the same per-modality CFG scales without STG/modality passes. - self._guidance_plan = LTXGuidancePlan.build( - LTXGuidanceSpec( - video=LTXModalityGuidance(cfg_scale=self._guidance_plan.spec.video.cfg_scale), - audio=LTXModalityGuidance(cfg_scale=self._guidance_plan.spec.audio.cfg_scale), - ) - ) - if self.do_classifier_free_guidance: - self.guidance_executor.validate_cfg_world_size(self._guidance_plan, cfg_world_size) - return self.do_classifier_free_guidance and cfg_world_size > 1 + guidance_world_size = get_guidance_parallel_world_size() + self.guidance_executor.validate_guidance_world_size(self._guidance_plan, guidance_world_size) + self.guidance_executor.warn_if_imbalanced( + self._guidance_plan, + guidance_world_size, + getattr(self, "_active_phase_name", "generate"), + ) + return self.do_guidance and guidance_world_size > 1 def _check_forward_inputs( self, @@ -594,14 +588,15 @@ class LTXRuntime( padded_length = latent_ops.get_sp_padded_audio_latent_length(requested_length, int(sp_size)) return requested_length if provided_length in {requested_length, padded_length} else provided_length - def _prepare_denoise_context_for_cfg( + def _prepare_denoise_context_for_guidance( self, forward_ctx: LTXForwardContext, denoise_ctx: LTXDenoiseContext, ) -> LTXDenoiseContext: return self.guidance_executor.prepare_denoise_context( self._guidance_plan, - forward_ctx.cfg_parallel_ready, + forward_ctx.guidance_parallel_ready, + get_guidance_parallel_world_size(), denoise_ctx, ) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 8f6a88190..83e975528 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -678,8 +678,7 @@ class OmniServeCommand(CLISubcommand): "--cfg-parallel-size", type=int, default=1, - choices=[1, 2], - help="Number of devices for CFG parallel computation for diffusion models. " + help="Number of devices used to execute diffusion guidance passes in parallel. " "Equivalent to setting DiffusionParallelConfig.cfg_parallel_size.", ) omni_config_group.add_argument(