Raphaël RiallandandSimon Veitner
c32495d58f
Signed-off-by: Raphael Rialland <raphael.rialland@mistral.ai> Signed-off-by: Simon Veitner <sveitner@redhat.com> Co-authored-by: Simon Veitner <sveitner@redhat.com>
107 lines
3.1 KiB
Python
107 lines
3.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
|
|
from tests.utils import multi_gpu_test
|
|
from vllm import SamplingParams
|
|
from vllm.distributed import cleanup_dist_env_and_memory
|
|
from vllm.engine.arg_utils import AsyncEngineArgs
|
|
from vllm.v1.engine.async_llm import AsyncLLM
|
|
|
|
MODEL = "hmellor/tiny-random-LlamaForCausalLM"
|
|
|
|
|
|
@multi_gpu_test(num_gpus=2)
|
|
@pytest.mark.parametrize(
|
|
"parallel_kwargs",
|
|
[
|
|
{"tensor_parallel_size": 2, "enable_batch_sharded_sampling": True},
|
|
{"pipeline_parallel_size": 2},
|
|
],
|
|
ids=["tp2-sharded-sampling", "pp2"],
|
|
)
|
|
def test_watermarked_generation_with_model_parallelism(vllm_runner, parallel_kwargs):
|
|
prompts = ["Parallel watermark smoke"] * 2
|
|
|
|
with vllm_runner(
|
|
MODEL,
|
|
dtype="half",
|
|
enforce_eager=True,
|
|
max_model_len=128,
|
|
max_num_seqs=4,
|
|
watermark_config={"algorithm": "gumbel", "key": 42},
|
|
**parallel_kwargs,
|
|
) as runner:
|
|
first = runner.llm.generate(
|
|
prompts,
|
|
SamplingParams(
|
|
temperature=0.8,
|
|
seed=1,
|
|
max_tokens=8,
|
|
ignore_eos=True,
|
|
),
|
|
)
|
|
second = runner.llm.generate(
|
|
prompts,
|
|
SamplingParams(
|
|
temperature=0.8,
|
|
seed=2,
|
|
max_tokens=8,
|
|
ignore_eos=True,
|
|
),
|
|
)
|
|
|
|
first_tokens = [output.outputs[0].token_ids for output in first]
|
|
second_tokens = [output.outputs[0].token_ids for output in second]
|
|
assert all(len(token_ids) == 8 for token_ids in first_tokens)
|
|
assert first_tokens[0] == first_tokens[1]
|
|
assert first_tokens == second_tokens
|
|
|
|
|
|
@multi_gpu_test(num_gpus=2)
|
|
def test_watermarked_generation_with_data_parallelism():
|
|
async def run() -> None:
|
|
engine = AsyncLLM.from_engine_args(
|
|
AsyncEngineArgs(
|
|
model=MODEL,
|
|
dtype="half",
|
|
data_parallel_size=2,
|
|
enforce_eager=True,
|
|
max_model_len=128,
|
|
max_num_seqs=4,
|
|
watermark_config={"algorithm": "gumbel", "key": 42},
|
|
)
|
|
)
|
|
|
|
async def generate(index: int) -> tuple[int, ...]:
|
|
final_output = None
|
|
async for output in engine.generate(
|
|
"Data-parallel watermark smoke",
|
|
SamplingParams(
|
|
temperature=0.8,
|
|
seed=index,
|
|
max_tokens=8,
|
|
ignore_eos=True,
|
|
),
|
|
request_id=f"watermark-dp-{index}",
|
|
data_parallel_rank=index,
|
|
):
|
|
final_output = output
|
|
assert final_output is not None
|
|
return tuple(final_output.outputs[0].token_ids)
|
|
|
|
try:
|
|
token_ids = await asyncio.gather(generate(0), generate(1))
|
|
assert len(token_ids[0]) == 8
|
|
assert token_ids[0] == token_ids[1]
|
|
finally:
|
|
engine.shutdown()
|
|
|
|
try:
|
|
asyncio.run(run())
|
|
finally:
|
|
cleanup_dist_env_and_memory()
|