[feat][Ascend][Mindspore]: support model-impl of mindspore (#9234)
This commit is contained in:
@@ -0,0 +1,162 @@
|
|||||||
|
# MindSpore Models
|
||||||
|
|
||||||
|
## Introduction
|
||||||
|
|
||||||
|
SGLang support run MindSpore framework models, this doc guide users to run mindspore models with SGLang.
|
||||||
|
|
||||||
|
## Requirements
|
||||||
|
|
||||||
|
MindSpore with SGLang current only support Ascend Npu device, users need first install Ascend CANN software packages.
|
||||||
|
The CANN software packages can download from the [Ascend Official Websites](https://www.hiascend.com). The version depends on the MindSpore version [MindSpore Installation](https://www.mindspore.cn/install)
|
||||||
|
|
||||||
|
## Supported Models
|
||||||
|
|
||||||
|
Currently, the following models are supported:
|
||||||
|
|
||||||
|
- **Qwen3**: Dense models supported. MoE models coming soon.
|
||||||
|
- *More models coming soon...*
|
||||||
|
|
||||||
|
## Installation
|
||||||
|
|
||||||
|
> **Note**: Currently, MindSpore models are provided by an independent package `sgl-mindspore`, which needs to be installed separately.
|
||||||
|
|
||||||
|
```shell
|
||||||
|
git clone https://github.com/chz34/sgl-mindspore.git
|
||||||
|
cd sgl-mindspore
|
||||||
|
pip install -e .
|
||||||
|
```
|
||||||
|
|
||||||
|
You will need to install the following packages, due to the support of tensor conversion through `dlpack` on 3rd devices, the minimum version of `PyTorch` is 2.7.1
|
||||||
|
|
||||||
|
```shell
|
||||||
|
pip install mindspore
|
||||||
|
pip install "torch>=2.7.1"
|
||||||
|
pip install "torch_npu>=2.7.1"
|
||||||
|
pip install triton_ascend
|
||||||
|
```
|
||||||
|
|
||||||
|
```shell
|
||||||
|
pip install -e "python[all_npu]"
|
||||||
|
```
|
||||||
|
|
||||||
|
## Run Model
|
||||||
|
|
||||||
|
Current SGLang-MindSpore support Qwen3 dense model, this doc uses Qwen3-8B as example.
|
||||||
|
|
||||||
|
### Offline infer
|
||||||
|
|
||||||
|
Use the following script for offline infer:
|
||||||
|
|
||||||
|
```python
|
||||||
|
import sglang as sgl
|
||||||
|
|
||||||
|
# Initialize the engine with MindSpore backend
|
||||||
|
llm = sgl.Engine(
|
||||||
|
model_path="/path/to/your/model", # Local model path
|
||||||
|
device="npu", # Use NPU device
|
||||||
|
model_impl="mindspore", # MindSpore implementation
|
||||||
|
attention_backend="ascend", # Attention backend
|
||||||
|
tp_size=1, # Tensor parallelism size
|
||||||
|
dp_size=1 # Data parallelism size
|
||||||
|
)
|
||||||
|
|
||||||
|
# Generate text
|
||||||
|
prompts = [
|
||||||
|
"Hello, my name is",
|
||||||
|
"The capital of France is",
|
||||||
|
"The future of AI is"
|
||||||
|
]
|
||||||
|
|
||||||
|
sampling_params = {"temperature": 0.01, "top_p": 0.9}
|
||||||
|
outputs = llm.generate(prompts, sampling_params)
|
||||||
|
|
||||||
|
for prompt, output in zip(prompts, outputs):
|
||||||
|
print(f"Prompt: {prompt}")
|
||||||
|
print(f"Generated: {output['text']}")
|
||||||
|
print("---")
|
||||||
|
```
|
||||||
|
|
||||||
|
### Start server
|
||||||
|
|
||||||
|
Launch a server with MindSpore backend:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Basic server startup
|
||||||
|
python3 -m sglang.launch_server \
|
||||||
|
--model-path /path/to/your/model \
|
||||||
|
--host 0.0.0.0 \
|
||||||
|
--device npu \
|
||||||
|
--model-impl mindspore \
|
||||||
|
--attention-backend ascend \
|
||||||
|
--tp-size 1 \
|
||||||
|
--dp-size 1
|
||||||
|
```
|
||||||
|
|
||||||
|
For distributed server with multiple nodes:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Multi-node distributed server
|
||||||
|
python3 -m sglang.launch_server \
|
||||||
|
--model-path /path/to/your/model \
|
||||||
|
--host 0.0.0.0 \
|
||||||
|
--device npu \
|
||||||
|
--model-impl mindspore \
|
||||||
|
--attention-backend ascend \
|
||||||
|
--dist-init-addr 127.0.0.1:29500 \
|
||||||
|
--nnodes 2 \
|
||||||
|
--node-rank 0 \
|
||||||
|
--tp-size 4 \
|
||||||
|
--dp-size 2
|
||||||
|
```
|
||||||
|
|
||||||
|
## Troubleshooting
|
||||||
|
|
||||||
|
#### Debug Mode
|
||||||
|
|
||||||
|
Enable sglang debug logging by log-level argument.
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python3 -m sglang.launch_server \
|
||||||
|
--model-path /path/to/your/model \
|
||||||
|
--host 0.0.0.0 \
|
||||||
|
--device npu \
|
||||||
|
--model-impl mindspore \
|
||||||
|
--attention-backend ascend \
|
||||||
|
--log-level DEBUG
|
||||||
|
```
|
||||||
|
|
||||||
|
Enable mindspore info and debug logging by setting environments.
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export GLOG_v=1 # INFO
|
||||||
|
export GLOG_v=0 # DEBUG
|
||||||
|
```
|
||||||
|
|
||||||
|
#### Explicitly select devices
|
||||||
|
|
||||||
|
Use the following environment variable to explicitly select the devices to use.
|
||||||
|
|
||||||
|
```shell
|
||||||
|
export ASCEND_RT_VISIBLE_DEVICES=4,5,6,7 # to set device
|
||||||
|
```
|
||||||
|
|
||||||
|
#### Some communication environment issues
|
||||||
|
|
||||||
|
In case of some environment with special communication environment, users need set some environment variables.
|
||||||
|
|
||||||
|
```shell
|
||||||
|
export MS_ENABLE_LCCL=off # current not support LCCL communication mode in SGLang-MindSpore
|
||||||
|
```
|
||||||
|
|
||||||
|
#### Some dependencies of protobuf
|
||||||
|
|
||||||
|
In case of some environment with special protobuf version, users need set some environment variables to avoid binary version mismatch.
|
||||||
|
|
||||||
|
```shell
|
||||||
|
export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python # to avoid protobuf binary version mismatch
|
||||||
|
```
|
||||||
|
|
||||||
|
## Support
|
||||||
|
For MindSpore-specific issues:
|
||||||
|
|
||||||
|
- Refer to the [MindSpore documentation](https://www.mindspore.cn/)
|
||||||
@@ -47,6 +47,7 @@ class ModelImpl(str, Enum):
|
|||||||
AUTO = "auto"
|
AUTO = "auto"
|
||||||
SGLANG = "sglang"
|
SGLANG = "sglang"
|
||||||
TRANSFORMERS = "transformers"
|
TRANSFORMERS = "transformers"
|
||||||
|
MINDSPORE = "mindspore"
|
||||||
|
|
||||||
|
|
||||||
def is_deepseek_nsa(config: PretrainedConfig) -> bool:
|
def is_deepseek_nsa(config: PretrainedConfig) -> bool:
|
||||||
|
|||||||
@@ -0,0 +1,118 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the SGLang project
|
||||||
|
"""ms_runner launch MindSpore distributed modules."""
|
||||||
|
|
||||||
|
import multiprocessing as mp
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import mindspore as ms
|
||||||
|
import torch
|
||||||
|
from mindspore._c_expression import GroupOptions
|
||||||
|
from mindspore.communication import create_group
|
||||||
|
|
||||||
|
from sglang.srt.distributed.parallel_state import _groups
|
||||||
|
|
||||||
|
|
||||||
|
class _Tmp:
|
||||||
|
def __init__(self):
|
||||||
|
self.sched_p = None
|
||||||
|
|
||||||
|
def set_sched_process(self, p):
|
||||||
|
self.sched_p = p
|
||||||
|
|
||||||
|
def __del__(self):
|
||||||
|
if self.sched_p:
|
||||||
|
self.sched_p.kill()
|
||||||
|
|
||||||
|
|
||||||
|
_tmp = _Tmp()
|
||||||
|
|
||||||
|
|
||||||
|
def _get_host_and_ip(distributed_init_method):
|
||||||
|
try:
|
||||||
|
_, ip_str, port_str = distributed_init_method.split(":")
|
||||||
|
ip = ip_str.split("/")[-1]
|
||||||
|
port = int(port_str)
|
||||||
|
except Exception as e:
|
||||||
|
raise RuntimeError(
|
||||||
|
"Cannot get host and port information from %s, error: %s!"
|
||||||
|
% (distributed_init_method, str(e))
|
||||||
|
)
|
||||||
|
|
||||||
|
return ip, port
|
||||||
|
|
||||||
|
|
||||||
|
def run_scheduler_init(rank, local_rank, world_size, master_addr, master_port):
|
||||||
|
with open(str(Path() / "schedule.log"), "w") as scheduler_f:
|
||||||
|
# For Python outputs.
|
||||||
|
sys.stdout = scheduler_f
|
||||||
|
sys.stderr = scheduler_f
|
||||||
|
# For C++ outputs.
|
||||||
|
os.dup2(scheduler_f.fileno(), 1)
|
||||||
|
os.dup2(scheduler_f.fileno(), 2)
|
||||||
|
os.environ["DEVICE_ID"] = str(local_rank)
|
||||||
|
os.environ["MS_WORKER_NUM"] = str(world_size)
|
||||||
|
os.environ["MS_ROLE"] = "MS_SCHED"
|
||||||
|
os.environ["MS_NODE_ID"] = str(rank)
|
||||||
|
os.environ["MS_SCHED_HOST"] = str(master_addr)
|
||||||
|
os.environ["MS_SCHED_PORT"] = str(master_port)
|
||||||
|
# This function is blocked until the whole cluster exits.
|
||||||
|
ms.communication.init()
|
||||||
|
|
||||||
|
|
||||||
|
def set_ms_parallel_env(rank, local_rank, world_size, init_method):
|
||||||
|
master_addr, master_port = _get_host_and_ip(init_method)
|
||||||
|
# change port avoiding port conflicts with torch
|
||||||
|
master_port = master_port + 35 if master_port < 65500 else master_port - 35
|
||||||
|
if not os.getenv("MS_ROLE"):
|
||||||
|
if rank == 0:
|
||||||
|
# Create a subprocess for scheduler of MindSpore, just for internal collaboration, not for collective communication
|
||||||
|
sched_p = mp.Process(
|
||||||
|
target=run_scheduler_init,
|
||||||
|
args=(rank, local_rank, world_size, master_addr, master_port),
|
||||||
|
)
|
||||||
|
sched_p.start()
|
||||||
|
global _tmp
|
||||||
|
_tmp.set_sched_process(sched_p)
|
||||||
|
|
||||||
|
os.environ["DEVICE_ID"] = str(local_rank)
|
||||||
|
os.environ["MS_WORKER_NUM"] = str(world_size)
|
||||||
|
os.environ["MS_ROLE"] = "MS_WORKER"
|
||||||
|
os.environ["MS_NODE_ID"] = str(rank)
|
||||||
|
os.environ["MS_SCHED_HOST"] = str(master_addr)
|
||||||
|
os.environ["MS_SCHED_PORT"] = str(master_port)
|
||||||
|
|
||||||
|
|
||||||
|
def reuse_hccl_comm():
|
||||||
|
for group_name, group in _groups.items():
|
||||||
|
# Torch ProcessGroupHccl
|
||||||
|
device_group = group().device_group
|
||||||
|
hccl_comm_handle = device_group._get_backend(torch.device("npu")).get_hccl_comm(
|
||||||
|
group().local_rank
|
||||||
|
)
|
||||||
|
print(
|
||||||
|
f"MindSpore reuse torch group: {device_group}, group_name: {group_name}, local rank: {group().local_rank},"
|
||||||
|
f"hccl communicator handle: {hex(hccl_comm_handle)}",
|
||||||
|
flush=True,
|
||||||
|
)
|
||||||
|
# Create MS communication group by hccl comm handle to reuse Torch group.
|
||||||
|
group_options = GroupOptions()
|
||||||
|
group_options.hccl_config = {"hccl_comm": hccl_comm_handle}
|
||||||
|
create_group(group_name, group().ranks, group_options)
|
||||||
|
|
||||||
|
|
||||||
|
def init_ms_distributed(world_size, rank, local_rank, server_args, port):
|
||||||
|
if server_args.dist_init_addr:
|
||||||
|
dist_init_method = f"tcp://{server_args.dist_init_addr}"
|
||||||
|
else:
|
||||||
|
dist_init_method = f"tcp://{server_args.host}:{port}"
|
||||||
|
set_ms_parallel_env(rank, local_rank, world_size, dist_init_method)
|
||||||
|
|
||||||
|
ms.set_context(infer_boost="on", jit_level="O0")
|
||||||
|
ms.set_context(mode=ms.context.PYNATIVE_MODE)
|
||||||
|
ms.set_device("Ascend", local_rank)
|
||||||
|
ms.communication.init("hccl")
|
||||||
|
# After distributed job is initialized, reuse hccl comms for MindSpore.
|
||||||
|
reuse_hccl_comm()
|
||||||
@@ -42,6 +42,7 @@ from sglang.srt.configs.load_config import LoadConfig, LoadFormat
|
|||||||
from sglang.srt.configs.model_config import (
|
from sglang.srt.configs.model_config import (
|
||||||
AttentionArch,
|
AttentionArch,
|
||||||
ModelConfig,
|
ModelConfig,
|
||||||
|
ModelImpl,
|
||||||
get_nsa_index_head_dim,
|
get_nsa_index_head_dim,
|
||||||
is_deepseek_nsa,
|
is_deepseek_nsa,
|
||||||
)
|
)
|
||||||
@@ -317,6 +318,8 @@ class ModelRunner:
|
|||||||
|
|
||||||
if get_bool_env_var("SGLANG_DETECT_SLOW_RANK"):
|
if get_bool_env_var("SGLANG_DETECT_SLOW_RANK"):
|
||||||
slow_rank_detector.execute()
|
slow_rank_detector.execute()
|
||||||
|
# Init mindspore running environment when model impl is "mindspore"
|
||||||
|
self.init_mindspore_runner()
|
||||||
|
|
||||||
# Update deep gemm configure
|
# Update deep gemm configure
|
||||||
if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
|
if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
|
||||||
@@ -364,6 +367,20 @@ class ModelRunner:
|
|||||||
else:
|
else:
|
||||||
self.piecewise_cuda_graph_runner = None
|
self.piecewise_cuda_graph_runner = None
|
||||||
|
|
||||||
|
def init_mindspore_runner(self):
|
||||||
|
# Init the mindspore runner
|
||||||
|
# for now, there is only some communication initialization work
|
||||||
|
if self.server_args.model_impl.lower() == ModelImpl.MINDSPORE and _is_npu:
|
||||||
|
from sglang.srt.model_executor.mindspore_runner import init_ms_distributed
|
||||||
|
|
||||||
|
init_ms_distributed(
|
||||||
|
world_size=self.tp_size * self.pp_size,
|
||||||
|
rank=self.tp_size * self.pp_rank + self.tp_rank,
|
||||||
|
local_rank=self.gpu_id,
|
||||||
|
server_args=self.server_args,
|
||||||
|
port=self.dist_port,
|
||||||
|
)
|
||||||
|
|
||||||
def initialize(self, min_per_gpu_memory: float):
|
def initialize(self, min_per_gpu_memory: float):
|
||||||
server_args = self.server_args
|
server_args = self.server_args
|
||||||
|
|
||||||
@@ -2018,6 +2035,9 @@ class ModelRunner:
|
|||||||
# TODO: Currently, cuda graph only captures decode steps, which only exists for generation models
|
# TODO: Currently, cuda graph only captures decode steps, which only exists for generation models
|
||||||
return
|
return
|
||||||
|
|
||||||
|
if self.server_args.model_impl.lower() == ModelImpl.MINDSPORE:
|
||||||
|
return
|
||||||
|
|
||||||
if self.device != "cpu" and self.server_args.disable_cuda_graph:
|
if self.device != "cpu" and self.server_args.disable_cuda_graph:
|
||||||
return
|
return
|
||||||
|
|
||||||
|
|||||||
@@ -99,7 +99,9 @@ def get_model_architecture(model_config: ModelConfig) -> Tuple[Type[nn.Module],
|
|||||||
supported_archs = ModelRegistry.get_supported_archs()
|
supported_archs = ModelRegistry.get_supported_archs()
|
||||||
is_native_supported = any(arch in supported_archs for arch in architectures)
|
is_native_supported = any(arch in supported_archs for arch in architectures)
|
||||||
|
|
||||||
if not is_native_supported or model_config.model_impl == ModelImpl.TRANSFORMERS:
|
if model_config.model_impl == ModelImpl.MINDSPORE:
|
||||||
|
architectures = ["MindSporeForCausalLM"]
|
||||||
|
elif not is_native_supported or model_config.model_impl == ModelImpl.TRANSFORMERS:
|
||||||
architectures = resolve_transformers_arch(model_config, architectures)
|
architectures = resolve_transformers_arch(model_config, architectures)
|
||||||
return ModelRegistry.resolve_model_cls(architectures)
|
return ModelRegistry.resolve_model_cls(architectures)
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,305 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the SGLang project
|
||||||
|
import logging
|
||||||
|
from typing import Any, Iterable, Optional, Tuple
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from sglang.srt.distributed import (
|
||||||
|
get_tensor_model_parallel_rank,
|
||||||
|
get_tensor_model_parallel_world_size,
|
||||||
|
)
|
||||||
|
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
|
||||||
|
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
||||||
|
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||||
|
from sglang.srt.models.registry import import_model_classes
|
||||||
|
from sglang.srt.utils import is_npu
|
||||||
|
|
||||||
|
_is_npu = is_npu()
|
||||||
|
|
||||||
|
if _is_npu:
|
||||||
|
import mindspore as ms
|
||||||
|
import numpy as np
|
||||||
|
import torch_npu
|
||||||
|
from mindspore import Tensor, mint, mutable
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def tensor_torch2ms(x: torch.Tensor):
|
||||||
|
if x is None or not isinstance(x, torch.Tensor):
|
||||||
|
return x
|
||||||
|
|
||||||
|
# torch tensor -> dlpack -> mindspore tensor
|
||||||
|
pt_dlpack = torch.utils.dlpack.to_dlpack(x)
|
||||||
|
ms_tensor = ms.utils.dlpack.from_dlpack(pt_dlpack)
|
||||||
|
return ms_tensor
|
||||||
|
|
||||||
|
|
||||||
|
def tensor_ms2torch(x: ms.Tensor):
|
||||||
|
if x is None or not isinstance(x, ms.Tensor):
|
||||||
|
return x
|
||||||
|
|
||||||
|
# ms tensor -> dlpack -> torch tensor
|
||||||
|
ms_dlpack = ms.utils.dlpack.to_dlpack(x)
|
||||||
|
torch_tensor = torch.utils.dlpack.from_dlpack(ms_dlpack)
|
||||||
|
torch_npu.npu.synchronize()
|
||||||
|
return torch_tensor
|
||||||
|
|
||||||
|
|
||||||
|
# Adapt from: https://gitee.com/mindspore/vllm-mindspore/blob/master/vllm_mindspore/model_executor/models/attention_mask.py
|
||||||
|
class LowerTriangularMask:
|
||||||
|
r"""
|
||||||
|
Provide Infer model attention mask.
|
||||||
|
Args:
|
||||||
|
dtype (ms dtype): The compute type of Infer model.
|
||||||
|
max_model_len (int): The max model length of Infer model.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, dtype, max_model_len, decode_mask_coeff=-10000.0):
|
||||||
|
self.dtype = dtype
|
||||||
|
self.max_model_len = max_model_len
|
||||||
|
self.cached_mask_len = 8 * 1024
|
||||||
|
self.decode_mask_coeff = decode_mask_coeff
|
||||||
|
|
||||||
|
prefill_mask_coeff = 1.0 if self.dtype == ms.bfloat16 else -10000.0
|
||||||
|
self.prefill_mask = Tensor(
|
||||||
|
np.triu(np.ones(shape=(128, 128), dtype=np.float16), k=1)
|
||||||
|
* prefill_mask_coeff,
|
||||||
|
dtype=self.dtype,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.hard_mask = mint.zeros((1, 1), dtype=dtype)
|
||||||
|
self.decode_mask = (
|
||||||
|
Tensor(
|
||||||
|
np.triu(
|
||||||
|
np.ones(
|
||||||
|
shape=(self.cached_mask_len, self.cached_mask_len),
|
||||||
|
dtype=np.int8,
|
||||||
|
),
|
||||||
|
k=1,
|
||||||
|
),
|
||||||
|
dtype=self.dtype,
|
||||||
|
)
|
||||||
|
* self.decode_mask_coeff
|
||||||
|
)
|
||||||
|
|
||||||
|
def create_mask(self, query_lens_np, seq_lens_np):
|
||||||
|
"""
|
||||||
|
when query_lens_np = [3], seq_lens_np = [6], decode_mask_coeff = 1
|
||||||
|
init attention mask
|
||||||
|
0 0 0 0 0 0
|
||||||
|
0 0 0 0 0 0
|
||||||
|
0 0 0 0 0 0
|
||||||
|
"""
|
||||||
|
max_seq_len = seq_lens_np.max().item()
|
||||||
|
total_q_len = query_lens_np.sum().item()
|
||||||
|
attention_mask = mint.zeros((total_q_len, max_seq_len), dtype=self.dtype)
|
||||||
|
|
||||||
|
req_num = query_lens_np.shape[0]
|
||||||
|
current_row = 0
|
||||||
|
for i in range(req_num):
|
||||||
|
q_len = query_lens_np[i].item()
|
||||||
|
current_row += q_len
|
||||||
|
# skip row when q_len <= 1, to decrease execute time
|
||||||
|
if q_len <= 1:
|
||||||
|
continue
|
||||||
|
seq_len = seq_lens_np[i].item()
|
||||||
|
context_len = seq_len - q_len
|
||||||
|
"""
|
||||||
|
set the right half to 1
|
||||||
|
0 0 0 1 1 1
|
||||||
|
0 0 0 1 1 1
|
||||||
|
0 0 0 1 1 1
|
||||||
|
"""
|
||||||
|
attention_mask[current_row - q_len : current_row, context_len:] = (
|
||||||
|
self.decode_mask_coeff
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
set the lower triangle of the right half to 0
|
||||||
|
0 0 0 0 1 1
|
||||||
|
0 0 0 0 0 1
|
||||||
|
0 0 0 0 0 0
|
||||||
|
"""
|
||||||
|
right_tensor = attention_mask[
|
||||||
|
current_row - q_len : current_row, context_len:seq_len
|
||||||
|
]
|
||||||
|
|
||||||
|
# use masked_fill_ to inplace modify attention_mask
|
||||||
|
right_tensor.masked_fill_(right_tensor.tril() == self.decode_mask_coeff, 0)
|
||||||
|
|
||||||
|
return attention_mask
|
||||||
|
|
||||||
|
def gen_attention_mask(
|
||||||
|
self,
|
||||||
|
is_prefill: bool,
|
||||||
|
position_ids: Tensor,
|
||||||
|
query_lens_np: np.ndarray,
|
||||||
|
seq_lens_np: np.ndarray,
|
||||||
|
):
|
||||||
|
max_query_len = query_lens_np.max()
|
||||||
|
max_seq_len = seq_lens_np.max()
|
||||||
|
if is_prefill:
|
||||||
|
attention_mask = self.prefill_mask
|
||||||
|
elif max_query_len > 1:
|
||||||
|
if max_seq_len <= self.cached_mask_len:
|
||||||
|
attention_mask = mint.index_select(self.decode_mask, 0, position_ids)
|
||||||
|
else:
|
||||||
|
attention_mask = self.create_mask(query_lens_np, seq_lens_np)
|
||||||
|
else:
|
||||||
|
attention_mask = self.hard_mask
|
||||||
|
return attention_mask
|
||||||
|
|
||||||
|
|
||||||
|
class MindSporeForCausalLM(torch.nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
config: Any,
|
||||||
|
quant_config: Optional[QuantizationConfig] = None,
|
||||||
|
prefix: str = "",
|
||||||
|
) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.config = config
|
||||||
|
|
||||||
|
ms.set_context(graph_kernel_flags="--disable_pass=gather_pre_rms_norm_fusion")
|
||||||
|
ms.set_kernel_launch_capture(False)
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"MindSporeForCausalLM tp size %d tp rank %d",
|
||||||
|
get_tensor_model_parallel_world_size(),
|
||||||
|
get_tensor_model_parallel_rank(),
|
||||||
|
)
|
||||||
|
if get_tensor_model_parallel_world_size() not in (1, 2, 4, 8):
|
||||||
|
# MatMulAllReduce only support tp size in (1, 2, 4, 8)
|
||||||
|
ms.set_context(graph_kernel_flags="--disable_pass=MatMulAllReduce")
|
||||||
|
|
||||||
|
arch = self.get_arch(self.config)
|
||||||
|
self.model = arch(config=config, quant_config=quant_config)
|
||||||
|
|
||||||
|
self.casual_mask = LowerTriangularMask(
|
||||||
|
self.config.param_dtype, self.config.max_position_embeddings
|
||||||
|
)
|
||||||
|
self.key_cache = []
|
||||||
|
self.value_cache = []
|
||||||
|
|
||||||
|
def get_arch(self, config):
|
||||||
|
# Get all implemented models
|
||||||
|
mindspore_models = import_model_classes("sgl_mindspore.models")
|
||||||
|
|
||||||
|
# Get arch from config
|
||||||
|
architectures = config.architectures
|
||||||
|
if isinstance(architectures, str):
|
||||||
|
architectures = [architectures]
|
||||||
|
if not architectures:
|
||||||
|
logger.warning("No model architectures are specified")
|
||||||
|
|
||||||
|
for arch in architectures:
|
||||||
|
if arch in mindspore_models:
|
||||||
|
return mindspore_models[arch]
|
||||||
|
if arch is None:
|
||||||
|
raise ValueError(f"Unsupported arch {architectures}")
|
||||||
|
|
||||||
|
@property
|
||||||
|
def use_mla(self):
|
||||||
|
return self.config.architectures[0] in ("DeepseekV3ForCausalLM")
|
||||||
|
|
||||||
|
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
|
||||||
|
self.model.load_weights(weights)
|
||||||
|
for _, cell in self.model.cells_and_names():
|
||||||
|
quant_method = getattr(cell, "quant_method", None)
|
||||||
|
if quant_method is not None:
|
||||||
|
quant_method.process_weights_after_loading(cell)
|
||||||
|
|
||||||
|
def get_kvcache(self, forward_batch: ForwardBatch):
|
||||||
|
def prepare_cache(cache_list, is_key_cache):
|
||||||
|
for i in range(self.config.num_hidden_layers):
|
||||||
|
if is_key_cache:
|
||||||
|
cache = forward_batch.token_to_kv_pool.get_key_buffer(i)
|
||||||
|
else:
|
||||||
|
cache = forward_batch.token_to_kv_pool.get_value_buffer(i)
|
||||||
|
cache_ms = tensor_torch2ms(cache)
|
||||||
|
if cache_ms.ndim == 3:
|
||||||
|
cache_ms = mint.unsqueeze(cache_ms, 2)
|
||||||
|
cache_list.append(cache_ms)
|
||||||
|
|
||||||
|
if self.use_mla:
|
||||||
|
if not self.key_cache:
|
||||||
|
prepare_cache(self.key_cache, is_key_cache=True)
|
||||||
|
return mutable(self.key_cache)
|
||||||
|
|
||||||
|
if self.key_cache and self.value_cache:
|
||||||
|
return mutable(self.key_cache), mutable(self.value_cache)
|
||||||
|
|
||||||
|
prepare_cache(self.key_cache, is_key_cache=True)
|
||||||
|
prepare_cache(self.value_cache, is_key_cache=False)
|
||||||
|
|
||||||
|
return mutable(self.key_cache), mutable(self.value_cache)
|
||||||
|
|
||||||
|
def prepare_inputs(self, input_ids, positions, forward_batch):
|
||||||
|
if self.use_mla:
|
||||||
|
key_cache = self.get_kvcache(forward_batch)
|
||||||
|
else:
|
||||||
|
key_cache, value_cache = self.get_kvcache(forward_batch)
|
||||||
|
|
||||||
|
# Different processing for the mindspore attention operator
|
||||||
|
# Without any prefix cache => Use FlashAttentionScore
|
||||||
|
# With cache => Use PagedAttention, no matter the query length is 1 or not
|
||||||
|
is_prefill = forward_batch.forward_mode.is_extend()
|
||||||
|
is_prefill = is_prefill and forward_batch.extend_prefix_lens.sum().item() == 0
|
||||||
|
|
||||||
|
batch_valid_length = forward_batch.seq_lens.cpu().numpy()
|
||||||
|
|
||||||
|
if forward_batch.extend_seq_lens is not None:
|
||||||
|
q_seq_lens = forward_batch.extend_seq_lens.cpu().numpy()
|
||||||
|
else:
|
||||||
|
q_seq_lens = np.ones([forward_batch.batch_size], dtype=np.int32)
|
||||||
|
|
||||||
|
page_size = forward_batch.token_to_kv_pool.page_size
|
||||||
|
block_tables = tensor_torch2ms(
|
||||||
|
(
|
||||||
|
forward_batch.req_to_token_pool.req_to_token[
|
||||||
|
forward_batch.req_pool_indices, : forward_batch.seq_lens.max()
|
||||||
|
][:, ::page_size]
|
||||||
|
// page_size
|
||||||
|
)
|
||||||
|
).to(ms.int32)
|
||||||
|
|
||||||
|
model_inputs = {}
|
||||||
|
model_inputs["input_ids"] = tensor_torch2ms(input_ids).to(ms.int32)
|
||||||
|
model_inputs["batch_valid_length"] = ms.Tensor(
|
||||||
|
batch_valid_length, dtype=ms.int32
|
||||||
|
)
|
||||||
|
model_inputs["position_ids"] = tensor_torch2ms(positions)
|
||||||
|
model_inputs["q_seq_lens"] = ms.Tensor(q_seq_lens, dtype=ms.int32)
|
||||||
|
model_inputs["attention_mask"] = self.casual_mask.gen_attention_mask(
|
||||||
|
is_prefill, model_inputs["position_ids"], q_seq_lens, batch_valid_length
|
||||||
|
).contiguous()
|
||||||
|
model_inputs["out_cache_loc"] = tensor_torch2ms(forward_batch.out_cache_loc).to(
|
||||||
|
ms.int32
|
||||||
|
)
|
||||||
|
model_inputs["is_prefill"] = is_prefill
|
||||||
|
model_inputs["key_cache"] = key_cache
|
||||||
|
if not self.use_mla:
|
||||||
|
model_inputs["value_cache"] = value_cache
|
||||||
|
model_inputs["block_tables"] = block_tables
|
||||||
|
return model_inputs
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
input_ids: torch.Tensor,
|
||||||
|
positions: torch.Tensor,
|
||||||
|
forward_batch: ForwardBatch,
|
||||||
|
) -> Tensor:
|
||||||
|
# prepare base inputs
|
||||||
|
model_inputs = self.prepare_inputs(input_ids, positions, forward_batch)
|
||||||
|
# prepare model inputs
|
||||||
|
model_inputs = self.model.prepare_inputs(forward_batch, model_inputs)
|
||||||
|
|
||||||
|
logits = self.model(**model_inputs)
|
||||||
|
|
||||||
|
# TODO: npu tensor ms2torch error to be fix, remain issues of torch_npu to get tensor from dlpack
|
||||||
|
logits_result = LogitsProcessorOutput(next_token_logits=tensor_ms2torch(logits))
|
||||||
|
return logits_result
|
||||||
|
|
||||||
|
|
||||||
|
EntryClass = [MindSporeForCausalLM]
|
||||||
@@ -2059,6 +2059,7 @@ class ServerArgs:
|
|||||||
"implementation is available.\n"
|
"implementation is available.\n"
|
||||||
'* "sglang" will use the SGLang model implementation.\n'
|
'* "sglang" will use the SGLang model implementation.\n'
|
||||||
'* "transformers" will use the Transformers model '
|
'* "transformers" will use the Transformers model '
|
||||||
|
'* "mindspore" will use the MindSpore model '
|
||||||
"implementation.\n",
|
"implementation.\n",
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -3957,6 +3958,9 @@ class ServerArgs:
|
|||||||
self.schedule_conservativeness >= 0
|
self.schedule_conservativeness >= 0
|
||||||
), "schedule_conservativeness must be non-negative"
|
), "schedule_conservativeness must be non-negative"
|
||||||
|
|
||||||
|
if self.model_impl == "mindspore":
|
||||||
|
assert is_npu(), "MindSpore model impl is only supported on Ascend npu."
|
||||||
|
|
||||||
def check_lora_server_args(self):
|
def check_lora_server_args(self):
|
||||||
assert self.max_loras_per_batch > 0, "max_loras_per_batch must be positive"
|
assert self.max_loras_per_batch > 0, "max_loras_per_batch must be positive"
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,64 @@
|
|||||||
|
"""
|
||||||
|
Usage:
|
||||||
|
python3 -m unittest test_mindspore_models.TestMindSporeQwen3.test_gsm8k
|
||||||
|
"""
|
||||||
|
|
||||||
|
import unittest
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
from sglang.srt.utils import kill_process_tree
|
||||||
|
from sglang.test.few_shot_gsm8k import run_eval
|
||||||
|
from sglang.test.test_utils import (
|
||||||
|
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||||
|
DEFAULT_URL_FOR_TEST,
|
||||||
|
CustomTestCase,
|
||||||
|
popen_launch_server,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class TestMindSporeQwen3(CustomTestCase):
|
||||||
|
@classmethod
|
||||||
|
def setUpClass(cls):
|
||||||
|
cls.model = "Qwen/Qwen3-8B"
|
||||||
|
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||||
|
cls.process = popen_launch_server(
|
||||||
|
cls.model,
|
||||||
|
cls.base_url,
|
||||||
|
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||||
|
other_args=[
|
||||||
|
"--device",
|
||||||
|
"npu",
|
||||||
|
"--model-impl",
|
||||||
|
"mindspore",
|
||||||
|
"--attention-backend",
|
||||||
|
"ascend",
|
||||||
|
"--tp-size",
|
||||||
|
"1",
|
||||||
|
"--dp-size",
|
||||||
|
"1",
|
||||||
|
"--mem-fraction-static",
|
||||||
|
0.8,
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def tearDownClass(cls):
|
||||||
|
kill_process_tree(cls.process.pid)
|
||||||
|
|
||||||
|
def test_gsm8k(self):
|
||||||
|
args = SimpleNamespace(
|
||||||
|
num_shots=5,
|
||||||
|
data_path=None,
|
||||||
|
num_questions=200,
|
||||||
|
max_new_tokens=512,
|
||||||
|
parallel=128,
|
||||||
|
host="http://127.0.0.1",
|
||||||
|
port=int(self.base_url.split(":")[-1]),
|
||||||
|
)
|
||||||
|
metrics = run_eval(args)
|
||||||
|
print(f"{metrics=}")
|
||||||
|
self.assertGreater(metrics["accuracy"], 0.78)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
@@ -536,6 +536,9 @@ suite_ascend = {
|
|||||||
TestFile("ascend/test_ascend_deepep.py", 400),
|
TestFile("ascend/test_ascend_deepep.py", 400),
|
||||||
TestFile("ascend/test_ascend_deepseek_mtp.py", 400),
|
TestFile("ascend/test_ascend_deepseek_mtp.py", 400),
|
||||||
],
|
],
|
||||||
|
"__not_in_ascend_ci__": [
|
||||||
|
TestFile("ascend/test_mindspore_models.py"),
|
||||||
|
],
|
||||||
}
|
}
|
||||||
|
|
||||||
suites.update(suite_amd)
|
suites.update(suite_amd)
|
||||||
|
|||||||
Reference in New Issue
Block a user