diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 4673a798b..6ce092d85 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -12,11 +12,20 @@ from setproctitle import setproctitle from sglang.multimodal_gen import envs from sglang.multimodal_gen.runtime.distributed import ( get_sp_group, + get_tp_rank, + get_tp_world_size, maybe_init_distributed_environment_and_model_parallel, + model_parallel_is_initialized, ) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_cfg_group, + get_classifier_free_guidance_rank, + get_classifier_free_guidance_world_size, + get_ring_parallel_rank, + get_ring_parallel_world_size, get_tp_group, + get_ulysses_parallel_rank, + get_ulysses_parallel_world_size, ) from sglang.multimodal_gen.runtime.entrypoints.utils import save_outputs from sglang.multimodal_gen.runtime.pipelines_core import ( @@ -70,7 +79,6 @@ class GPUWorker: def init_device_and_model(self) -> None: """Initialize the device and load the model.""" - setproctitle(f"sgl_diffusion::scheduler_TP{self.local_rank}") torch.get_device_module().set_device(self.local_rank) # Set environment variables for distributed initialization os.environ["MASTER_ADDR"] = "localhost" @@ -78,7 +86,7 @@ class GPUWorker: os.environ["LOCAL_RANK"] = str(self.local_rank) os.environ["RANK"] = str(self.rank) os.environ["WORLD_SIZE"] = str(self.server_args.num_gpus) - # Initialize the distributed environment + # initialize the distributed environment maybe_init_distributed_environment_and_model_parallel( tp_size=self.server_args.tp_size, enable_cfg_parallel=self.server_args.enable_cfg_parallel, @@ -90,6 +98,25 @@ class GPUWorker: dist_timeout=self.server_args.dist_timeout, ) + # set proc title + if model_parallel_is_initialized(): + suffix = "" + if get_tp_world_size() != 1: + tp_rank = get_tp_rank() + suffix += f"_TP{tp_rank}" + if get_ulysses_parallel_world_size() != 1: + u_rank = get_ulysses_parallel_rank() + suffix += f"_U{u_rank}" + if get_ring_parallel_world_size() != 1: + r_rank = get_ring_parallel_rank() + suffix += f"_R{r_rank}" + if get_classifier_free_guidance_world_size() != 1: + c_rank = get_classifier_free_guidance_rank() + suffix += f"_C{c_rank}" + setproctitle(f"sgl_diffusion::scheduler{suffix}") + else: + setproctitle(f"sgl_diffusion::scheduler_{self.local_rank}") + self.pipeline = build_pipeline(self.server_args) # apply layerwise offload after lora is applied while building LoRAPipeline