diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py b/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py index 255a49af6587..6befe80b492d 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py @@ -14,9 +14,9 @@ class ModelDeploymentConfig: dit_layerwise_offload_modes: tuple[Literal["auto", "memory"], ...] = () auto_dit_offload_prefetch_size: float | None = None keep_resident_min_available_gb: float | None = None - # Per-model resident defaults. Auto mode additionally keeps an image DiT - # resident above the image workload memory threshold; video DiT placement - # stays with the model's FSDP/layerwise policy. + # Fallback residency for offline/no-warmup deployments. Eligible server + # warmup paths start load-safe and replace these coarse thresholds with a + # measured per-phase placement plan before reporting ready. keep_resident_components: tuple[OffloadComponentName, ...] = ("vae",) fsdp_auto_min_available_memory_gb: float | None = None fsdp_auto_requires_cfg: bool = True diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py b/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py index 1b99d00e9ffe..99615f0ecbd5 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py @@ -165,11 +165,15 @@ def __post_init__(self): def get_model_deployment_config(self) -> ModelDeploymentConfig: return ModelDeploymentConfig( dit_layerwise_offload_modes=("memory",), + # The two-stage path is not numerically invariant when its text + # encoders or VAE use layerwise offload. Keep them resident when + # there is enough headroom, while preserving the low-memory + # layerwise path on consumer GPUs. + keep_resident_min_available_gb=70, + keep_resident_components=("text_encoder", "vae"), # Conservative auto-FSDP gate for the 720p world-model path. Users # can still force FSDP explicitly on smaller cards. fsdp_auto_min_available_memory_gb=60, - keep_resident_min_available_gb=120, - keep_resident_components=("dit", "vae"), ) # --- Latent shape --- diff --git a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py index 1e135b3e82c4..9d69bfd7c8a6 100644 --- a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py +++ b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py @@ -231,10 +231,12 @@ def set_lora_weights( If False, append to existing list (for multi-LoRA support). output_offset: Optional constant output term paired with this adapter """ - lora_A_param = torch.nn.Parameter( - A - ) # share storage with weights in the pipeline - lora_B_param = torch.nn.Parameter(B) + # This runtime never trains adapters. In particular, weights loaded + # during a warmup may be inference tensors; leaving Parameter's + # requires_grad default enabled makes a later dynamic-LoRA request try + # to read their unavailable version counters while sharding views. + lora_A_param = torch.nn.Parameter(A, requires_grad=False) + lora_B_param = torch.nn.Parameter(B, requires_grad=False) output_offset_param = ( torch.nn.Parameter(output_offset, requires_grad=False) if output_offset is not None diff --git a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py index 6a750db23d33..d3929e66c8f0 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py @@ -22,6 +22,7 @@ ComponentResidencyStrategy, ComponentUse, ResidencyState, + build_component_residency_strategy, ) from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, @@ -456,6 +457,12 @@ class LTX2TwoStageResidencyStrategy(ComponentResidencyStrategy): def __init__(self, manager: "LTX2TwoStageResidencyController") -> None: self.manager = manager + self._placement_strategies: dict[ + tuple[str, int, str], ComponentResidencyStrategy + ] = {} + + def supports_auto_residency(self) -> bool: + return True @property def pipeline(self) -> "LTX2TwoStagePipeline": @@ -473,6 +480,19 @@ def _phase(self, use: ComponentUse) -> str: def initialize(self) -> None: pass + def _placement_strategy( + self, module: torch.nn.Module, use: ComponentUse + ) -> ComponentResidencyStrategy: + mode = self.server_args.residency_mode(use.component_name) + key = (use.component_name, id(module), mode) + strategy = self._placement_strategies.get(key) + if strategy is None: + strategy = build_component_residency_strategy( + use.component_name, module, self.server_args + ) + self._placement_strategies[key] = strategy + return strategy + def prepare_for_use( self, module: torch.nn.Module, @@ -482,6 +502,20 @@ def prepare_for_use( phase = self._phase(use) if phase != self.manager._active_phase: self.enter_phase(phase) + self._placement_strategy(module, use).prepare_for_use(module, use, state) + + def prefetch_for_use( + self, + module: torch.nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> bool: + phase = self._phase(use) + if phase != self.manager._active_phase: + self.enter_phase(phase) + return self._placement_strategy(module, use).prefetch_for_use( + module, use, state + ) def wait_for_use( self, @@ -489,6 +523,7 @@ def wait_for_use( use: ComponentUse, state: ResidencyState, ) -> None: + self._placement_strategy(module, use).wait_for_use(module, use, state) self.ensure_phase_ready(self._phase(use)) def finish_use( @@ -497,6 +532,7 @@ def finish_use( use: ComponentUse, state: ResidencyState, ) -> None: + self._placement_strategy(module, use).finish_use(module, use, state) self.exit_phase(self._phase(use)) def finish_request( @@ -507,8 +543,11 @@ def finish_request( *, preferred: bool, ) -> None: + self._placement_strategy(module, use).finish_request( + module, use, state, preferred=preferred + ) if not preferred: - self.finish_use(module, use, state) + self.exit_phase(self._phase(use)) return phase = self._phase(use) if phase != self.manager._active_phase: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py index 5f6620b8b3f0..fe7e0b663814 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py @@ -293,26 +293,48 @@ def _temporarily_disable_offload( yield [] return - # clear device cache to free up unused memory - if torch.get_device_module().is_available(): - torch.get_device_module().synchronize() - torch.get_device_module().empty_cache() + offloaded_module_names = [ + module_name + for module_name in module_names + if ( + (module := self.modules.get(module_name)) is not None + and is_layerwise_offloaded_module(module) + ) + ] + # Record every target, including coarse component offload. The current + # coarse path can already materialize the full component here, and the + # planner must preserve that phase if it later chooses layerwise mode. + residency_manager = getattr(self, "component_residency_manager", None) + residency_transition = ( + residency_manager.full_weight_transition(module_names) + if residency_manager is not None + else nullcontext() + ) - offload_disabled_modules = [] - for module_name in module_names: - module = self.modules.get(module_name) - if isinstance(module, torch.nn.Module): - restore_weight_snapshot(module) - if module is not None and is_layerwise_offloaded_module(module): + with residency_transition: + # clear device cache to free up unused memory + if torch.get_device_module().is_available(): + torch.get_device_module().synchronize() + torch.get_device_module().empty_cache() + + # snapshot-offloaded components merge into their restored weights + for module_name in module_names: + module = self.modules.get(module_name) + if isinstance(module, torch.nn.Module): + restore_weight_snapshot(module) + + offload_disabled_modules = [] + for module_name in offloaded_module_names: + module = self.modules[module_name] module.disable_offload() offload_disabled_modules.append(module) - try: - yield offload_disabled_modules - finally: - # Re-enable layerwise offload: sync weights to CPU and restore hooks - for module in offload_disabled_modules: - module.enable_offload() + try: + yield offload_disabled_modules + finally: + # Re-enable layerwise offload: sync weights to CPU and restore hooks + for module in offload_disabled_modules: + module.enable_offload() def _needs_lora_weight_update_context( self, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py index ca680fc9afb9..49fa729a1110 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py @@ -246,7 +246,7 @@ def _encode_delight_prompts( return torch.cat([positive, negative, negative]) @torch.no_grad() - def _run_delight(self, image: Image.Image) -> Image.Image: + def _run_delight(self, image: Image.Image, batch: Req) -> Image.Image: assert self.delight_transformer is not None assert self.delight_vae is not None assert self.delight_scheduler is not None @@ -278,7 +278,14 @@ def _run_delight(self, image: Image.Image) -> Image.Image: ) scheduler = self.delight_scheduler - scheduler.set_timesteps(self.config.delight_num_inference_steps, device=device) + target_steps = self.config.delight_num_inference_steps + measured_steps = ( + min(target_steps, max(1, int(batch.num_inference_steps))) + if batch.is_warmup + else target_steps + ) + batch.record_stage_iterations(measured_steps, target_steps) + scheduler.set_timesteps(measured_steps, device=device) generator = torch.Generator(device="cpu").manual_seed(42) latent_channels = self.delight_transformer.config.out_channels latents = randn_tensor( @@ -350,11 +357,11 @@ def _run_delight(self, image: Image.Image) -> Image.Image: (composited.clamp(0, 1).cpu().numpy() * 255).astype(np.uint8) ) - def _prepare_reference_image(self, image_path: str) -> Image.Image: + def _prepare_reference_image(self, image_path: str, batch: Req) -> Image.Image: image = self._load_input_image(image_path) if not self.config.delight_enable: return image - return self._run_delight(image) + return self._run_delight(image, batch) def _render_multiview( self, mesh: Any @@ -404,7 +411,9 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: mesh_future = executor.submit(self._unwrap_mesh, mesh) - image_future = executor.submit(self._prepare_reference_image, image_path) + image_future = executor.submit( + self._prepare_reference_image, image_path, batch + ) paint_mesh = mesh_future.result() delighted_image = image_future.result() @@ -542,15 +551,27 @@ def _camera_index(azimuth: float, elevation: float) -> int: base, divisor = 12, 1 return base + base_index // divisor - def _timesteps(self, device: torch.device) -> torch.Tensor: + def _timesteps( + self, device: torch.device, batch: Req | None = None + ) -> torch.Tensor: + target_steps = ( + 10 + if self.config.paint_turbo_mode + else (self.config.paint_num_inference_steps) + ) + measured_steps = ( + min(target_steps, max(1, int(batch.num_inference_steps))) + if batch is not None and batch.is_warmup + else target_steps + ) + if batch is not None: + batch.record_stage_iterations(measured_steps, target_steps) if not self.config.paint_turbo_mode: - self.scheduler.set_timesteps( - self.config.paint_num_inference_steps, device=device - ) + self.scheduler.set_timesteps(measured_steps, device=device) return self.scheduler.timesteps self.scheduler.set_timesteps( - num_inference_steps=10, + num_inference_steps=measured_steps, original_inference_steps=30, device=device, ) @@ -639,7 +660,7 @@ def _prepare_denoising_inputs(self, batch: Req) -> PaintDenoisingInputs: if position_attention_mask is not None: model_kwargs["position_attn_mask"] = position_attention_mask - timesteps = self._timesteps(device) + timesteps = self._timesteps(device, batch) latent_channels = self.transformer.config.in_channels latent_size = render_size // self.vae_scale_factor latents = randn_tensor( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/upsampling.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/upsampling.py index 9afa06359bee..75d1a26a0b1c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/upsampling.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/upsampling.py @@ -50,7 +50,10 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: if self.pipeline.should_skip_ltx2_lora_switch_stage(): batch.extra["ltx2_phase"] = self.phase return batch - self.pipeline.switch_lora_phase(self.phase, batch=batch) + # The installed adapter survives this stage. Give its tensors version + # counters for later offload stages, which run outside inference mode. + with torch.inference_mode(False), torch.no_grad(): + self.pipeline.switch_lora_phase(self.phase, batch=batch) batch.extra["ltx2_phase"] = self.phase return batch diff --git a/python/sglang/multimodal_gen/test/server/consistency_thresholds/5090.json b/python/sglang/multimodal_gen/test/server/consistency_thresholds/5090.json index 9f6f79fa4ca1..f575e698d4c4 100644 --- a/python/sglang/multimodal_gen/test/server/consistency_thresholds/5090.json +++ b/python/sglang/multimodal_gen/test/server/consistency_thresholds/5090.json @@ -5,6 +5,12 @@ "ssim_threshold": 0.86, "psnr_threshold": 19.5, "mean_abs_diff_threshold": 8.5 + }, + "wan2_1_t2v_1.3b": { + "clip_threshold": 0.95, + "ssim_threshold": 0.85, + "psnr_threshold": 25.0, + "mean_abs_diff_threshold": 8.0 } } } diff --git a/python/sglang/multimodal_gen/test/unit/test_hunyuan3d_native_texture_models.py b/python/sglang/multimodal_gen/test/unit/test_hunyuan3d_native_texture_models.py index cbfec5d64d0e..4f17edb9f482 100644 --- a/python/sglang/multimodal_gen/test/unit/test_hunyuan3d_native_texture_models.py +++ b/python/sglang/multimodal_gen/test/unit/test_hunyuan3d_native_texture_models.py @@ -2,6 +2,7 @@ import unittest from types import SimpleNamespace +from unittest.mock import Mock import torch from diffusers import AutoencoderKL as DiffusersAutoencoderKL @@ -222,6 +223,24 @@ def test_uses_standard_lcm_schedule_without_custom_timesteps(self): ) self.assertFalse(stage.scheduler.custom_timesteps) + def test_warmup_caps_paint_steps_and_records_full_target(self): + stage = Hunyuan3DPaintTexGenStage.__new__(Hunyuan3DPaintTexGenStage) + stage.config = Hunyuan3D2PipelineConfig(paint_turbo_mode=True) + stage.scheduler = LCMScheduler( + num_train_timesteps=1000, + original_inference_steps=50, + ) + batch = SimpleNamespace( + is_warmup=True, + num_inference_steps=4, + record_stage_iterations=Mock(), + ) + + timesteps = stage._timesteps(torch.device("cpu"), batch) + + self.assertEqual(len(timesteps), 4) + batch.record_stage_iterations.assert_called_once_with(4, 10) + if __name__ == "__main__": unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_lora_inference_mode.py b/python/sglang/multimodal_gen/test/unit/test_lora_inference_mode.py index ae1b7db42242..7ffea359aabc 100644 --- a/python/sglang/multimodal_gen/test/unit/test_lora_inference_mode.py +++ b/python/sglang/multimodal_gen/test/unit/test_lora_inference_mode.py @@ -1,3 +1,5 @@ +from types import SimpleNamespace + import torch from torch import nn @@ -5,6 +7,9 @@ LinearWithLoRA, _compute_lora_delta, ) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.upsampling import ( + LTX2LoRASwitchStage, +) def test_stacked_lora_delta_preserves_projection_order(): @@ -49,3 +54,47 @@ def test_lora_merge_unmerge_handles_inference_base_weight(): assert not layer.merged assert not layer.base_layer.weight.is_inference() assert torch.allclose(layer.base_layer.weight, base_weight) + + +def test_dynamic_lora_reuses_inference_weights_without_autograd_tracking(): + base_layer = nn.Linear(4, 3, bias=False) + layer = LinearWithLoRA(base_layer, lora_rank=2, lora_alpha=2) + + with torch.inference_mode(): + layer.set_lora_weights( + torch.ones(2, 4), + torch.ones(3, 2), + clear_existing=True, + merge_weights=False, + ) + + assert layer.lora_A.is_inference() + assert layer.lora_B.is_inference() + assert not layer.lora_A.requires_grad + assert not layer.lora_B.requires_grad + with torch.no_grad(): + sharded_view = layer.lora_B[:2, :] + assert sharded_view.shape == (2, 2) + + +def test_ltx2_lora_switch_creates_versioned_adapter_tensors(): + class _Pipeline: + adapter = None + + @staticmethod + def should_skip_ltx2_lora_switch_stage(): + return False + + def switch_lora_phase(self, phase, *, batch): + self.adapter = torch.ones(1) + + pipeline = _Pipeline() + stage = LTX2LoRASwitchStage(pipeline, "stage2") + batch = SimpleNamespace(extra={}) + + with torch.inference_mode(): + stage.forward(batch, SimpleNamespace()) + + assert pipeline.adapter is not None + assert not pipeline.adapter.is_inference() + assert batch.extra["ltx2_phase"] == "stage2" diff --git a/python/sglang/multimodal_gen/test/unit/test_lora_pipeline.py b/python/sglang/multimodal_gen/test/unit/test_lora_pipeline.py index 546f3ea5c0c0..3c2c51fc9bb9 100644 --- a/python/sglang/multimodal_gen/test/unit/test_lora_pipeline.py +++ b/python/sglang/multimodal_gen/test/unit/test_lora_pipeline.py @@ -1,7 +1,7 @@ from collections import defaultdict from contextlib import contextmanager, nullcontext from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest import torch @@ -48,12 +48,59 @@ def _make_pipeline(layer: BaseLayerWithLoRA) -> _TestLoRAPipeline: pipeline.lora_layers_transformer_2 = {} pipeline.lora_layers_critic = {} pipeline.is_lora_merged = {} + pipeline.component_residency_manager = None pipeline.lora_adapters["adapter"]["linear.lora_A"] = torch.ones(1, 2) pipeline.lora_adapters["adapter"]["linear.lora_B"] = torch.ones(2, 1) return pipeline +def test_temporary_offload_materialization_is_reported_to_residency_manager(): + pipeline = _make_pipeline(_make_layer()) + module = Mock() + pipeline.modules = {"transformer": module} + residency_manager = Mock() + residency_manager.full_weight_transition.return_value = nullcontext() + pipeline.component_residency_manager = residency_manager + + with patch( + "sglang.multimodal_gen.runtime.pipelines_core.lora.pipeline." + "is_layerwise_offloaded_module", + return_value=True, + ): + with pipeline._temporarily_disable_offload( + target="transformer", use_module_names_only=True + ): + pass + + residency_manager.full_weight_transition.assert_called_once_with(["transformer"]) + module.disable_offload.assert_called_once_with() + module.enable_offload.assert_called_once_with() + + +def test_component_offload_weight_update_is_reported_without_disabling_offload(): + pipeline = _make_pipeline(_make_layer()) + module = Mock() + pipeline.modules = {"transformer": module} + residency_manager = Mock() + residency_manager.full_weight_transition.return_value = nullcontext() + pipeline.component_residency_manager = residency_manager + + with patch( + "sglang.multimodal_gen.runtime.pipelines_core.lora.pipeline." + "is_layerwise_offloaded_module", + return_value=False, + ): + with pipeline._temporarily_disable_offload( + target="transformer", use_module_names_only=True + ): + pass + + residency_manager.full_weight_transition.assert_called_once_with(["transformer"]) + module.disable_offload.assert_not_called() + module.enable_offload.assert_not_called() + + def test_merge_cache_only_accepts_cpu_backed_weights(): pipeline = _make_pipeline(_make_layer()) cpu_cache = pipeline._merge_cache_for( diff --git a/python/sglang/multimodal_gen/test/unit/test_ltx2_5_config.py b/python/sglang/multimodal_gen/test/unit/test_ltx2_5_config.py index 256344885adb..daff47f8426d 100644 --- a/python/sglang/multimodal_gen/test/unit/test_ltx2_5_config.py +++ b/python/sglang/multimodal_gen/test/unit/test_ltx2_5_config.py @@ -699,5 +699,40 @@ def test_explicit_component_path_wins(self): self.assertEqual(args.component_paths["transformer"], "/some/other/transformer") +class TestLTX2TwoStageAutoResidency(unittest.TestCase): + def test_custom_strategy_delegates_after_runtime_placement_change(self): + from sglang.multimodal_gen.runtime.pipelines import ltx_2_pipeline + + mode = {"value": "resident"} + server_args = SimpleNamespace( + residency_mode=lambda _name: mode["value"], + ) + manager = SimpleNamespace( + pipeline=SimpleNamespace(), + server_args=server_args, + _active_phase="stage1", + ) + strategy = ltx_2_pipeline.LTX2TwoStageResidencyStrategy(manager) + resident = mock.Mock() + offloaded = mock.Mock() + module = mock.Mock() + use = SimpleNamespace(component_name="transformer", phase="stage1") + state = SimpleNamespace() + + with mock.patch.object( + ltx_2_pipeline, + "build_component_residency_strategy", + side_effect=(resident, offloaded), + ) as build: + strategy.prepare_for_use(module, use, state) + mode["value"] = "component-offload" + strategy.prepare_for_use(module, use, state) + + self.assertTrue(strategy.supports_auto_residency()) + self.assertEqual(build.call_count, 2) + resident.prepare_for_use.assert_called_once_with(module, use, state) + offloaded.prepare_for_use.assert_called_once_with(module, use, state) + + if __name__ == "__main__": unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index 81e43ef77df2..009866441783 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -36,6 +36,7 @@ ) from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( LTX2PipelineConfig, + LTX23PipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import ( MiniMaxH3PipelineConfig, @@ -1639,8 +1640,10 @@ def test_sana_wm_and_helios_declare_model_residency_hints(self): sana_wm_deployment = SanaWMPipelineConfig().get_model_deployment_config() self.assertEqual(sana_wm_deployment.fsdp_auto_min_available_memory_gb, 60) self.assertEqual(sana_wm_deployment.dit_layerwise_offload_modes, ("memory",)) - self.assertEqual(sana_wm_deployment.keep_resident_min_available_gb, 120) - self.assertEqual(sana_wm_deployment.keep_resident_components, ("dit", "vae")) + self.assertEqual(sana_wm_deployment.keep_resident_min_available_gb, 70) + self.assertEqual( + sana_wm_deployment.keep_resident_components, ("text_encoder", "vae") + ) helios_deployment = HeliosDistilledConfig().get_model_deployment_config() self.assertEqual(helios_deployment.keep_resident_min_available_gb, 120) @@ -1656,6 +1659,108 @@ def test_realtime_server_keeps_model_residency_fallback(self): self.assertEqual(args.residency_mode("transformer"), RESIDENT) + def test_pipeline_configs_declare_auto_tune_hints(self): + qwen_deployment = QwenImagePipelineConfig().get_model_deployment_config() + cosmos3_deployment = Cosmos3Config( + model_path="nvidia/Cosmos3-Nano" + ).get_model_deployment_config() + cosmos3_super_deployment = Cosmos3Config( + model_path="nvidia/Cosmos3-Super" + ).get_model_deployment_config() + local_cosmos3_nano_deployment = Cosmos3Config( + model_path="/models/custom-checkpoint", is_nano=True + ).get_model_deployment_config() + wan_deployment = WanT2V480PConfig().get_model_deployment_config() + mova_deployment = MOVAPipelineConfig().get_model_deployment_config() + zimage_deployment = ZImagePipelineConfig().get_model_deployment_config() + lingbot_deployment = LingBotWorldCausalDMDConfig().get_model_deployment_config() + ltx_deployment = LTX2PipelineConfig().get_model_deployment_config() + ltx23_config = LTX23PipelineConfig() + longlive_deployment = LongLive2T2VConfig().get_model_deployment_config() + sana_wm_deployment = SanaWMPipelineConfig().get_model_deployment_config() + + self.assertIsNone(qwen_deployment.fsdp_auto_min_available_memory_gb) + self.assertEqual(qwen_deployment.dit_layerwise_offload_modes, ()) + + self.assertEqual(cosmos3_deployment.keep_resident_min_available_gb, 90) + self.assertEqual(cosmos3_deployment.keep_resident_components, ("dit", "vae")) + self.assertEqual(cosmos3_super_deployment.keep_resident_min_available_gb, 120) + self.assertEqual( + local_cosmos3_nano_deployment.keep_resident_min_available_gb, 90 + ) + + self.assertIsNone(wan_deployment.fsdp_auto_min_available_memory_gb) + self.assertEqual(wan_deployment.dit_layerwise_offload_modes, ("memory",)) + self.assertEqual(wan_deployment.keep_resident_min_available_gb, 60) + self.assertEqual(wan_deployment.keep_resident_components, ("dit",)) + + self.assertIsNone(mova_deployment.fsdp_auto_min_available_memory_gb) + self.assertEqual( + mova_deployment.dit_layerwise_offload_modes, ("auto", "memory") + ) + self.assertEqual(mova_deployment.keep_resident_min_available_gb, 130) + self.assertEqual(mova_deployment.keep_resident_components, ("dit", "vae")) + + self.assertEqual(zimage_deployment.fsdp_auto_min_available_memory_gb, 40) + self.assertEqual(zimage_deployment.keep_resident_min_available_gb, 30) + self.assertTrue(zimage_deployment.fsdp_auto_requires_cfg) + self.assertEqual(zimage_deployment.dit_layerwise_offload_modes, ()) + + self.assertEqual(lingbot_deployment.dit_layerwise_offload_modes, ("memory",)) + self.assertEqual(lingbot_deployment.keep_resident_min_available_gb, 70) + self.assertEqual(lingbot_deployment.keep_resident_components, ("dit",)) + + self.assertEqual(ltx_deployment.keep_resident_min_available_gb, 70) + self.assertEqual(ltx_deployment.keep_resident_components, ("dit",)) + self.assertEqual( + ltx_deployment.auto_cfg_parallel_degree_by_num_gpus, ((4, 1), (8, 1)) + ) + self.assertEqual(ltx_deployment.get_auto_cfg_parallel_degree(4), 1) + self.assertEqual(ltx_deployment.get_auto_cfg_parallel_degree(8), 1) + self.assertEqual(ltx_deployment.get_auto_cfg_parallel_degree(2), 2) + self.assertEqual(longlive_deployment.keep_resident_min_available_gb, 60) + self.assertEqual( + longlive_deployment.keep_resident_components, + ("dit", "text_encoder", "vae"), + ) + self.assertFalse( + LTX2PipelineConfig().dit_config.arch_config.enable_packed_qkv_input_a2a + ) + self.assertFalse( + ltx23_config.dit_config.arch_config.enable_packed_qkv_input_a2a + ) + self.assertFalse(ltx23_config.supports_auto_residency) + self.assertFalse(QwenImageLayeredPipelineConfig().supports_auto_residency) + + self.assertEqual(sana_wm_deployment.fsdp_auto_min_available_memory_gb, 60) + self.assertEqual(sana_wm_deployment.dit_layerwise_offload_modes, ("memory",)) + self.assertEqual(sana_wm_deployment.keep_resident_min_available_gb, 70) + self.assertEqual( + sana_wm_deployment.keep_resident_components, ("text_encoder", "vae") + ) + + fast_hunyuan_deployment = FastHunyuanConfig().get_model_deployment_config() + self.assertEqual(fast_hunyuan_deployment.keep_resident_min_available_gb, 60) + self.assertEqual( + fast_hunyuan_deployment.keep_resident_components, ("dit", "vae") + ) + + fast_wan_deployment = FastWan2_2_TI2V_5B_Config().get_model_deployment_config() + self.assertEqual(fast_wan_deployment.keep_resident_min_available_gb, 60) + self.assertEqual(fast_wan_deployment.keep_resident_components, ("dit",)) + + for dual_dit_config in ( + Wan2_2_T2V_A14B_Config(), + Wan2_2_I2V_A14B_Config(), + ): + dual_dit_deployment = dual_dit_config.get_model_deployment_config() + self.assertIsNone(dual_dit_deployment.keep_resident_min_available_gb) + self.assertEqual(dual_dit_deployment.keep_resident_components, ("vae",)) + + # default keeps only vae resident (encoders are large, dit owned by FSDP) + self.assertEqual(qwen_deployment.keep_resident_components, ("vae",)) + self.assertIsNone(qwen_deployment.keep_resident_min_available_gb) + def test_longlive_residency_scales_with_available_memory(self): high_memory_args = self._from_dict_with_pipeline_config( LongLive2T2VConfig(), @@ -1675,6 +1780,25 @@ def test_longlive_residency_scales_with_available_memory(self): self.assertIn("text_encoder", constrained_offload) self.assertIn("vae", constrained_offload) + def test_sana_wm_residency_scales_with_available_memory(self): + high_memory_args = self._from_dict_with_pipeline_config( + SanaWMPipelineConfig(), + memory_gb=80, + kwargs={"performance_mode": "auto"}, + ) + high_memory_offload = high_memory_args.layerwise_offload_components or [] + self.assertNotIn("text_encoder", high_memory_offload) + self.assertNotIn("vae", high_memory_offload) + + constrained_args = self._from_dict_with_pipeline_config( + SanaWMPipelineConfig(), + memory_gb=60, + kwargs={"performance_mode": "auto"}, + ) + constrained_offload = constrained_args.layerwise_offload_components or [] + self.assertIn("text_encoder", constrained_offload) + self.assertIn("vae", constrained_offload) + def test_qwen_ar_generation_residency_scales_with_available_memory(self): pipeline_configs = ( QwenImageLayeredPipelineConfig(),