diff --git a/python/sglang/srt/distributed/bootstrap.py b/python/sglang/srt/distributed/bootstrap.py index 6c91c1efc..f6efb4e5a 100644 --- a/python/sglang/srt/distributed/bootstrap.py +++ b/python/sglang/srt/distributed/bootstrap.py @@ -3,14 +3,12 @@ import os import time from typing import List, Optional -import msgspec import torch import torch.distributed as dist from sglang.srt.configs.model_config import ModelConfig from sglang.srt.distributed import ( get_default_distributed_backend, - get_pp_group, get_tp_group, get_world_group, init_distributed_environment, @@ -60,13 +58,6 @@ _is_cpu_arm64 = is_host_cpu_arm64() _TP_ALL_TO_ALL_WARMUP_BYTES_PER_PEER = 4 << 20 -class TorchDistributedResult(msgspec.Struct, frozen=True, kw_only=True): - tp_group: object - pp_group: object - attention_tp_group: object - pre_model_load_memory: float - - def init_torch_distributed( *, server_args: ServerArgs, @@ -140,10 +131,6 @@ def init_torch_distributed( distributed=get_world_group().world_size > 1 and not is_draft_worker, cpu_group=get_world_group().cpu_group, ) - tp_group = get_tp_group() - pp_group = get_pp_group() - attention_tp_group = get_parallel().attn_tp_group - # Check memory for tensor parallelism local_gpu_memory = get_available_gpu_memory(device, ps.gpu_id) if ps.tp_size > 1 and not is_draft_worker: @@ -156,12 +143,7 @@ def init_torch_distributed( f"Init torch distributed ends. elapsed={time.perf_counter() - tic:.2f} s, " f"mem usage={(before_avail_memory - local_gpu_memory):.2f} GB" ) - return TorchDistributedResult( - tp_group=tp_group, - pp_group=pp_group, - attention_tp_group=attention_tp_group, - pre_model_load_memory=pre_model_load_memory, - ) + return pre_model_load_memory def _resolve_backend(*, device: str) -> str: diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index a0cc3c22e..b0130f353 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -601,7 +601,6 @@ class Scheduler( and get_parallel().attn_tp_rank == 0 and get_parallel().attn_cp_rank == 0 ), - ps=self.ps, tp_group=self.tp_group, pp_group=self.pp_group, enable_hierarchical_cache=self.enable_hierarchical_cache, diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index f2484bb53..e0529c281 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -64,7 +64,6 @@ if TYPE_CHECKING: from sglang.srt.configs.model_config import ModelConfig from sglang.srt.distributed.parallel_state import GroupCoordinator - from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.managers.tp_worker import BaseTpWorker from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.base_spec_worker import HiCacheDraftPlan @@ -263,12 +262,14 @@ def build_kv_cache( attn_cp_cpu_group: ProcessGroup, enable_metrics: bool, enable_kv_cache_events: bool, - ps: ParallelState, tp_group: GroupCoordinator, pp_group: GroupCoordinator, enable_hierarchical_cache: bool, hicache_draft_plan: Optional[HiCacheDraftPlan] = None, ) -> KVCacheBuildResult: + # Built from the scheduler loop, outside any draft scope, so the context + # answers for the process this cache belongs to. + parallel = get_parallel() sliding_window_size: Optional[int] = None full_tokens_per_layer: Optional[int] = None swa_tokens_per_layer: Optional[int] = None @@ -369,10 +370,10 @@ def build_kv_cache( enable_session_radix_cache=get_memory().enable_session_radix_cache, enable_mamba_extra_buffer=get_exec().mamba.enable_mamba_extra_buffer, enable_mamba_extra_buffer_lazy=get_exec().mamba.enable_mamba_extra_buffer_lazy, - pp_rank=ps.pp_rank, - pp_size=ps.pp_size, - attn_cp_rank=ps.attn_cp_rank, - attn_cp_size=ps.attn_cp_size, + pp_rank=parallel.pp_rank, + pp_size=parallel.pp_size, + attn_cp_rank=parallel.attn_cp_rank, + attn_cp_size=parallel.attn_cp_size, chunked_prefill_size=effective_chunked_prefill_size, sliding_window_size=sliding_window_size, mtp_draft_device_pools=mtp_draft_device_pools, @@ -390,8 +391,8 @@ def build_kv_cache( effective_chunked_prefill_size=effective_chunked_prefill_size, tp_worker=tp_worker, model_config=model_config, - tp_size=ps.tp_size, - tp_rank=ps.tp_rank, + tp_size=parallel.tp_size, + tp_rank=parallel.tp_rank, tp_group=tp_group, ) with auto_size_hicache( diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 49e552c61..ef86199ba 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -206,7 +206,6 @@ def _pp_local_per_request_bytes( if TYPE_CHECKING: - from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.mem_cache.unified_memory_pool import ( UnifiedKVPool, UnifiedPoolBundle, @@ -260,7 +259,10 @@ class _PoolSizes(msgspec.Struct, frozen=True, kw_only=True): class KVCacheConfigurator: device: str gpu_id: int - ps: ParallelState + # Frozen at construction, not asked for later: this configurator is built + # inside the scope that describes a draft runner and used outside it. + attn_dp_size: int + pp_size: int pp_group: Any model: Any model_config: ModelConfig @@ -2324,7 +2326,7 @@ class KVCacheConfigurator: max_num_reqs = get_schedule().max_running_requests if max_num_reqs is not None: - requested_per_worker = max_num_reqs // self.ps.attn_dp_size + requested_per_worker = max_num_reqs // self.attn_dp_size max_num_reqs = min(requested_per_worker, token_capacity // 2) else: requested_per_worker = None @@ -2433,16 +2435,16 @@ class KVCacheConfigurator: # allocates its own [start_layer, end_layer) slice. Charge the largest # per-stage share so every rank derives the same pool without a collective. all_mamba_layers = config.mamba2_cache_params.layers - if self.ps.pp_size > 1 and all_mamba_layers: + if self.pp_size > 1 and all_mamba_layers: max_stage_mamba_layers = max( sum(1 for i in all_mamba_layers if start <= i < end) for start, end in ( get_pp_indices( self.model_config.num_hidden_layers, rank, - self.ps.pp_size, + self.pp_size, ) - for rank in range(self.ps.pp_size) + for rank in range(self.pp_size) ) ) else: @@ -2474,7 +2476,7 @@ class KVCacheConfigurator: replayssm_ring_per_req = int(replayssm_ring_per_req * pp_layer_scale) if replayssm_active and self.hybrid_kda_config is None: replay_req_slots = ( - get_schedule().max_running_requests // self.ps.attn_dp_size + 1 + get_schedule().max_running_requests // self.attn_dp_size + 1 ) replayssm_fixed_bytes = replayssm_ring_per_req * replay_req_slots replayssm_ring_per_slot = 0 @@ -2490,7 +2492,7 @@ class KVCacheConfigurator: get_context().override( "mamba_pool.per_dp_shard", max_mamba_cache_size=get_schedule().max_mamba_cache_size - // self.ps.attn_dp_size, + // self.attn_dp_size, ) # Reserve intermediate memory based on capped max_num_reqs (+1: the # pool's padding slot, see memory_pool.py). Skipped under replayssm @@ -2498,7 +2500,7 @@ class KVCacheConfigurator: if has_spec_dec and not replayssm_active: ratio = self._calculate_mamba_ratio() capped_reqs = min( - get_schedule().max_running_requests // self.ps.attn_dp_size, + get_schedule().max_running_requests // self.attn_dp_size, get_schedule().max_mamba_cache_size // ratio, ) intermediate_size = ( @@ -2515,7 +2517,7 @@ class KVCacheConfigurator: get_context().override( "mamba_pool.from_max_running_requests", max_mamba_cache_size=get_schedule().max_running_requests - // self.ps.attn_dp_size, + // self.attn_dp_size, ) # Reserve intermediate memory based on capped max_num_reqs (+1: the # pool's padding slot). Skipped under replayssm. @@ -2555,7 +2557,7 @@ class KVCacheConfigurator: # Intermediate memory is included in mamba_budget, subtract it # so the return value only has main_state subtracted from total capped_reqs = min( - get_schedule().max_running_requests // self.ps.attn_dp_size, + get_schedule().max_running_requests // self.attn_dp_size, get_schedule().max_mamba_cache_size // ratio, ) intermediate_size = per_req * (capped_reqs + 1) * D diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 7c61d494d..948b95487 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -421,7 +421,7 @@ class ModelRunner: import os logger.warning( - f"Context: {self.device=} {ps.gpu_id=} {os.environ.get('CUDA_VISIBLE_DEVICES')=} {ps.tp_rank=} {ps.tp_size=}" + f"Context: {self.device=} {ps.gpu_id=} {os.environ.get('CUDA_VISIBLE_DEVICES')=} {get_parallel().tp_rank=} {get_parallel().tp_size=}" ) raise @@ -447,7 +447,7 @@ class ModelRunner: # CPU offload set_offloader(create_offloader(dp_rank=get_parallel().dp_rank)) - self._weight_checker = WeightChecker(get_model=lambda: self.model, ps=self.ps) + self._weight_checker = WeightChecker(get_model=lambda: self.model) if envs.SGLANG_DETECT_SLOW_RANK.get(): slow_rank_detector.execute() @@ -503,9 +503,9 @@ class ModelRunner: ): return - join_effective_ep_size = get_parallel().ep_join_rank_offset + self.ps.tp_size + join_effective_ep_size = get_parallel().ep_join_rank_offset + self.tp_size dist.barrier(group=self.tp_group.cpu_group) - if self.ps.tp_rank == 0: + if self.tp_rank == 0: register_scale_cohort( get_parallel().ep_join_rank_offset, join_effective_ep_size, @@ -513,7 +513,7 @@ class ModelRunner: join_scale_process_group() get_context().override("elastic_ep.scale_join", ep_size=join_effective_ep_size) - global_ep_rank = self.ps.tp_rank + get_parallel().ep_join_rank_offset + global_ep_rank = self.tp_rank + get_parallel().ep_join_rank_offset broadcast_global_expert_location_metadata( model_config=self.model_config, moe_ep_rank=global_ep_rank, @@ -563,7 +563,7 @@ class ModelRunner: def init_weight_updater(self): self.weight_updater = WeightUpdater( - tp_rank=self.ps.tp_rank, + tp_rank=self.tp_rank, device=self.device, gpu_id=self.gpu_id, model_config=self.model_config, @@ -586,8 +586,8 @@ class ModelRunner: def init_weight_exporter(self): self.weight_exporter = WeightExporter( - tp_rank=self.ps.tp_rank, - tp_size=self.ps.tp_size, + tp_rank=self.tp_rank, + tp_size=self.tp_size, gpu_id=self.gpu_id, get_model_path=lambda: self.model_config.model_path, get_model=lambda: self.model, @@ -596,7 +596,7 @@ class ModelRunner: def init_remote_instance_weight_transporter(self): self.remote_instance_weight_transporter = RemoteInstanceWeightTransporter( get_model=lambda: self.model, - tp_rank=self.ps.tp_rank, + tp_rank=get_parallel().tp_rank, gpu_id=self.gpu_id, ) @@ -610,10 +610,13 @@ class ModelRunner: ) def init_kv_cache_configurator(self): + # The replica count that shares this KV budget is the deployment's, not + # the one a draft scope reports; the pool is allocated outside any scope. self.kv_cache_configurator = KVCacheConfigurator( device=self.device, gpu_id=self.gpu_id, - ps=self.ps, + attn_dp_size=get_parallel().attn_dp_size, + pp_size=self.pp_size, pp_group=self.pp_group, model=self.model, model_config=self.model_config, @@ -647,8 +650,8 @@ class ModelRunner: from sglang.srt.model_executor.mindspore_runner import init_ms_distributed init_ms_distributed( - world_size=self.ps.tp_size * get_parallel().pp_size, - rank=self.ps.tp_size * get_parallel().pp_rank + self.ps.tp_rank, + world_size=self.tp_size * get_parallel().pp_size, + rank=self.tp_size * get_parallel().pp_rank + self.tp_rank, local_rank=self.gpu_id, port=self.dist_port, ) @@ -718,7 +721,7 @@ class ModelRunner: moe_ep_rank=expert_rank, ) ) - if self.ps.tp_rank == 0 and envs.SGLANG_LOG_EXPERT_LOCATION_METADATA.get(): + if self.tp_rank == 0 and envs.SGLANG_LOG_EXPERT_LOCATION_METADATA.get(): logger.info( "Initial expert_location_metadata:\n%s", format_expert_location_layout(get_global_expert_location_metadata()), @@ -779,7 +782,7 @@ class ModelRunner: def maybe_apply_post_load_model_transforms(self): supports_torch_tp = getattr(self.model, "supports_torch_tp", False) - if self.ps.tp_size > 1 and supports_torch_tp: + if self.tp_size > 1 and supports_torch_tp: self.apply_torch_tp() def maybe_init_lora_manager(self): @@ -1015,7 +1018,7 @@ class ModelRunner: def post_capture_elastic_ep_recover(self): join_process_groups() - global_ep_rank = self.ps.tp_rank + get_parallel().ep_join_rank_offset + global_ep_rank = self.tp_rank + get_parallel().ep_join_rank_offset broadcast_global_expert_location_metadata( model_config=self.model_config, moe_ep_rank=global_ep_rank, @@ -1162,13 +1165,13 @@ class ModelRunner: def check_quantized_moe_compatibility(self): check_quantized_moe_compatibility( model_config=self.model_config, - tp_size=self.ps.tp_size, + tp_size=self.tp_size, moe_ep_size=get_parallel().moe_ep_size, moe_dp_size=get_parallel().moe_dp_size, ) def init_torch_distributed(self): - result = bootstrap.init_torch_distributed( + self.pre_model_load_memory = bootstrap.init_torch_distributed( server_args=self.server_args, model_config=self.model_config, device=self.device, @@ -1177,10 +1180,23 @@ class ModelRunner: is_draft_worker=self.is_draft_worker, local_omp_cpuid=self.local_omp_cpuid if self.device == "cpu" else None, ) - self.tp_group = result.tp_group - self.pp_group = result.pp_group - self.attention_tp_group = result.attention_tp_group - self.pre_model_load_memory = result.pre_model_load_memory + # Read once, here: a draft runner is constructed inside the scope that + # states its topology and used outside it, so what it holds has to be + # the placement it was built for rather than whatever the context + # answers later. Groups and widths alike -- a runner asked about its own + # shape after the scope has closed must still describe itself. + parallel = get_parallel() + self.tp_group = parallel.tp_group + self.pp_group = parallel.pp_group + self.attention_tp_group = parallel.attn_tp_group + self.tp_rank = parallel.tp_rank + self.tp_size = parallel.tp_size + self.dp_size = parallel.dp_size + self.attn_dp_size = parallel.attn_dp_size + self.pp_rank = parallel.pp_rank + self.pp_size = parallel.pp_size + self.attn_cp_rank = parallel.attn_cp_rank + self.attn_cp_size = parallel.attn_cp_size def init_shared_mooncake_transfer_engine(self): maybe_init_shared_mooncake_transfer_engine(gpu_id=self.gpu_id) @@ -1204,7 +1220,7 @@ class ModelRunner: self.load_config = build_load_config( server_args=self.server_args, load_format=draft_load_format, - tp_rank=self.ps.tp_rank, + tp_rank=self.tp_rank, remote_instance_weight_transporter_engine=self.remote_instance_weight_transporter.engine, remote_instance_weight_transporter_session_id=self.remote_instance_weight_transporter.session_id, draft_model_idx=self.draft_model_idx, @@ -1217,11 +1233,11 @@ class ModelRunner: ) if self.device == "cpu": self.model_config = adjust_config_with_unaligned_cpu_tp( - self.model_config, self.load_config, self.ps.tp_size + self.model_config, self.load_config, self.tp_size ) maybe_trigger_remote_instance_nccl_send_group( - tp_rank=self.ps.tp_rank, + tp_rank=self.tp_rank, load_format=draft_load_format, ) @@ -1297,8 +1313,8 @@ class ModelRunner: model=self.model, spec_algorithm=self.spec_algorithm, is_draft_worker=self.is_draft_worker, - tp_size=self.ps.tp_size, - tp_rank=self.ps.tp_rank, + tp_size=self.tp_size, + tp_rank=self.tp_rank, pp_rank=get_parallel().pp_rank, ) @@ -1316,7 +1332,7 @@ class ModelRunner: if self.startup_weight_load is None: dist_barrier_after_load( elastic_ep_backend=get_exec().moe.elastic_ep_backend, - tp_rank=self.ps.tp_rank, + tp_rank=self.tp_rank, is_ep_joiner=get_exec().moe.is_ep_joiner, ) @@ -1339,7 +1355,7 @@ class ModelRunner: self.startup_weight_load.finalize() dist_barrier_after_load( elastic_ep_backend=get_exec().moe.elastic_ep_backend, - tp_rank=self.ps.tp_rank, + tp_rank=self.tp_rank, is_ep_joiner=get_exec().moe.is_ep_joiner, ) self.startup_weight_load = None @@ -1367,8 +1383,8 @@ class ModelRunner: dtype=self.dtype, server_args=self.server_args, lora_backend=get_lora().lora_backend, - tp_size=self.ps.tp_size, - tp_rank=self.ps.tp_rank, + tp_size=self.tp_size, + tp_rank=self.tp_rank, max_lora_rank=get_lora().max_lora_rank, target_modules=get_lora().lora_target_modules, lora_paths=get_lora().lora_paths, @@ -1583,15 +1599,16 @@ class ModelRunner: # rather than spawning additional processes, so dp_size must not be # multiplied into the process count here (unlike regular DP, where # dp_size * tp_size * pp_size is the true worker count). - dp_size = 1 if get_parallel().enable_dp_attention else self.ps.dp_size + parallel = get_parallel() + dp_size = 1 if parallel.enable_dp_attention else parallel.dp_size self.local_omp_cpuid = numa_utils.init_threads_binding( numa_index=self.gpu_id, - world_size=dp_size * self.ps.tp_size * get_parallel().pp_size, + world_size=dp_size * parallel.tp_size * parallel.pp_size, ) def apply_torch_tp(self): model_parallel.apply_torch_tp( - model=self.model, device=self.device, tp_size=self.ps.tp_size + model=self.model, device=self.device, tp_size=self.tp_size ) def update_decode_attn_backend(self, stream_idx: int): @@ -1717,7 +1734,7 @@ class ModelRunner: if self.msprobe_debugger is not None: rank_id = ( self.gpu_id - if self.ps.attn_dp_size is not None and self.ps.attn_dp_size > 1 + if self.attn_dp_size is not None and self.attn_dp_size > 1 else None ) self.msprobe_debugger.start(model=self.model, rank_id=rank_id) @@ -2098,7 +2115,7 @@ class ModelRunner: set_global_expert_location_metadata(new_metadata, allow_overwrite=True) def _elastic_global_rank(self) -> int: - return self.ps.tp_rank + get_parallel().ep_join_rank_offset + return self.tp_rank + get_parallel().ep_join_rank_offset def _rearm_eplb_after_elastic_scale(self) -> None: if self.eplb_manager is None: @@ -2120,7 +2137,7 @@ class ModelRunner: self._rearm_eplb_after_elastic_scale() def _report_elastic_scale_failure(self, error: str, effective_size: int) -> None: - if self.ps.tp_rank != 0 or get_exec().moe.is_ep_scale_joiner: + if self.tp_rank != 0 or get_exec().moe.is_ep_scale_joiner: return from sglang.srt.managers.io_struct import ElasticScaleUpdateReq @@ -2131,7 +2148,7 @@ class ModelRunner: ) def _elastic_scale_ready_barrier(self, target_size: int, log_tag: str) -> None: - if self.ps.tp_rank == 0: + if self.tp_rank == 0: logger.debug( "[Elastic EP][scale] %s entering post-scale WORLD barrier " "(target_ep_size=%d)", @@ -2139,7 +2156,7 @@ class ModelRunner: target_size, ) dist.barrier(group=dist.group.WORLD) - if self.ps.tp_rank == 0: + if self.tp_rank == 0: logger.debug( "[Elastic EP][scale] %s passed post-scale WORLD barrier " "(target_ep_size=%d)", @@ -2204,7 +2221,7 @@ class ModelRunner: ElasticEPStateManager.commit_scale() self._rearm_eplb_after_elastic_scale() - if self.ps.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner: + if self.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner: from sglang.srt.managers.io_struct import ElasticScaleUpdateReq self._pending_elastic_scale_update = ElasticScaleUpdateReq( @@ -2237,7 +2254,7 @@ class ModelRunner: ) ElasticEPStateManager.fail_recovery(error) self._report_elastic_scale_failure(error, effective_size) - if self.ps.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner: + if self.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner: logger.error("[Elastic EP] %s", error) return @@ -2263,7 +2280,7 @@ class ModelRunner: ElasticEPStateManager.fail_scale(error) self._reset_eplb_after_elastic_scale_failure() self._report_elastic_scale_failure(error, effective_size) - if self.ps.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner: + if self.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner: logger.error("[Elastic EP] %s", error) return @@ -2279,7 +2296,7 @@ class ModelRunner: ElasticEPStateManager.fail_scale(error) self._reset_eplb_after_elastic_scale_failure() self._report_elastic_scale_failure(error, effective_size) - if self.ps.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner: + if self.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner: logger.error("[Elastic EP] %s", error) return if not ElasticEPStateManager.begin_scale(): diff --git a/python/sglang/srt/utils/weight_checker.py b/python/sglang/srt/utils/weight_checker.py index 6e0b5581e..4b554c6a8 100644 --- a/python/sglang/srt/utils/weight_checker.py +++ b/python/sglang/srt/utils/weight_checker.py @@ -8,6 +8,7 @@ import torch.distributed as dist from pydantic import BaseModel, ConfigDict from sglang.srt.managers.mm_utils import tensor_hash +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils.weight_checker_comparator import ( CHUNK_NUMEL, ComparableWeight, @@ -67,9 +68,23 @@ def _is_non_persistent_buffer_name(name: str) -> bool: class WeightChecker: - def __init__(self, *, get_model: Callable[[], Any], ps: Any): + def __init__(self, *, get_model: Callable[[], Any]): self._get_model = get_model - self._ps = ps + # A check is served on demand from the scheduler loop, which is outside + # the scope that describes a draft runner. The report has to name the + # runner it was built for, so the placement is read here, at + # construction, rather than asked for when the request arrives. + parallel = get_parallel() + self._placement = ParallelismInfo( + tp_rank=parallel.tp_rank, + tp_size=parallel.tp_size, + dp_rank=parallel.dp_rank if parallel.dp_rank is not None else 0, + dp_size=parallel.attn_dp_size, + pp_rank=parallel.pp_rank, + pp_size=parallel.pp_size, + rank=0, + size=1, + ) self._snapshot_tensors = None def handle(self, action: str, allow_quant_error: bool = False) -> Optional[Dict]: @@ -161,16 +176,14 @@ class WeightChecker: return info.model_dump() def _parallelism_info(self) -> ParallelismInfo: - ps = self._ps - return ParallelismInfo( - tp_rank=ps.tp_rank, - tp_size=ps.tp_size, - dp_rank=ps.dp_rank if ps.dp_rank is not None else 0, - dp_size=ps.attn_dp_size, - pp_rank=ps.pp_rank, - pp_size=ps.pp_size, - rank=dist.get_rank() if dist.is_initialized() else 0, - size=dist.get_world_size() if dist.is_initialized() else 1, + # The WORLD position is asked for now rather than frozen: unlike the + # runner's placement it is a property of the process, and an elastic + # scale-up moves it. + return self._placement.model_copy( + update={ + "rank": dist.get_rank() if dist.is_initialized() else 0, + "size": dist.get_world_size() if dist.is_initialized() else 1, + } ) def _model_state(self): diff --git a/test/registered/layers/mamba/test_mamba2_mixer.py b/test/registered/layers/mamba/test_mamba2_mixer.py index 7ac76d0c4..771536a40 100644 --- a/test/registered/layers/mamba/test_mamba2_mixer.py +++ b/test/registered/layers/mamba/test_mamba2_mixer.py @@ -129,8 +129,27 @@ def mixer2_gated_norm_tensor_parallel( ) mixer.weight.weight_loader(mixer.weight, weight) - # m2 reads tp via get_parallel().tp_size/rank — force it through the context. - with get_parallel().override(tp_size=1, tp_rank=0): + # m2 reads tp via get_parallel().tp_size/rank — state a single-rank topology + # through the context. Every width that follows from `tp_size` is named: + # narrowing one leaf and leaving the quotients behind describes no layout. + with get_parallel().override( + tp_size=1, + tp_rank=0, + tp_group=None, + attn_tp_size=1, + attn_tp_rank=0, + attn_tp_group=None, + attn_dp_size=1, + attn_dp_rank=0, + attn_cp_size=1, + attn_cp_rank=0, + attn_cp_group=None, + moe_ep_size=1, + moe_ep_rank=0, + moe_ep_group=None, + moe_dp_size=1, + moe_tp_size=1, + ): # create gated-norm without TP to compute reference mixer_single_gpu = m2.Mixer2RMSNormGated( full_hidden_size=hidden_size, diff --git a/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py b/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py index 277a38d5b..91a286724 100644 --- a/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py +++ b/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py @@ -84,7 +84,7 @@ def _stub( stub._max_mamba_cache_size = max_mamba_cache_size stub._disable_radix_cache = disable_radix_cache stub.max_total_num_tokens = max_total_num_tokens - stub.ps = SimpleNamespace(attn_dp_size=dp_size) + stub.attn_dp_size = dp_size return stub @@ -94,7 +94,7 @@ def _hybrid_stub_for_initialize( """A stub carrying what the real initialize() reads (hybrid path).""" stub = MlxModelRunnerStub.__new__(MlxModelRunnerStub) stub._mlx_pool_size = pool - stub.ps = SimpleNamespace(attn_dp_size=1) + stub.attn_dp_size = 1 stub.device = "cpu" # read by init_ngram_embedding_manager # Evaluated as a call argument in init_ngram_embedding_manager before # the use_ngram_embedding short-circuit; never read. diff --git a/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py b/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py index c3789e787..2094e8c0c 100644 --- a/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py +++ b/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py @@ -91,7 +91,8 @@ def _make_prefill_aware_swa_runner( kv_cache_dtype=torch.float16, kv_cache_dtype_str="auto", page_size=1, - ps=SimpleNamespace(attn_cp_size=1, tp_size=1), + attn_cp_size=1, + tp_size=1, is_draft_worker=False, server_args=server_args, attention_chunk_size=None, diff --git a/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py b/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py index 1816cd510..0c3ee44f7 100644 --- a/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py +++ b/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py @@ -671,7 +671,7 @@ class TestModelRunnerStartupWeightLoadOwnership(CustomTestCase): elastic_ep_backend=None, is_ep_joiner=False, ) - runner.ps = SimpleNamespace(tp_rank=0) + runner.tp_rank = 0 return runner def test_start_delegates_to_the_manager(self): diff --git a/test/registered/unit/spec/test_dflash_logits.py b/test/registered/unit/spec/test_dflash_logits.py index df343070f..9ebbc9c2e 100644 --- a/test/registered/unit/spec/test_dflash_logits.py +++ b/test/registered/unit/spec/test_dflash_logits.py @@ -248,7 +248,7 @@ def test_worker_folds_a_gate_admitted_quantized_selector_head(monkeypatch): worker = SimpleNamespace( block_size=8, selector=object(), - ps=SimpleNamespace(tp_rank=0), + model_runner=SimpleNamespace(tp_rank=0), draft_model=SimpleNamespace(lm_head=None), device="cpu", _selector_sampling_enabled=True, @@ -282,7 +282,7 @@ def test_worker_warns_once_when_selector_sampling_is_disabled(monkeypatch): selector=object(), _selector_sampling_enabled=False, _warned_sampling_fallback=False, - ps=SimpleNamespace(tp_rank=0), + model_runner=SimpleNamespace(tp_rank=0), ) batch = SimpleNamespace(sampling_info=SimpleNamespace(is_all_greedy=False)) diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 6c50ff3c1..00c43ba70 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -2781,17 +2781,154 @@ class TestWhoAnswersDuringADraftScope(CustomTestCase): def test_a_report_built_for_a_runner_follows_that_runner(self): """A weight check is an on-demand request served from the scheduler loop, so it runs outside the scope that describes a draft runner. Its - report has to name the runner it was built for, which is why it holds - a record instead of asking the context.""" - from sglang.srt.distributed.parallel_state_wrapper import ParallelState + report has to name the runner it was built for, which is why it reads + the placement once, where it is built, instead of asking again when the + request arrives.""" + from sglang.srt.distributed import parallel_state from sglang.srt.utils.weight_checker import WeightChecker - draft = ParallelState.trivial(pp_rank=0, pp_size=1) - checker = WeightChecker(get_model=lambda: None, ps=draft) self._two_stage_pipeline() + self.assertEqual(get_parallel().pp_size, 2) + group = self._single_member_group() + with patch.object(parallel_state, "_PP", group): + with parallel_state.patch_pipeline_parallel_group(group): + checker = WeightChecker(get_model=lambda: None) + + # The scope has closed and the context answers the target's shape again. + self.assertEqual(get_parallel().pp_size, 2) info = checker._parallelism_info() self.assertEqual((info.pp_rank, info.pp_size), (0, 1)) +class TestNothingReadsThePlacementBeforeItIsFrozen(CustomTestCase): + """`ModelRunner.__init__` freezes its placement partway through. + + A method called before that point reads an attribute that does not exist + yet, and only on the configuration that reaches it -- a remote weight + transporter, a NUMA binding -- so the suites say nothing and a GPU job is + where it surfaces. The order is what makes it wrong, so the order is what + is checked. + """ + + def _model_runner(self): + import ast as _ast + + source = (_SRT / "model_executor" / "model_runner.py").read_text() + for node in _ast.parse(source).body: + if isinstance(node, _ast.ClassDef) and node.name == "ModelRunner": + return {m.name: m for m in node.body if isinstance(m, _ast.FunctionDef)} + raise AssertionError("ModelRunner not found") + + def _frozen_names(self, methods): + import ast as _ast + + return { + target.attr + for node in _ast.walk(methods["init_torch_distributed"]) + if isinstance(node, _ast.Assign) + for target in node.targets + if isinstance(target, _ast.Attribute) + and isinstance(target.value, _ast.Name) + and target.value.id == "self" + } + + def _reads(self, methods, fn, frozen, depth=0): + import ast as _ast + + found = set() + for node in _ast.walk(fn): + if ( + isinstance(node, _ast.Attribute) + and isinstance(node.value, _ast.Name) + and node.value.id == "self" + and isinstance(node.ctx, _ast.Load) + and node.attr in frozen + ): + found.add(node.attr) + if ( + depth < 2 + and isinstance(node, _ast.Call) + and isinstance(node.func, _ast.Attribute) + and isinstance(node.func.value, _ast.Name) + and node.func.value.id == "self" + and node.func.attr in methods + and node.func.attr != "init_torch_distributed" + ): + found |= self._reads( + methods, methods[node.func.attr], frozen, depth + 1 + ) + return found + + def _calls_before_the_freeze(self, methods): + import ast as _ast + + calls = [] + for statement in methods["__init__"].body: + for node in _ast.walk(statement): + if ( + isinstance(node, _ast.Call) + and isinstance(node.func, _ast.Attribute) + and isinstance(node.func.value, _ast.Name) + and node.func.value.id == "self" + ): + calls.append((node.lineno, node.func.attr)) + calls.sort() + names = [name for _, name in calls] + self.assertIn( + "init_torch_distributed", + names, + "the freeze moved; this census is keyed on where it happens", + ) + return calls[: names.index("init_torch_distributed")] + + def test_no_method_called_before_the_freeze_reads_what_it_freezes(self): + methods = self._model_runner() + frozen = self._frozen_names(methods) + self.assertGreater(len(frozen), 5, "found no frozen names; census is broken") + offenders = [] + for lineno, name in self._calls_before_the_freeze(methods): + fn = methods.get(name) + if fn is None: + continue + read = self._reads(methods, fn, frozen) + if read: + offenders.append( + f"__init__:{lineno} self.{name}() reads {sorted(read)}" + ) + self.assertEqual( + offenders, + [], + "these run before init_torch_distributed and read what it sets; " + "ask get_parallel() there, or move the call after the freeze:\n " + + "\n ".join(offenders), + ) + + def test_the_census_would_notice_one(self): + """The positive control: the walk has to find a read that is there.""" + import ast as _ast + import textwrap + + methods = { + m.name: m + for m in _ast.parse( + textwrap.dedent( + """ + class R: + def init_torch_distributed(self): + self.tp_rank = 0 + + def early(self): + return self.tp_rank + """ + ) + ) + .body[0] + .body + } + frozen = self._frozen_names(methods) + self.assertEqual(frozen, {"tp_rank"}) + self.assertEqual(self._reads(methods, methods["early"], frozen), {"tp_rank"}) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/utils/test_weight_checker.py b/test/registered/unit/utils/test_weight_checker.py index 702923f51..5db02aff7 100644 --- a/test/registered/unit/utils/test_weight_checker.py +++ b/test/registered/unit/utils/test_weight_checker.py @@ -44,7 +44,7 @@ from sglang.srt.utils.weight_checker_comparator import ( RawComparable, ) from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci -from sglang.test.test_utils import CustomTestCase +from sglang.test.test_utils import CustomTestCase, enter_scope, published_topology register_amd_ci(est_time=30, suite="stage-b-test-1-gpu-small-amd") register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small") @@ -551,7 +551,8 @@ class _WeightCheckerTestBase(CustomTestCase): torch.manual_seed(0) self.model = _TinyModel().cuda() runner = _FakeModelRunner(self.model) - self.checker = WeightChecker(get_model=lambda: runner.model, ps=runner.ps) + enter_scope(self, published_topology()) + self.checker = WeightChecker(get_model=lambda: runner.model) class TestSnapshot(_WeightCheckerTestBase): @@ -760,9 +761,16 @@ class _ChecksumTestBase(CustomTestCase): pp_rank=0, pp_size=1, ) - self.checker = WeightChecker( - get_model=lambda: self.runner.model, ps=self.runner.ps + enter_scope( + self, + published_topology( + tp_size=4, + dp_size=2, + enable_dp_attention=True, + ranks={"world_rank": 2, "dp_rank": 1}, + ), ) + self.checker = WeightChecker(get_model=lambda: self.runner.model) class TestComputeChecksum(_ChecksumTestBase):