6a4b81e2d9
Signed-off-by: chenzhenyang <andy271828@163.com> Signed-off-by: chenzhenyang <chenzhenyang@moonshot.cn> Co-authored-by: chenzhenyang <chenzhenyang@moonshot.cn> Co-authored-by: zhaochenyang20 <zhaochen20@outlook.com>
105 lines
3.3 KiB
Python
105 lines
3.3 KiB
Python
import unittest
|
|
|
|
import torch
|
|
|
|
from sglang.srt.configs.qwen3_vl import Qwen3VLConfig
|
|
from sglang.srt.distributed.parallel_state import (
|
|
init_distributed_environment,
|
|
initialize_model_parallel,
|
|
)
|
|
from sglang.srt.layers.dp_attention import initialize_dp_attention
|
|
from sglang.srt.layers.quantization.unquant import (
|
|
LinearMethodBase,
|
|
UnquantizedLinearMethod,
|
|
)
|
|
from sglang.srt.models.qwen3_vl import Qwen3VLMoeVisionModel
|
|
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
|
|
|
|
|
|
def unpack(tensor, dim_len, pack_len):
|
|
dim_part = dim_len // pack_len
|
|
ret_val = tensor.reshape(dim_part, dim_part, pack_len, pack_len, -1)
|
|
ret_val = ret_val.permute(4, 0, 2, 1, 3).reshape(1, -1, dim_len, dim_len)
|
|
return ret_val
|
|
|
|
|
|
class TestEmbedInterpolate(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.pDevice = torch.get_default_device()
|
|
torch.set_default_device("npu")
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
torch.set_default_device(cls.pDevice)
|
|
|
|
def test_embed_interpolate(self):
|
|
self.assertTrue(issubclass(UnquantizedLinearMethod, LinearMethodBase))
|
|
t_dim = [16, 32]
|
|
s_dim = [192, 574]
|
|
sarg = ServerArgs(model_path="dummy", device="npu")
|
|
mconf = Qwen3VLConfig(
|
|
hidden_size=64,
|
|
num_heads=1,
|
|
num_position_embeddings=2304,
|
|
patch_size=16,
|
|
spatial_merge_size=2,
|
|
temporal_patch_size=2,
|
|
deepstack_visual_indexes=[5, 11, 17],
|
|
in_channels=3,
|
|
depth=24,
|
|
intermediate_size=256,
|
|
hidden_act="gelu_pytorch_tanh",
|
|
out_hidden_size=2560,
|
|
)
|
|
set_global_server_args_for_scheduler(sarg)
|
|
init_distributed_environment(
|
|
backend="gloo",
|
|
world_size=1,
|
|
rank=0,
|
|
local_rank=0,
|
|
distributed_init_method="tcp://127.0.0.1:2646",
|
|
)
|
|
initialize_model_parallel()
|
|
initialize_dp_attention(
|
|
server_args=sarg,
|
|
model_config=mconf,
|
|
)
|
|
model = Qwen3VLMoeVisionModel(
|
|
mconf,
|
|
quant_config=None,
|
|
norm_eps=1e-6,
|
|
prefix="visual",
|
|
)
|
|
grid_thw = torch.tensor(
|
|
[(t, s, s) for t, s in zip(t_dim, s_dim)], dtype=torch.int32
|
|
)
|
|
embeddings = model.fast_pos_embed_interpolate(grid_thw)
|
|
|
|
embeddings_s0 = embeddings[: s_dim[0] * s_dim[0], :]
|
|
embeddings_s1 = embeddings[s_dim[0] * s_dim[0] : 2 * s_dim[0] * s_dim[0], :]
|
|
self.assertTrue(torch.allclose(embeddings_s0, embeddings_s1, atol=5e-5))
|
|
|
|
embeddings_l = embeddings[
|
|
t_dim[0] * s_dim[0] * s_dim[0] : t_dim[0] * s_dim[0] * s_dim[0]
|
|
+ s_dim[1] * s_dim[1],
|
|
:,
|
|
]
|
|
embeddings_s0 = torch.nn.functional.interpolate(
|
|
unpack(embeddings_s0, s_dim[0], 2),
|
|
size=(48, 48),
|
|
mode="area",
|
|
)
|
|
embeddings_r = torch.nn.functional.interpolate(
|
|
unpack(embeddings_l, s_dim[1], 2),
|
|
size=(48, 48),
|
|
mode="area",
|
|
)
|
|
self.assertTrue(
|
|
torch.allclose(embeddings_s0, embeddings_r, atol=5e-1, rtol=5e-1)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|