[diffusion] chore: revise process title (#18446)
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user