117 lines
4.1 KiB
Python
117 lines
4.1 KiB
Python
# Copyright 2023-2024 SGLang Team
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""
|
|
End-to-end tests for the --enable-lora-overlap-loading server argument.
|
|
"""
|
|
|
|
import multiprocessing as mp
|
|
import unittest
|
|
|
|
from sglang.test.ci.ci_register import register_cuda_ci
|
|
from sglang.test.lora_utils import (
|
|
CI_MULTI_LORA_MODELS,
|
|
TEST_MULTIPLE_BATCH_PROMPTS,
|
|
TORCH_DTYPES,
|
|
LoRAModelCase,
|
|
ensure_reproducibility,
|
|
)
|
|
from sglang.test.runners import SRTRunner
|
|
from sglang.test.test_utils import CustomTestCase, calculate_rouge_l
|
|
|
|
register_cuda_ci(est_time=300, suite="stage-b-test-small-1-gpu")
|
|
|
|
|
|
class TestLoRAPipelineLoading(CustomTestCase):
|
|
|
|
def _run_mixed_batch_test(
|
|
self,
|
|
model_case: LoRAModelCase,
|
|
torch_dtype,
|
|
):
|
|
base_path = model_case.base
|
|
adaptor_paths = [a.name for a in model_case.adaptors]
|
|
print(
|
|
f"\n========== Testing mixed batch LoRA overlap loading on base '{base_path}' "
|
|
f"with dtype={torch_dtype} ==========\n"
|
|
)
|
|
ensure_reproducibility()
|
|
max_new_tokens = 32
|
|
|
|
prompts = TEST_MULTIPLE_BATCH_PROMPTS[:3]
|
|
configs = [
|
|
[None, adaptor_paths[0], adaptor_paths[1]],
|
|
[adaptor_paths[0], None, adaptor_paths[1]],
|
|
[adaptor_paths[0], adaptor_paths[1], None],
|
|
[adaptor_paths[1], adaptor_paths[0], adaptor_paths[1]],
|
|
]
|
|
common_args = dict(
|
|
torch_dtype=torch_dtype,
|
|
model_type="generation",
|
|
tp_size=model_case.tp_size,
|
|
lora_paths=adaptor_paths,
|
|
max_loras_per_batch=model_case.max_loras_per_batch,
|
|
max_loaded_loras=model_case.max_loaded_loras,
|
|
disable_cuda_graph=True,
|
|
disable_radix_cache=True,
|
|
mem_fraction_static=0.65,
|
|
sleep_on_idle=True,
|
|
)
|
|
|
|
results_no_overlap_loading = []
|
|
with SRTRunner(
|
|
base_path, enable_lora_overlap_loading=False, **common_args
|
|
) as runner:
|
|
for lora_paths in configs:
|
|
results_no_overlap_loading.append(
|
|
runner.batch_forward(
|
|
prompts, max_new_tokens=max_new_tokens, lora_paths=lora_paths
|
|
).output_strs
|
|
)
|
|
|
|
results_overlap_loading = []
|
|
with SRTRunner(
|
|
base_path, enable_lora_overlap_loading=True, **common_args
|
|
) as runner:
|
|
for lora_paths in configs:
|
|
results_overlap_loading.append(
|
|
runner.batch_forward(
|
|
prompts, max_new_tokens=max_new_tokens, lora_paths=lora_paths
|
|
).output_strs
|
|
)
|
|
|
|
for i, (res_no_overlap_loading, res_overlap_loading) in enumerate(
|
|
zip(results_no_overlap_loading, results_overlap_loading)
|
|
):
|
|
scores = calculate_rouge_l(res_overlap_loading, res_no_overlap_loading)
|
|
for j, score in enumerate(scores):
|
|
assert score >= model_case.rouge_l_tolerance, (
|
|
f"Batch {i} prompt {j} mismatch: {score}\n"
|
|
f"Overlap loading: {res_overlap_loading[j]}\n"
|
|
f"No overlap loading: {res_no_overlap_loading[j]}"
|
|
)
|
|
|
|
def test_mixed_batch(self):
|
|
for model_case in CI_MULTI_LORA_MODELS:
|
|
for dtype in TORCH_DTYPES:
|
|
self._run_mixed_batch_test(model_case, dtype)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
mp.set_start_method("spawn")
|
|
except RuntimeError:
|
|
pass
|
|
|
|
unittest.main(warnings="ignore")
|