From 49adb37e3756608243258b353885c9f22bfa1306 Mon Sep 17 00:00:00 2001 From: yh0903 Date: Tue, 30 Dec 2025 02:09:38 -0800 Subject: [PATCH] [diffusion] chore: fix ZMQ binding and model loading for FastWan compatibility (#13978) Co-authored-by: Han Yu Co-authored-by: Mick Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../runtime/entrypoints/cli/serve.py | 1 - .../entrypoints/diffusion_generator.py | 4 +-- .../runtime/loader/component_loader.py | 1 + .../runtime/loader/fsdp_load.py | 15 +++++++--- .../runtime/managers/scheduler.py | 2 +- .../runtime/scheduler_client.py | 8 +++--- .../multimodal_gen/runtime/server_args.py | 28 ++++++++----------- 7 files changed, 31 insertions(+), 28 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/cli/serve.py b/python/sglang/multimodal_gen/runtime/entrypoints/cli/serve.py index 40207b400..5dc285c87 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/cli/serve.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/cli/serve.py @@ -31,7 +31,6 @@ def add_multimodal_gen_serve_args(parser: argparse.ArgumentParser): def execute_serve_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None): """The entry point for the serve command.""" server_args = ServerArgs.from_cli_args(args, unknown_args) - server_args.post_init_serve() launch_server(server_args) if server_args.webui: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 1b1cb8d16..fea68a3d2 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -145,12 +145,12 @@ class DiffGenerator: if not sync_scheduler_client.ping(): raise ConnectionError( f"Could not connect to remote scheduler at " - f"{self.server_args.scheduler_endpoint()} with `local mode` as False. " + f"{self.server_args.scheduler_endpoint} with `local mode` as False. " "Please ensure the server is running." ) logger.info( f"Successfully connected to remote scheduler at " - f"{self.server_args.scheduler_endpoint()}." + f"{self.server_args.scheduler_endpoint}." ) def generate( diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loader.py index 98f597541..100bdc321 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loader.py @@ -728,6 +728,7 @@ class TransformerLoader(ComponentLoader): param_dtype=torch.bfloat16, reduce_dtype=torch.float32, output_dtype=None, + strict=False, ) total_params = sum(p.numel() for p in model.parameters()) diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index 7f326ba31..89ee92883 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -79,6 +79,7 @@ def maybe_load_fsdp_model( fsdp_inference: bool = False, output_dtype: torch.dtype | None = None, pin_cpu_memory: bool = True, + strict: bool = True, ) -> torch.nn.Module: """ Load the model with FSDP if is training, else load the model without FSDP. @@ -138,7 +139,7 @@ def maybe_load_fsdp_model( weight_iterator, device, default_dtype, - strict=True, + strict=strict, cpu_offload=cpu_offload, param_names_mapping=param_names_mapping_fn, ) @@ -255,9 +256,15 @@ def load_model_from_full_model_state_dict( for target_param_name, full_tensor in custom_param_sd.items(): meta_sharded_param = meta_sd.get(target_param_name) if meta_sharded_param is None: - raise ValueError( - f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect." - ) + if strict: + raise ValueError( + f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect." + ) + else: + logger.warning( + f"Parameter '{target_param_name}' from checkpoint not found in model; skipping. This is expected for optional parameters." + ) + continue if not hasattr(meta_sharded_param, "device_mesh"): full_tensor = full_tensor.to(device=device, dtype=param_dtype) actual_param = param_dict.get(target_param_name) diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index c3e0c2890..cc931968f 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/python/sglang/multimodal_gen/runtime/managers/scheduler.py @@ -49,7 +49,7 @@ class Scheduler: # Inter-process Communication self.context = zmq.Context(io_threads=2) - endpoint = server_args.scheduler_endpoint() + endpoint = server_args.scheduler_endpoint if gpu_id == 0: # router allocates identify (envelope) for each connection self.receiver, actual_endpoint = get_zmq_socket( diff --git a/python/sglang/multimodal_gen/runtime/scheduler_client.py b/python/sglang/multimodal_gen/runtime/scheduler_client.py index 5bfc62533..caec33da4 100644 --- a/python/sglang/multimodal_gen/runtime/scheduler_client.py +++ b/python/sglang/multimodal_gen/runtime/scheduler_client.py @@ -70,7 +70,7 @@ class SchedulerClient: # 100 minute timeout for generation self.scheduler_socket.setsockopt(zmq.RCVTIMEO, 6000000) - scheduler_endpoint = self.server_args.scheduler_endpoint() + scheduler_endpoint = self.server_args.scheduler_endpoint self.scheduler_socket.connect(scheduler_endpoint) logger.debug( f"SchedulerClient connected to backend scheduler at {scheduler_endpoint}" @@ -98,7 +98,7 @@ class SchedulerClient: ping_socket.setsockopt(zmq.LINGER, 0) ping_socket.setsockopt(zmq.RCVTIMEO, 2000) # 2-second timeout for pings - endpoint = self.server_args.scheduler_endpoint() + endpoint = self.server_args.scheduler_endpoint try: ping_socket.connect(endpoint) @@ -157,7 +157,7 @@ class AsyncSchedulerClient: # 100 minute timeout socket.setsockopt(zmq.RCVTIMEO, 6000000) - endpoint = self.server_args.scheduler_endpoint() + endpoint = self.server_args.scheduler_endpoint socket.connect(endpoint) try: @@ -182,7 +182,7 @@ class AsyncSchedulerClient: ping_socket.setsockopt(zmq.LINGER, 0) ping_socket.setsockopt(zmq.RCVTIMEO, 2000) - endpoint = self.server_args.scheduler_endpoint() + endpoint = self.server_args.scheduler_endpoint try: ping_socket.connect(endpoint) diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 9161dff02..e89690fb7 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -220,9 +220,9 @@ class ServerArgs: # TODO: do not hard code master_port: int | None = None - # http server endpoint config, would be ignored in local mode - host: str | None = None - port: int | None = None + # http server endpoint config + host: str | None = "127.0.0.1" + port: int | None = 30000 # TODO: webui and their endpoint, check if webui_port is available. webui: bool = False @@ -290,6 +290,8 @@ class ServerArgs: if self.attention_backend in ["fa3", "fa4"]: self.attention_backend = "fa" + # network initialization: port and host + self.port = self.settle_port(self.port) # Add randomization to avoid race condition when multiple servers start simultaneously initial_scheduler_port = self.scheduler_port + random.randint(0, 100) self.scheduler_port = self.settle_port(initial_scheduler_port) @@ -306,6 +308,7 @@ class ServerArgs: "Failed to load V-MoBA config from %s: %s", self.moba_config_path, e ) raise + self.check_server_args() # log clean server_args @@ -575,12 +578,15 @@ class ServerArgs: else: return f"http://{self.host}:{self.port}" + @property def scheduler_endpoint(self): """ - Internal endpoint for scheduler - + Internal endpoint for scheduler. + Prefers the configured host but normalizes localhost -> 127.0.0.1 to avoid ZMQ issues. """ - scheduler_host = self.host or "localhost" + scheduler_host = self.host + if scheduler_host is None or scheduler_host == "localhost": + scheduler_host = "127.0.0.1" return f"tcp://{scheduler_host}:{self.scheduler_port}" def settle_port( @@ -623,16 +629,6 @@ class ServerArgs: f"(started from port {original_port})" ) - def post_init_serve(self): - """ - Post init when in serve mode - """ - if self.host is None: - self.host = "localhost" - if self.port is None: - self.port = 3000 - self.port = self.settle_port(self.port) - @classmethod def from_cli_args( cls, args: argparse.Namespace, unknown_args: list[str] | None = None