import os import logging import multiprocessing as mp import traceback from queue import Empty import numpy as np import torch import torch.distributed as dist from vllm import distributed as vllm_dist from transformers import AutoConfig _original_from_pretrained = AutoConfig.from_pretrained from vllm.config import ParallelConfig from vllm.config import VllmConfig, set_current_vllm_config, get_current_vllm_config import socket from typing import Any, Optional, List, Optional, Tuple import random import time from contextlib import closing from .parallel_strategy import ThresholdParallelDecoder, CreditThresholdParallelDecoder, HierarchyDecoder from .utils import KVCacheFactory, BlockIteratorFactory from .generate_uniform import IterSmoothWithVicinityCacheDiffusionLLM, IterSmoothDiffusionLLM, VicinityCacheDiffusionLLM, BlockWiseDiffusionLLM, BlockDiffusionLLM from ..model.modeling_fused_olmoe import FusedOlmoeForCausalLM from ..model.modeling_llada import LLaDAModelLM import time import json import sys from transformers.configuration_utils import PretrainedConfig logger = logging.getLogger(__name__) def is_port_available(port: int) -> bool: """ Check if a single port is available. :param port: The port number to check. :return: Whether the port is available. """ sock_type = socket.SOCK_STREAM try: with closing(socket.socket(socket.AF_INET, sock_type)) as s: s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) s.bind(('0.0.0.0', port)) return True except OSError: return False def find_continuous_ports( num_ports: int, start_port: int = 30000, end_port: int = 65535, protocol: str = 'tcp', max_attempts: int = 300, retry_delay: float = 0.1 ) -> Optional[Tuple[int, List[int]]]: """ Find a sequence of consecutive free ports within a specified range. :param num_ports: Number of consecutive ports required. :param start_port: Starting port number (default: 1024). :param end_port: Ending port number (default: 65535). :param protocol: Protocol type ('tcp' only). :param max_attempts: Maximum number of retry attempts. :param retry_delay: Delay between retries (in seconds). :return: A tuple of (start_port, list_of_ports) if found, otherwise raise ValueError. """ # Validate input parameters if num_ports <= 0: raise ValueError("Number of ports must be greater than 0") if start_port < 0 or end_port > 65535: raise ValueError("Port range must be between 0 and 65535") if start_port > end_port: raise ValueError("Start port cannot be greater than end port") if (end_port - start_port + 1) < num_ports: raise ValueError(f"Port range ({start_port}-{end_port}) is too small to accommodate {num_ports} consecutive ports") # calculate the max start port max_start = end_port - num_ports + 1 for attempt in range(max_attempts): # Seeding with a combination of process ID and high-resolution timestamp pid = os.getpid() seed = pid + int(time.time() * 1_000_000) random.seed(seed) current = random.randint(start_port, max_start-1) while current <= max_start: # check current port if not is_port_available(current): current += 1 continue # check following n+1 port all_available = True for offset in range(1, num_ports): port_to_check = current + offset if not is_port_available(port_to_check): all_available = False current = port_to_check + 1 break # find available port and then lock these port if all_available: # check port again final_check = True for port in range(current, current + num_ports): if not is_port_available(port): final_check = False current = port + 1 break if final_check: ports = list(range(current, current + num_ports)) return current, ports # No port found in this attempt; sleep and retry. time.sleep(retry_delay) # No port found in all attempts; raise valueerror raise ValueError(f"No available ports found after {max_attempts} attempts. Range: {start_port}-{end_port}, Number of ports: {num_ports}") def load_local_config(model_dir): # load config.json from model_path cfg_path = os.path.join(model_dir, "config.json") with open(cfg_path, "r", encoding="utf-8") as f: data = json.load(f) data.pop("auto_map", None) # dtype rename if "torch_dtype" in data and "dtype" not in data: td = data["torch_dtype"] if isinstance(td, str): m = { "float16": torch.float16, "fp16": torch.float16, "bfloat16": torch.bfloat16, "bf16": torch.bfloat16, "float32": torch.float32, "fp32": torch.float32, } data["dtype"] = m.get(td.lower(), td) else: data["dtype"] = td # return with PretrainedConfig format. cfg = PretrainedConfig.from_dict(data) cfg.name_or_path = model_dir return cfg # When multiple workers run under DP/TP, concurrent AutoConfig calls can cause race conditions during config loading due to compilation overhead. # To avoid this, we override AutoConfig to directly load and convert config.json from the model path. def _patched(cls, name_or_path, *args, **kwargs): subfolder = kwargs.get("subfolder") local_dir = os.path.join(name_or_path, subfolder) if subfolder else name_or_path assert os.path.isdir(local_dir), f"Model path does not exist: {local_dir}" cfg_file = os.path.join(local_dir, "config.json") assert os.path.isfile(cfg_file), f"Config.json file does not exist: {cfg_file}" try: return load_local_config(local_dir) except Exception as e: # Fall back to original loader on local config failure. logger.warning(f"Failed to load config from local path '{local_dir}': {e}") logger.warning("Falling back to Hugging Face's original loader...") return _original_from_pretrained(name_or_path, *args, **kwargs) class SamplingParams: """ The parameters used for sampling a sequence. Parameters ---------- threshold : float The threshold used for threshold-based parallel decoding algorithm. cache : str The kv-cache type. Valid values include 'prefix', 'dual' and ''. temperature : float The temperature used for decoding tokens. early_stop : bool Whether to stop generating tokens after encountering an EOS. cont_weight : float This is used by IterSmooth algorithm. prefix_look : int This is used by vicinity KV-cache refresh algorithm. This determines the number of tokens before the decoding block that should recompute key and value states in every diffusion iteration. after_look : int This is used by vicinity KV-cache refresh algorithm. This determines the number of tokens after the decoding block that should recompute key and value states in every diffusion iteration. warmup_steps : int This is used by vicinity KV-cache refresh algorithm. This determines the number of steps at the beginning that we need to refresh key and value states of the entire sequence. enable_torch_compile : bool Whether to use torch compile for the model code. mask_id : int The mask ID eos_id : int The EOS ID """ def __init__(self, threshold=0.9, low_threshold=0.6, cache='dual', temperature=0., early_stop=True, cont_weight=0.3, prefix_look=16, after_look=16, warmup_steps=4, enable_torch_compile=True, mask_id=156895, eos_id=156892, parallel_decoding='threshold', use_credit=False, use_bd=True, max_length=4096, ep_size=1, prefilling_limit=256, mini_batch_size=1, batch_size=1, use_naive_batching=False): self.threshold = threshold self.low_threshold = low_threshold self.cache = cache self.temperature = temperature self.early_stop = early_stop self.cont_weight = cont_weight self.prefix_look = prefix_look self.after_look = after_look self.warmup_steps = warmup_steps self.mask_id = mask_id self.eos_id = eos_id self.enable_torch_compile = enable_torch_compile self.parallel_decoding = parallel_decoding self.use_credit = use_credit self.use_bd = use_bd self.max_length = max_length self.ep_size = ep_size self.prefilling_limit = prefilling_limit self.mini_batch_size = mini_batch_size self.batch_size = batch_size self.use_naive_batching = use_naive_batching def init_generator(model, sample_params, backend='vllm', max_length=4096): if sample_params.parallel_decoding == 'threshold': if sample_params.use_credit: decoder = CreditThresholdParallelDecoder(temperature=sample_params.temperature, threshold=sample_params.threshold, mask_id=sample_params.mask_id, eos_id=sample_params.eos_id) else: decoder = ThresholdParallelDecoder(temperature=sample_params.temperature, threshold=sample_params.threshold, mask_id=sample_params.mask_id, eos_id=sample_params.eos_id) else: decoder = HierarchyDecoder(temperature=sample_params.temperature, threshold=sample_params.threshold, low_threshold=sample_params.low_threshold, mask_id=sample_params.mask_id, eos_id=sample_params.eos_id) if sample_params.cache == 'prefix' or sample_params.cache == 'dual': cache_factory = KVCacheFactory(sample_params.cache, is_bd_model=sample_params.use_bd, backend=backend, max_length=max_length) else: cache_factory = None if not sample_params.use_bd: if cache_factory is not None and sample_params.cont_weight > 0: dllm = IterSmoothWithVicinityCacheDiffusionLLM(model, decoder, BlockIteratorFactory(), cache_factory=cache_factory, early_stop=sample_params.early_stop, cont_weight=sample_params.cont_weight, prefix_look=sample_params.prefix_look, after_look=sample_params.after_look, warmup_steps=sample_params.warmup_steps) elif cache_factory is not None: dllm = VicinityCacheDiffusionLLM(model, decoder, BlockIteratorFactory(), cache_factory=cache_factory, early_stop=sample_params.early_stop, prefix_look=sample_params.prefix_look, after_look=sample_params.after_look, warmup_steps=sample_params.warmup_steps) elif sample_params.cont_weight > 0: dllm = IterSmoothDiffusionLLM(model, decoder, BlockIteratorFactory(), cache_factory=None, early_stop=sample_params.early_stop, cont_weight=sample_params.cont_weight) else: dllm = BlockWiseDiffusionLLM(model, decoder, BlockIteratorFactory(), cache_factory=None, early_stop=sample_params.early_stop) else: dllm = BlockDiffusionLLM(model, decoder, BlockIteratorFactory(start_block_align=True, use_block_diffusion=True), cache_factory=cache_factory, early_stop=sample_params.early_stop, maximum_unroll=1, expected_tpf=15, backend=backend, prefilling_limit=sample_params.prefilling_limit, mini_batch_size=sample_params.mini_batch_size, use_naive_batching=sample_params.use_naive_batching) return dllm def generate(dllm, device, req_q, res_q): while True: data = req_q.get() if isinstance(data, str): assert data == 'stop' break else: input_ids, gen_len, block_len = data input_ids = input_ids.to(device) out = dllm.generate(input_ids, gen_length=gen_len, block_length=block_len) num_forwards = dllm.num_forwards if res_q is not None: res_q.put((out, num_forwards)) def sglang_llada2_server_process(model_path, sample_params, world_size, rank, gpu_id, q, res_q, master_port, error_q=None): try: _sglang_llada2_server_process(model_path, sample_params, world_size, rank, gpu_id, q, res_q, master_port) except Exception as e: error_msg = traceback.format_exc() logger.error(f"[SERVER PROCESS {rank} ERROR]: {error_msg}") if error_q is not None: try: error_q.put((e, error_msg)) except: logger.error(f"[ERROR_Q exception in {rank}].") def _sglang_llada2_server_process(model_path, sample_params, world_size, rank, gpu_id, q, res_q, master_port): torch.cuda.set_device(gpu_id) device = torch.device(gpu_id) logger.info(f'start v2 server. server port: {master_port}') os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = str(master_port) from sglang.srt import distributed distributed.init_distributed_environment(world_size, rank, 'env://', rank, 'nccl') distributed.initialize_model_parallel(world_size, sample_params.ep_size, 1, backend='nccl') from sglang.srt.server_args import ServerArgs from sglang.srt.layers.moe import initialize_moe_config from ..model.modeling_llada2_moe_sglang import LLaDA2SGLangLM from .diffusion_runner import ModelRunner from sglang.srt.layers.dp_attention import initialize_dp_attention # override AutoConfig to directly load and convert config.json from the model path AutoConfig.from_pretrained = classmethod(_patched) model_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) # ServerArgs call autoconfig in sglang/srt/utils/hf_transformers_utils.py. We use overrided autoconbfig here. server_args = ServerArgs(model_path=model_path, enable_dp_attention=True, trust_remote_code=True, tp_size=world_size, dp_size = 1, pp_size = 1, port=master_port+1, dist_init_addr="127.0.0.1:{}".format(master_port+2)) try: from sglang.srt.server_args import set_global_server_args_for_scheduler except ImportError: pass else: set_global_server_args_for_scheduler(server_args) initialize_dp_attention( server_args=server_args, model_config=model_config, ) initialize_moe_config(server_args) model = LLaDA2SGLangLM(config=model_config, expert_map_path='.').eval() torch.set_default_dtype(torch.bfloat16) model.load_weights(model_path, device=device) initialize_moe_config(server_args) model = model.to(device) max_length = sample_params.max_length prefill_lengths=[32*i for i in range(1, sample_params.prefilling_limit//32+1)] supported_batch_sizes = [2**i for i in range(int(np.log2(sample_params.mini_batch_size))+1)] model = ModelRunner(model, device, server_args=server_args, max_length=max_length, prefill_lengths=prefill_lengths, enable_compile=sample_params.enable_torch_compile, supported_batch_sizes=supported_batch_sizes, use_cross_block=sample_params.batch_size==1) dllm = init_generator(model, sample_params, backend='sglang', max_length=max_length) generate(dllm, model.device, req_q=q, res_q=res_q) def moe_server_process(model_path, sample_params, world_size, rank, gpu_id, q, res_q, master_port, error_q=None): try: _moe_server_process(model_path, sample_params, world_size, rank, gpu_id, q, res_q, master_port) except Exception as e: error_msg = traceback.format_exc() logger.error(f"[SERVER PROCESS {rank} ERROR]: {error_msg}") if error_q is not None: try: error_q.put((e, error_msg)) except: logger.error(f"[ERROR_Q exception in {rank}].") def _moe_server_process(model_path, sample_params, world_size, rank, gpu_id, q, res_q, master_port): torch.cuda.set_device(gpu_id) device = torch.device(gpu_id) logger.info(f'start MOE server. server port: {master_port}') os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = str(master_port) vllm_dist.init_distributed_environment(world_size, rank, 'env://', rank, 'nccl') vllm_dist.initialize_model_parallel(world_size, backend='nccl') # setup EP parallel_config = ParallelConfig(enable_expert_parallel = True) with set_current_vllm_config(VllmConfig(parallel_config = parallel_config)): model_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) model = FusedOlmoeForCausalLM(config=model_config).eval() model.load_weights(model_path, torch_dtype=torch.bfloat16) if world_size > 1: model.tensor_parallel(world_size) if sample_params.enable_torch_compile: model.forward = torch.compile(model.forward, mode='reduce-overhead', fullgraph=False, dynamic=True) model = model.to(device) dllm = init_generator(model, sample_params) generate(dllm, model.device, req_q=q, res_q=res_q) # TODO(zhengda) we should destroy the distributed environment. However, the function hangs if TP/EP is turned on. #vllm_dist.destroy_distributed_environment() def server_process(model_path, sample_params, world_size, rank, gpu_id, q, res_q, master_port, error_q=None): try: _server_process(model_path, sample_params, world_size, rank, gpu_id, q, res_q, master_port) except Exception as e: error_msg = traceback.format_exc() logger.error(f"[SERVER PROCESS {rank} ERROR]: {error_msg}") if error_q is not None: try: error_q.put((e, error_msg)) except: logger.error(f"[ERROR_Q exception in {rank}].") def _server_process(model_path, sample_params, world_size, rank, gpu_id, q, res_q, master_port): device = torch.device(gpu_id) os.environ['MASTER_ADDR'] = '127.0.0.1' os.environ['MASTER_PORT'] = str(master_port) dist.init_process_group(backend="nccl", rank=rank, world_size=world_size) config = AutoConfig.from_pretrained(model_path) config.flash_attention = True model = LLaDAModelLM.from_pretrained(model_path, torch_dtype=torch.bfloat16, config=config).eval() if world_size > 1: model.tensor_parallel(rank, world_size) if sample_params.enable_torch_compile: model.forward = torch.compile(model.forward, mode='reduce-overhead', fullgraph=False, dynamic=True) model = model.to(device) dllm = init_generator(model, sample_params) generate(dllm, model.device, req_q=q, res_q=res_q) dist.destroy_process_group() class ServerGroup: def __init__(self): self.procs = [] self.req_qs = [] self.res_q = None self.error_qs = [] def add_request(self, req): assert len(self.req_qs) != 0 and len(self.req_qs) == len(self.procs) for q in self.req_qs: q.put(req) def get_response(self, timeout=None): for i, eq in enumerate(self.error_qs): if not eq.empty(): (e, error_msg) = eq.get_nowait() print(f"Worker {i} failed: {error_msg}") raise e try: return self.res_q.get(timeout=timeout) except Empty: print(f"[FATAL ERROR] Response timeout - possible hang or crash.") raise RuntimeError("Inference timed out or worker crashed.") # def _handle_error(self, msg): # logger.error(f"[FATAL ERROR] {msg}") # self.stop_running() # raise RuntimeError(msg) def start_server(self, model_path, model_type, sample_params, server_port, gpus, backend): ctx = mp.get_context('spawn') assert len(self.procs) == 0, 'The server is already running.' procs = [] req_qs = [] error_qs = [] for i, gpu in enumerate(gpus): if i == 0: res_q = ctx.Queue() self.res_q = res_q else: res_q = None q = ctx.Queue() req_qs.append(q) error_q = ctx.Queue() error_qs.append(error_q) if backend=='sglang': assert model_type.startswith('llada2') p = ctx.Process(target=sglang_llada2_server_process, args=(model_path, sample_params, len(gpus), i, gpu, q, res_q, server_port, error_q)) elif model_type=='llada-moe': p = ctx.Process(target=moe_server_process, args=(model_path, sample_params, len(gpus), i, gpu, q, res_q, server_port, error_q)) else: p = ctx.Process(target=server_process, args=(model_path, sample_params, len(gpus), i, gpu, q, res_q, server_port, error_q)) p.daemon = True procs.append(p) p.start() self.procs = procs self.req_qs = req_qs self.error_qs = error_qs def is_running(self): return len(self.procs) != 0 def stop_running(self): for p in self.procs: p.kill() p.join() self.procs = [] self.req_qs = [] self.error_qs = [] self.res_q = None class ServerHandle: def __init__(self): self.groups = [] self.need_response = 0 def add_requests(self, reqs): prompts, gen_length, block_length = reqs self.need_response = 0 # assert len(self.groups) == prompts.shape[0], 'We cannot only use DP to support batch size > 1.' if len(self.groups) == prompts.shape[0]: for i, prompt in enumerate(prompts): self.groups[i].add_request((prompt.unsqueeze(0), gen_length, block_length)) self.need_response = len(prompts) else: partial_data = torch.chunk(prompts, len(self.groups), dim=0) for i in range(len(partial_data)): self.groups[i].add_request((partial_data[i], gen_length, block_length)) self.need_response = len(partial_data) def get_responses(self, timeout=None): res = [] for i in range(self.need_response): group = self.groups[i] res.append(group.get_response(timeout=timeout)) return res def start_server(self, model_path, model_type, sample_params, server_port, num_gpus, dp_size, tpep_size, backend): gpu = 0 assert num_gpus >= dp_size * tpep_size for i in range(dp_size): self.groups.append(ServerGroup()) gpus = [gpu + i for i in range(tpep_size)] logger.info(f'start server group on GPU {gpus}, server port: {server_port[i]}') self.groups[-1].start_server(model_path, model_type, sample_params, server_port[i], gpus, backend=backend) gpu = gpus[-1] + 1 def is_running(self): return len(self.groups) != 0 def stop_running(self): for group in self.groups: group.stop_running() self.groups = [] handle = ServerHandle() class DiffusionLLMServing: """ Serving dLLM inference. This is an experimental feature to enable serving in dInfer. This class creates multiple processes to enable dLLM inference in the background. A new request is sent to the background processes for model inference and the result is sent back to the main process. Parameters ---------- model : str The model path is_moe : bool Whether this is a MOE model. This leads to using different model code and inference code. sample_params : SamplingParams The parameters used in sampling. server_port: int The port for communication between the background process. num_gpus : int The number of GPUs used for parallel computation. max_retries : int The maximum number of retry attempts per DP group when searching for free ports. start_port : int The starting port number (inclusive) for the port search range. end_port : int The ending port number (inclusive) for the port search range. """ def __init__(self, model, model_type='llada2', sample_params=None, server_port=None, num_gpus=None, dp_size=None, tpep_size=None, backend='sglang', timeout = None, max_retries=3, start_port=30000, end_port=60000 ): if sample_params is None: sample_params = SamplingParams() self.sample_params = sample_params if num_gpus is None: num_gpus = torch.cuda.device_count() if dp_size is None: dp_size = 1 if tpep_size is None: tpep_size = num_gpus // dp_size assert dp_size * tpep_size <= num_gpus if server_port == None: # find a free port for each group server_port = self.port_selection(dp_size, start_port, end_port, max_retries=max_retries) else: # Keep compatibility with the existing implementation: assign ports starting from the current port + 1. server_port = [server_port+i for i in range(dp_size)] logger.info(f"find server port:{server_port}") if not handle.is_running(): handle.start_server(model, model_type, sample_params, server_port, num_gpus, dp_size, tpep_size, backend) self.num_forwards = 0 self.timeout = timeout def generate(self, prompts, gen_length=128, block_length=128): ''' Generate tokens with diffusion iterations. Parameters: ---------- prompts: Torch.Tensor A tensor of shape (b, L) that contains the input prompts. gen_length: int Generated answer length. block_length: int Block length, less than or equal to gen_length. If less than gen_length, it means using semi_autoregressive remasking. Returns ------- Torch.Tensor: A tensor of shape (b, L') that contains the prompt tokens and the generated tokens. The generation results of different lengths are padded with EOS. ''' prompts = prompts.cpu() handle.add_requests((prompts, gen_length, block_length)) rets = handle.get_responses(timeout=self.timeout) max_len = max([tensor.shape[1] for (tensor, _) in rets]) total_batch_size = sum([tensor.shape[0] for (tensor, _) in rets]) res = torch.zeros(total_batch_size, max_len, dtype=rets[0][0].dtype) res[:] = self.sample_params.eos_id sum_num_forwards = 0 p = 0 for i, (tensor, num_forwards) in enumerate(rets): sum_num_forwards = max(sum_num_forwards, num_forwards) out_len = int(tensor.shape[1]) res[p:p+len(tensor), :out_len] = tensor self.num_forwards = sum_num_forwards return res def stop_serving(self): """ Stop model serving. """ handle.stop_running() return None def port_selection(self, dp_size, start_port, end_port, max_retries=3): # select continuous port for each dp group ranging from start_port to end_port ports=[] for dp in range(dp_size): for retry in range(max_retries): try: # start with 7 ports: 1 nccl port; 6 continus port for server_args: 1 server_args default port + 5 extra ports if dp_attention is available # for simplity, we select 7 continuous_ports for each group port, dp_ports_list = find_continuous_ports(num_ports=7, start_port=start_port, end_port=end_port) ports.append(port) # record the first port for this group logger.info("Allocated port range [%d-%d] (%d ports) for DP group %d", dp_ports_list[0], dp_ports_list[-1], len(dp_ports_list), dp) break except ValueError as e: if retry < max_retries - 1: logger.warning("DP rank %d, attempt %d failed: %s. Retrying...",dp, retry + 1, e) time.sleep(1) else: raise RuntimeError( f"Failed to allocate ports for DP rank {dp} after {max_retries} attempts" ) from e return ports