diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 25288bea8..e78679527 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -272,11 +272,6 @@ def _discover_and_register_pipelines(): cls.pipeline_config_cls, cls.sampling_params_cls, ) - logger.debug( - f"Auto-registered config classes for pipeline '{cls.pipeline_name}': " - f"PipelineConfig={cls.pipeline_config_cls.__name__}, " - f"SamplingParams={cls.sampling_params_cls.__name__}" - ) logger.debug( f"Registering pipelines complete, {len(_PIPELINE_REGISTRY)} pipelines registered" ) diff --git a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py index 71da9c977..f76d53937 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py +++ b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py @@ -11,6 +11,8 @@ import torch.distributed as dist from torch import Tensor from torch.distributed import ProcessGroup, ReduceOp +from sglang.multimodal_gen.runtime.distributed.utils import all_gather_single + def _ipc_all_to_all_4d(group, input_, scatter_dim): """2-rank IPC path for AllToAll4D; None when the transport is unavailable.""" @@ -103,7 +105,7 @@ class DistributedAutograd: output_size, dtype=input_.dtype, device=input_.device ) - dist.all_gather_into_tensor(output_tensor, input_, group=group) + all_gather_single(output_tensor, input_, group=group) output_tensor = output_tensor.reshape((world_size,) + input_size) output_tensor = output_tensor.movedim(0, dim) diff --git a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cpu_communicator.py b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cpu_communicator.py index 08abee7b3..3345cc379 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cpu_communicator.py +++ b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cpu_communicator.py @@ -9,6 +9,8 @@ import os import torch from torch.distributed import ProcessGroup +from sglang.multimodal_gen.runtime.distributed.utils import all_gather_single + from .base_device_communicator import DeviceCommunicatorBase @@ -26,6 +28,7 @@ class CpuCommunicator(DeviceCommunicatorBase): super().__init__(cpu_group, device, device_group, unique_name) self.dist_module = torch.distributed + self._all_gather_single = all_gather_single if ( (current_platform.get_cpu_architecture() == CpuArchEnum.X86) @@ -33,6 +36,7 @@ class CpuCommunicator(DeviceCommunicatorBase): and unique_name.startswith("tp") ): self.dist_module = _CPUSHMDistributed(self) + self._all_gather_single = self.dist_module.all_gather_single def all_reduce( self, @@ -89,9 +93,7 @@ class CpuCommunicator(DeviceCommunicatorBase): output_size, dtype=input_.dtype, device=input_.device ) # All-gather. - self.dist_module.all_gather_into_tensor( - output_tensor, input_, group=self.device_group - ) + self._all_gather_single(output_tensor, input_, group=self.device_group) # Reshape output_tensor = output_tensor.reshape((self.world_size,) + input_size) @@ -153,7 +155,7 @@ class _CPUSHMDistributed: torch.distributed.get_group_rank(group, dst), ) - def all_gather_into_tensor( + def all_gather_single( self, output: torch.Tensor, input: torch.Tensor, diff --git a/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py b/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py index 62b1ff27b..1b660f9bd 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py +++ b/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py @@ -24,6 +24,7 @@ from sglang.multimodal_gen.runtime.distributed.device_communicators.base_device_ from sglang.multimodal_gen.runtime.distributed.device_communicators.cpu_communicator import ( CpuCommunicator, ) +from sglang.multimodal_gen.runtime.distributed.utils import all_gather_single from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import ( init_logger, @@ -417,9 +418,7 @@ class GroupCoordinator: ): return torch.ops.sgl_kernel.shm_allgather(input_, dim) else: - torch.distributed.all_gather_into_tensor( - output_tensor, input_, group=self.device_group - ) + all_gather_single(output_tensor, input_, group=self.device_group) if dim != 0: input_size[0] //= world_size diff --git a/python/sglang/multimodal_gen/runtime/distributed/utils.py b/python/sglang/multimodal_gen/runtime/distributed/utils.py index d56979efc..ef481c9e3 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/utils.py +++ b/python/sglang/multimodal_gen/runtime/distributed/utils.py @@ -18,8 +18,15 @@ from typing import Any import torch from torch.distributed import TCPStore +try: + from torch.distributed import all_gather_single as _all_gather_single +except ImportError: + from torch.distributed import all_gather_into_tensor as _all_gather_single + from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +all_gather_single = _all_gather_single + logger = init_logger(__name__) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py index 5c7a3c917..6422f2da3 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py @@ -343,16 +343,16 @@ class TransformerLoader(ComponentLoader): ) checkpoint_key_filter = _minimax_h3_adaln_cache_key_filter - if ( - init_params["quant_config"] is None - and component_server_args.transformer_weights_path is not None - ): + runtime_quant_config = init_params["quant_config"] + if runtime_quant_config is not None: + logger.debug( + "Runtime quantization: %s", type(runtime_quant_config).__name__ + ) + elif component_server_args.transformer_weights_path is not None: logger.info( "Using an unquantized transformer weight override from %s", component_server_args.transformer_weights_path, ) - else: - logger.debug("quantization config: %s", init_params["quant_config"]) local_torch_device = get_local_torch_device() checkpoint_load_device = ( diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py index ca986e801..3c6293f1e 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py @@ -257,9 +257,9 @@ def describe_host_memory() -> str: capped = cgroup_memory_limit_bytes() available = host_memory_available_bytes() if capped is None: - return f"host memory available: {available / GIB_BYTES:.1f} GiB (no cgroup cap)" + return f"{available / GIB_BYTES:.1f} GiB available (no cgroup cap)" limit, usage = capped return ( - f"host memory available: {available / GIB_BYTES:.1f} GiB " + f"{available / GIB_BYTES:.1f} GiB available " f"(cgroup cap {limit / GIB_BYTES:.1f} GiB, in use {usage / GIB_BYTES:.1f} GiB)" ) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py index 7234e1b31..94827f4fa 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py @@ -4,6 +4,7 @@ import re import threading from collections.abc import Mapping, Sequence from contextlib import nullcontext +from time import perf_counter from typing import Any, Dict, List, Optional, Set, Tuple import torch @@ -811,9 +812,6 @@ class LayerwiseOffloadManager: self.register_forward_hooks() self._configured = True - logger.debug( - f"LayerwiseOffloadManager initialized with num prefetched layer: {self.prefetch_size}, num resident layers: {self.resident_layers}, total num layers: {self.num_layers}, residency policy: {self.residency_policy}" - ) if self.residency_policy == RESIDENCY_POLICY_STRIDED and self._streamed_order: # Printed because the layout is the whole point of the policy, and # "did it actually stride?" is otherwise only answerable from a @@ -1657,7 +1655,7 @@ class LayerwiseOffloadableModuleMixin: ) self.layerwise_offload_managers = [] named_modules = dict(self.named_modules()) - configured_layer_names = [] + layer_specs = [] # `--dit-*` is the group default these fall back to, not a scope. prefetch_value, resident_value, residency_policy = ( server_args.layerwise_tuning_for( @@ -1687,6 +1685,32 @@ class LayerwiseOffloadableModuleMixin: else: resident_layers = min(num_layers, int(resident_value)) + layer_specs.append((layer_name, num_layers, prefetch_size, resident_layers)) + + if not layer_specs: + logger.debug( + "No layerwise-offloadable ModuleList found for %s. Candidates: %s", + self.__class__.__name__, + self.layer_names, + ) + return + + component_label = ( + f"{component_name} ({self.__class__.__name__})" + if component_name is not None + else self.__class__.__name__ + ) + logger.info( + "Configuring layerwise offload for %s: %s", + component_label, + ", ".join( + f"{layer_name} ({num_layers} layers)" + for layer_name, num_layers, _, _ in layer_specs + ), + ) + started_at = perf_counter() + + for layer_name, num_layers, prefetch_size, resident_layers in layer_specs: # Pinning these weights is what lets the copy stream run ahead of # compute, but pinned pages are the ones the kernel cannot reclaim, # so they are handed out only while the budget lasts. The budget goes @@ -1709,7 +1733,6 @@ class LayerwiseOffloadableModuleMixin: residency_policy=residency_policy, ) self.layerwise_offload_managers.append(manager) - configured_layer_names.append(layer_name) if current_platform.is_mps(): for manager in self.layerwise_offload_managers: @@ -1746,18 +1769,26 @@ class LayerwiseOffloadableModuleMixin: for manager in enabled_managers: manager._finalize_initialization() - if configured_layer_names: - logger.debug( - "Enabled layerwise offload for %s on modules: %s", - self.__class__.__name__, - configured_layer_names, - ) - else: - logger.debug( - "No layerwise-offloadable ModuleList found for %s. Candidates: %s", - self.__class__.__name__, - self.layer_names, - ) + managers = self.layerwise_offload_managers + prefetch_sizes = ", ".join( + str(value) + for value in sorted({manager.prefetch_size for manager in managers}) + ) + policies = ", ".join(sorted({manager.residency_policy for manager in managers})) + total_layers = sum(manager.num_layers for manager in managers) + resident_layers = sum(manager.resident_layers for manager in managers) + logger.info( + "Layerwise offload ready for %s in %.2fs: groups=%d, layers=%d, " + "prefetch/group=%s, resident=%d/%d, policy=%s", + component_label, + perf_counter() - started_at, + len(managers), + total_layers, + prefetch_sizes, + resident_layers, + total_layers, + policies, + ) def prepare_for_next_req(self): if self.layerwise_offload_managers is None: @@ -2073,7 +2104,7 @@ def configure_layerwise_offload_modules( reverse=True, ) pin_budget = HostPinBudget() - logger.info("Layerwise offload: %s", describe_host_memory()) + logger.info("Layerwise offload host memory: %s", describe_host_memory()) for component_name in selected_pipeline_component_names: module = modules[component_name] @@ -2108,7 +2139,7 @@ def configure_layerwise_offload_modules( ) logger.info( - "Enabled layerwise offload for pipeline components: %s", + "Layerwise offload summary: %s", ", ".join( f"{name} ({format_component_residency(modules[name])})" for name in configured_component_names diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/common.py b/python/sglang/multimodal_gen/runtime/models/vaes/common.py index e7e4edbeb..cbe9b6af1 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/common.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/common.py @@ -25,6 +25,7 @@ from sglang.multimodal_gen.runtime.distributed import ( get_sp_world_size, model_parallel_is_initialized, ) +from sglang.multimodal_gen.runtime.distributed.utils import all_gather_single from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, ) @@ -369,7 +370,7 @@ class ParallelTiledVAE(ABC, nn.Module, LayerwiseOffloadableModuleMixin): .repeat(world_size, *[1] * len(padded_results.shape)) .contiguous() ) - dist.all_gather_into_tensor(gathered_results, padded_results) + all_gather_single(gathered_results, padded_results) dist.all_gather_object(gathered_dim_metadata, local_dim_metadata) gathered_dim_metadata = cast(list[list[torch.Size]], gathered_dim_metadata) @@ -510,7 +511,7 @@ class ParallelTiledVAE(ABC, nn.Module, LayerwiseOffloadableModuleMixin): device=padded_results.device, dtype=padded_results.dtype, ) - dist.all_gather_into_tensor(gathered_results, padded_results) + all_gather_single(gathered_results, padded_results) dec = z.new_empty( ( diff --git a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py index 7f76c6ed2..01ecb819b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -994,6 +994,40 @@ def test_configure_resolves_residency_policy(monkeypatch): ) +def test_configure_logs_component_start_and_completion(monkeypatch): + _patch_fake_device(monkeypatch) + logs = [] + timestamps = iter((10.0, 12.345)) + monkeypatch.setattr( + layerwise_offload_mod.logger, + "info", + lambda message, *args: logs.append(message % args), + ) + monkeypatch.setattr( + layerwise_offload_mod, + "perf_counter", + lambda: next(timestamps), + ) + + comp = _ResidentComponent(8) + comp.configure_layerwise_offload( + _server_args( + dit_offload_prefetch_size=2, + dit_layerwise_resident_layers=3, + ), + component_name="transformer", + ) + + assert logs[0] == ( + "Configuring layerwise offload for transformer (_ResidentComponent): " + "blocks (8 layers)" + ) + assert logs[-1] == ( + "Layerwise offload ready for transformer (_ResidentComponent) in 2.35s: " + "groups=1, layers=8, prefetch/group=2, resident=3/8, policy=leading" + ) + + def test_configure_offloads_all_layer_groups_before_moving_non_layers(monkeypatch): _patch_fake_device(monkeypatch) model = _MultiGroupComponent()