diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 41044f06c..35a9080ec 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -182,7 +182,10 @@ from sglang.srt.utils.offloader import ( get_offloader, set_offloader, ) -from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions +from sglang.srt.utils.patch_torch import ( + monkey_patch_torch_reductions, + register_sgl_tp_rank, +) from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter from sglang.srt.utils.weight_checker import WeightChecker from sglang.srt.weight_sync.tensor_bucket import ( @@ -759,6 +762,8 @@ class ModelRunner: server_args=self.server_args, model_config=self.model_config, ) + if is_npu(): + register_sgl_tp_rank(self.gpu_id) min_per_gpu_memory = get_available_gpu_memory( self.device, diff --git a/python/sglang/srt/utils/patch_torch.py b/python/sglang/srt/utils/patch_torch.py index 9b4e21154..37868c80c 100644 --- a/python/sglang/srt/utils/patch_torch.py +++ b/python/sglang/srt/utils/patch_torch.py @@ -21,24 +21,42 @@ from sglang.srt.utils.common import is_npu _is_npu = is_npu() +if _is_npu: + from torch_npu.multiprocessing import reductions as npu_reductions + + def _rebuild_npu_tensor_modified(*args): + args = _modify_tuple(args, _REDUCE_TENSOR_ARG_DEVICE_INDEX, npu_verl_to_sglang) + return npu_reductions._rebuild_npu_tensor_original(*args) + + def npu_verl_to_sglang(device: int): + assert ( + SGLANG_TP_RANK is not None + ), "SGLANG_TP_RANK is not registered. Please call register_sgl_tp_rank() first." + return SGLANG_TP_RANK + + +SGLANG_TP_RANK = None + def monkey_patch_torch_reductions(): """Monkey patching before Torch https://github.com/pytorch/pytorch/pull/149248 is fixed""" - # Currently, NPU does not support UUID. This has been temporarily commented out, with support expected in the fourth quarter. - if _is_npu: - return + if not _is_npu: + if hasattr(reductions, "_reduce_tensor_original"): + return + reductions._reduce_tensor_original = reductions.reduce_tensor + reductions._rebuild_cuda_tensor_original = reductions.rebuild_cuda_tensor - if hasattr(reductions, "_reduce_tensor_original"): - return + reductions.reduce_tensor = _reduce_tensor_modified + reductions.rebuild_cuda_tensor = _rebuild_cuda_tensor_modified + reductions.init_reductions() + else: + # FIXME: This is a temp patch for npu as HDK does not support device uuid for now + if hasattr(npu_reductions, "_rebuild_npu_tensor_original"): + return - reductions._reduce_tensor_original = reductions.reduce_tensor - reductions._rebuild_cuda_tensor_original = reductions.rebuild_cuda_tensor - - reductions.reduce_tensor = _reduce_tensor_modified - reductions.rebuild_cuda_tensor = _rebuild_cuda_tensor_modified - - reductions.init_reductions() + npu_reductions._rebuild_npu_tensor_original = npu_reductions.rebuild_npu_tensor + npu_reductions.rebuild_npu_tensor = _rebuild_npu_tensor_modified # The signature has not been changed for years, and we will not need this when the next version is released, @@ -46,6 +64,11 @@ def monkey_patch_torch_reductions(): _REDUCE_TENSOR_ARG_DEVICE_INDEX = 6 +def register_sgl_tp_rank(rank: int): + global SGLANG_TP_RANK + SGLANG_TP_RANK = rank + + def _reduce_tensor_modified(*args, **kwargs): output_fn, output_args = reductions._reduce_tensor_original(*args, **kwargs) output_args = _modify_tuple(