From 1b56a886bb92f3d6a753c6b252d5e4f5872a7e3a Mon Sep 17 00:00:00 2001 From: Zhongdongming Dai Date: Mon, 26 Jan 2026 19:04:25 -0800 Subject: [PATCH] [chore]: improve time tracing of model loading process (#15426) Co-authored-by: Michael Shin Co-authored-by: ishandhanani <82981111+ishandhanani@users.noreply.github.com> --- python/sglang/srt/model_executor/model_runner.py | 11 ++++++++++- python/sglang/srt/model_loader/loader.py | 11 +++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 275f5164e..4490956c2 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -700,6 +700,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): ) def init_torch_distributed(self): + tic = time.perf_counter() logger.info("Init torch distributed begin.") try: @@ -807,11 +808,13 @@ class ModelRunner(ModelRunnerKVCacheMixin): logger.warning(msg) logger.info( - f"Init torch distributed ends. mem usage={(before_avail_memory - local_gpu_memory):.2f} GB" + 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 def load_model(self): + tic_total = time.perf_counter() before_avail_memory = get_available_gpu_memory(self.device, self.gpu_id) logger.info( f"Load weight begin. avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB" @@ -957,6 +960,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.weight_load_mem_usage = before_avail_memory - after_avail_memory logger.info( f"Load weight end. " + f"elapsed={time.perf_counter() - tic_total:.2f} s, " f"type={type(self.model).__name__}, " f"dtype={self.dtype}, " f"avail mem={after_avail_memory:.2f} GB, " @@ -1604,6 +1608,8 @@ class ModelRunner(ModelRunnerKVCacheMixin): def init_attention_backend(self): """Init attention kernel backend.""" + tic = time.perf_counter() + logger.info("Init attention backend begin.") if self.server_args.enable_pdmux: self.attn_backend = self._get_attention_backend(init_new_workspace=True) self.decode_attn_backend_group = [] @@ -1614,6 +1620,9 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.attn_backend = TboAttnBackend.init_new(self._get_attention_backend) else: self.attn_backend = self._get_attention_backend() + logger.info( + f"Init attention backend end. elapsed={time.perf_counter() - tic:.2f} s" + ) def _get_attention_backend(self, init_new_workspace: bool = False): """Init attention kernel backend.""" diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index cb6aa0d78..1b6658c6a 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -335,6 +335,9 @@ class DefaultModelLoader(BaseModelLoader): model_config=model_config, ) + counter_before_loading_weights: float = 0.0 + counter_after_loading_weights: float = 0.0 + def __init__(self, load_config: LoadConfig): super().__init__(load_config) extra_config = load_config.model_loader_extra_config @@ -544,6 +547,9 @@ class DefaultModelLoader(BaseModelLoader): filtered_weights.append((source.prefix + new_name, tensor)) return tuple(filtered_weights) + if self.counter_before_loading_weights == 0.0: + logger.info("Beginning to load weights") + self.counter_before_loading_weights = time.perf_counter() # Apply the prefix. return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator) @@ -663,6 +669,11 @@ class DefaultModelLoader(BaseModelLoader): model, self._get_all_weights(model_config, model), target_device ) + self.counter_after_loading_weights = time.perf_counter() + logger.info( + "Loading weights took %.2f seconds", + self.counter_after_loading_weights - self.counter_before_loading_weights, + ) return model.eval() @staticmethod