[CI] Run the Clef pooler test on GPU (#61409)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5.5
parent
de8291e24f
commit
61d9c777bf
@@ -15,9 +15,11 @@ from vllm.model_executor.models.clef_joint_schema import (
|
||||
JointSchemaHead,
|
||||
JointSchemaPooler,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.pooling_params import PoolingParams
|
||||
from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
|
||||
|
||||
DEVICE_TYPE = current_platform.device_type
|
||||
HIDDEN, VOCAB = 64, 500
|
||||
TYPES = [0, 1, 1, 2]
|
||||
COUNTS = [2, 3, 1, 5]
|
||||
@@ -55,7 +57,7 @@ def reference_logits(head, hidden, token_ids, embedding, question_spans, spans):
|
||||
base
|
||||
+ head.option_summary_norm(torch.stack(summaries))
|
||||
+ head.global_projection(global_vector)
|
||||
+ head.type_embedding(torch.tensor(TYPES))
|
||||
+ head.type_embedding(torch.tensor(TYPES, device=hidden.device))
|
||||
).unsqueeze(0)
|
||||
for layer in head.layers:
|
||||
fields = layer(fields, memory)
|
||||
@@ -90,7 +92,8 @@ def pooler():
|
||||
with torch.no_grad():
|
||||
for p in head.parameters():
|
||||
p.add_(torch.randn_like(p) * 0.2)
|
||||
embedding = torch.randn(VOCAB, HIDDEN)
|
||||
head.to(DEVICE_TYPE)
|
||||
embedding = torch.randn(VOCAB, HIDDEN, device=DEVICE_TYPE)
|
||||
config = SimpleNamespace(
|
||||
scheduler_config=SimpleNamespace(
|
||||
enable_chunked_prefill=False, async_scheduling=False
|
||||
@@ -132,7 +135,7 @@ def test_pooler_matches_reference(pooler):
|
||||
pooler, embedding = pooler
|
||||
lengths = [120, 90, 100]
|
||||
token_ids = [torch.randint(0, VOCAB, (n,)) for n in lengths]
|
||||
hidden = torch.randn(sum(lengths), HIDDEN)
|
||||
hidden = torch.randn(sum(lengths), HIDDEN, device=DEVICE_TYPE)
|
||||
first, first_questions, first_spans = schema(lengths[0])
|
||||
last, last_questions, last_spans = schema(lengths[2])
|
||||
params = [
|
||||
@@ -154,7 +157,7 @@ def test_pooler_matches_reference(pooler):
|
||||
metadata.build_pooling_cursor(
|
||||
torch.tensor(lengths).numpy(),
|
||||
seq_lens_cpu=torch.tensor(lengths),
|
||||
device=torch.device("cpu"),
|
||||
device=hidden.device,
|
||||
)
|
||||
probs, empty, logits = pooler(hidden, metadata)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user