diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 605089c76..9b615cce8 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -391,8 +391,8 @@ class ModelRunner(ModelRunnerKVCacheMixin): # Initialize MooncakeTransferEngine self.init_shared_mooncake_transfer_engine() - # Get memory before model loading - min_per_gpu_memory = self.init_torch_distributed() + # Get available memory before model loading + pre_model_load_memory = self.init_torch_distributed() # Init forward stream for overlap schedule self.forward_stream = torch.get_device_module(self.device).Stream() @@ -413,7 +413,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): deep_gemm_wrapper.update_deep_gemm_config(gpu_id, server_args) # Initialize the model runner - self.initialize(min_per_gpu_memory) + self.initialize(pre_model_load_memory) self.check_quantized_moe_compatibility() if self.is_multimodal: @@ -447,7 +447,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): port=self.dist_port, ) - def initialize(self, min_per_gpu_memory: float): + def initialize(self, pre_model_load_memory: float): server_args = self.server_args self.memory_saver_adapter = TorchMemorySaverAdapter.create( @@ -600,7 +600,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.configure_kv_cache_dtype() # Init memory pool and attention backends - self.init_memory_pool(min_per_gpu_memory) + self.init_memory_pool(pre_model_load_memory) # Init max running requests self.max_running_requests = min( @@ -826,7 +826,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): if is_npu(): register_sgl_tp_rank(self.gpu_id) - min_per_gpu_memory = get_available_gpu_memory( + pre_model_load_memory = get_available_gpu_memory( self.device, self.gpu_id, distributed=get_world_group().world_size > 1, @@ -839,9 +839,9 @@ class ModelRunner(ModelRunnerKVCacheMixin): # Check memory for tensor parallelism local_gpu_memory = get_available_gpu_memory(self.device, self.gpu_id) if self.tp_size > 1 and not self.is_draft_worker: - if min_per_gpu_memory < local_gpu_memory * 0.9: + if pre_model_load_memory < local_gpu_memory * 0.9: msg = "The memory capacity is unbalanced. Some GPUs may be occupied by other processes. " - msg += f"{min_per_gpu_memory=}, {local_gpu_memory=}, {local_gpu_memory * 0.9=}" + msg += f"{pre_model_load_memory=}, {local_gpu_memory=}, {local_gpu_memory * 0.9=}" if envs.SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK.get(): raise RuntimeError(msg) else: @@ -851,7 +851,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): f"Init torch distributed ends. elapsed={time.perf_counter() - tic:.2f} s, " f"mem usage={(before_avail_memory - local_gpu_memory):.2f} GB" ) - return min_per_gpu_memory + return pre_model_load_memory def init_shared_mooncake_transfer_engine(self): """ diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index 5a3e31a37..60da20313 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -113,8 +113,8 @@ class ModelRunnerKVCacheMixin: ) return cell_size - def profile_max_num_token(self: ModelRunner, total_gpu_memory: int): - available_gpu_memory = get_available_gpu_memory( + def profile_max_num_token(self: ModelRunner, pre_model_load_memory: int): + post_model_load_memory = get_available_gpu_memory( self.device, self.gpu_id, distributed=get_world_group().world_size > 1, @@ -140,7 +140,7 @@ class ModelRunnerKVCacheMixin: cell_size = self.get_cell_size_per_token(num_layers) - rest_memory = available_gpu_memory - total_gpu_memory * ( + rest_memory = post_model_load_memory - pre_model_load_memory * ( 1 - self.mem_fraction_static ) if self.mambaish_config is not None: @@ -675,10 +675,10 @@ class ModelRunnerKVCacheMixin: self.token_to_kv_pool_allocator.full_to_swa_index_mapping ) - def init_memory_pool(self: ModelRunner, total_gpu_memory: int): + def init_memory_pool(self: ModelRunner, pre_model_load_memory: int): max_num_reqs = self.server_args.max_running_requests max_total_tokens = self.server_args.max_total_tokens - self.max_total_num_tokens = self.profile_max_num_token(total_gpu_memory) + self.max_total_num_tokens = self.profile_max_num_token(pre_model_load_memory) if max_num_reqs is None: max_num_reqs = min(