Skip to content

test: [TRTLLM-4000] Port multi GPU changes to GitHub #3027

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Funatiq merged 1 commit into from
Mar 26, 2025
Merged

test: [TRTLLM-4000] Port multi GPU changes to GitHub #3027

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions cpp/tests/executor/executorTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4117,10 +4117,16 @@ TEST_P(TimeoutTest, TimeoutStreamingTest)
}
else
{
if (val != NULL && !isMultiGpu)
{
GTEST_SKIP() << "Skipping SingleGpu tests";
}

if (!isMultiGpu && !useOrchestratorMode)
{
GTEST_SKIP() << "Leader mode on single GPU crashes";
}

// Check that it was launched with right number of MPI ranks
if (!useOrchestratorMode && COMM_SESSION.getSize() != 4)
{
Expand Down Expand Up @@ -4324,10 +4330,16 @@ TEST_P(TimeoutTest, TimeoutNonstreamingTest)
}
else
{
if (val != NULL && !isMultiGpu)
{
GTEST_SKIP() << "Skipping SingleGpu tests";
}

if (!isMultiGpu && !useOrchestratorMode)
{
GTEST_SKIP() << "Leader mode on single GPU crashes";
}

// Check that it was launched with right number of MPI ranks
if (!useOrchestratorMode && COMM_SESSION.getSize() != 4)
{
Expand Down
150 changes: 81 additions & 69 deletions tests/integration/defs/cpp_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,14 @@
import os as _os
import pathlib as _pl
import subprocess as _sp
import sys as _sys
from typing import Generator, List, Optional, Sequence

build_script_dir = _pl.Path(
__file__).parent.resolve().parent.parent.parent / "scripts"
assert build_script_dir.is_dir()
_sys.path.append(str(build_script_dir))

from build_wheel import get_build_dir as get_trt_llm_build_dir

default_test_parallel = 2
Expand Down Expand Up @@ -237,8 +243,7 @@ def produce_mpirun_command(*, global_commands, nranks, local_commands,
return l[:-1]


def run_multi_gpu_tests(build_dir: _pl.Path, timeout=1500):

def run_simple_multi_gpu_tests(build_dir: _pl.Path, timeout=1500):
tests_dir = build_dir / "tests"
cpp_env = {**_os.environ}
# Utils tests
Expand Down Expand Up @@ -278,17 +283,27 @@ def run_multi_gpu_tests(build_dir: _pl.Path, timeout=1500):
env=new_env,
timeout=600)

xml_output_file = build_dir / "results-multi-gpu-real-decoder.xml"
trt_model_test = produce_mpirun_command(
global_commands=["mpirun", "--allow-run-as-root"],
nranks=4,
local_commands=[
"batch_manager/trtGptModelRealDecoderTest",
"--gtest_filter=*TP*:*PP*"
],
leader_commands=[f"--gtest_output=xml:{xml_output_file}"])
run_command(trt_model_test, cwd=tests_dir, env=cpp_env,
timeout=timeout) # expecting ~ 1200s
# UCX transceiver tests, the test may not be built if ENABLE_UCX is 0
if _os.path.exists(
_os.path.join(tests_dir, "batch_manager/ucxDataTransceiverTest")):
ucx_env = copy.copy(cpp_env)
ucx_env["UCX_MEMTYPE_CACHE"] = "n"
ucx_trans_test = [
"mpirun",
"-n",
"2",
"--allow-run-as-root",
"batch_manager/ucxDataTransceiverTest",
]
run_command(ucx_trans_test, cwd=tests_dir, env=ucx_env, timeout=300)
else:
_logger.info(
"batch_manager/ucxDataTransceiverTest not found, so skipping.")


def run_llama_executor_multi_gpu_tests(build_dir: _pl.Path, timeout=1500):
tests_dir = build_dir / "tests"
cpp_env = {**_os.environ}

mgpu_env = copy.copy(cpp_env)
mgpu_env["RUN_LLAMA_MULTI_GPU"] = "true"
Expand Down Expand Up @@ -316,19 +331,6 @@ def run_multi_gpu_tests(build_dir: _pl.Path, timeout=1500):
]
run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500)

#EncDec test in leader mode
xml_output_file = build_dir / "results-multi-gpu-t5-exec-leader-mode.xml"
trt_model_test = produce_mpirun_command(
global_commands=["mpirun", "--allow-run-as-root"],
nranks=4,
local_commands=[
"executor/encDecTest",
"--gtest_filter=T5MultiGPUTest/EncDecParamsTest.Forward*"
],
leader_commands=[f"--gtest_output=xml:{xml_output_file}"],
)
run_command(trt_model_test, cwd=tests_dir, env=cpp_env, timeout=1500)

#Logits processor and guided decoding test in leader mode
xml_output_file = build_dir / "results-multi-gpu-logits-proc.xml"
tp_pp_sizes = [(4, 1), (2, 2), (1, 4)]
Expand All @@ -350,24 +352,44 @@ def run_multi_gpu_tests(build_dir: _pl.Path, timeout=1500):
leader_commands=[f"--gtest_output=xml:{xml_output_file}"])
run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500)

# UCX transceiver tests, the test may not be built if ENABLE_UCX is 0
if _os.path.exists(
_os.path.join(tests_dir, "batch_manager/ucxDataTransceiverTest")):
ucx_env = copy.copy(cpp_env)
ucx_env["UCX_MEMTYPE_CACHE"] = "n"
ucx_trans_test = [
"mpirun",
"-n",
"2",
"--allow-run-as-root",
"batch_manager/ucxDataTransceiverTest",
]
run_command(ucx_trans_test, cwd=tests_dir, env=ucx_env, timeout=300)

run_disagg_tests(build_dir)
def run_t5_multi_gpu_tests(build_dir: _pl.Path, timeout=1500):
tests_dir = build_dir / "tests"
cpp_env = {**_os.environ}

#EncDec test in leader mode
xml_output_file = build_dir / "results-multi-gpu-t5-exec-leader-mode.xml"
trt_model_test = produce_mpirun_command(
global_commands=["mpirun", "--allow-run-as-root"],
nranks=4,
local_commands=[
"executor/encDecTest",
"--gtest_filter=T5MultiGPUTest/EncDecParamsTest.Forward*"
],
leader_commands=[f"--gtest_output=xml:{xml_output_file}"],
)
run_command(trt_model_test, cwd=tests_dir, env=cpp_env, timeout=1500)


def run_trt_gpt_model_real_decoder_multi_gpu_tests(build_dir: _pl.Path,
timeout=1500):
tests_dir = build_dir / "tests"
cpp_env = {**_os.environ}

xml_output_file = build_dir / "results-multi-gpu-real-decoder.xml"
trt_model_test = produce_mpirun_command(
global_commands=["mpirun", "--allow-run-as-root"],
nranks=4,
local_commands=[
"batch_manager/trtGptModelRealDecoderTest",
"--gtest_filter=*TP*:*PP*"
],
leader_commands=[f"--gtest_output=xml:{xml_output_file}"])
run_command(trt_model_test, cwd=tests_dir, env=cpp_env,
timeout=timeout) # expecting ~ 1200s

def run_disagg_tests(build_dir: _pl.Path):

def run_disagg_multi_gpu_tests(build_dir: _pl.Path):

tests_dir = build_dir / "tests"
cpp_env = {**_os.environ}
Expand Down Expand Up @@ -549,38 +571,28 @@ def prepare_model_tests(model_name: str,
timeout=600)


def prepare_multi_gpu_model_tests(python_exe: str,
def prepare_multi_gpu_model_tests(test_list: List[str],
python_exe: str,
root_dir: _pl.Path,
resources_dir: _pl.Path,
model_cache: Optional[str] = None):
model_cache_arg = ["--model_cache", model_cache] if model_cache else []
only_multi_gpu_arg = ["--only_multi_gpu"]

prepare_model_tests(model_name="llama",
python_exe=python_exe,
root_dir=root_dir,
resources_dir=resources_dir,
model_cache_arg=model_cache_arg,
only_multi_gpu_arg=only_multi_gpu_arg)

prepare_model_tests(model_name="llama",
python_exe=python_exe,
root_dir=root_dir,
resources_dir=resources_dir,
model_cache_arg=model_cache_arg)

prepare_model_tests(model_name="t5",
python_exe=python_exe,
root_dir=root_dir,
resources_dir=resources_dir,
model_cache_arg=model_cache_arg,
only_multi_gpu_arg=['--tp', '4', '--pp', '1'])

prepare_model_tests(model_name="gpt",
python_exe=python_exe,
root_dir=root_dir,
resources_dir=resources_dir,
model_cache_arg=model_cache_arg)

if "llama" in test_list:
prepare_model_tests(model_name="llama",
python_exe=python_exe,
root_dir=root_dir,
resources_dir=resources_dir,
model_cache_arg=model_cache_arg,
only_multi_gpu_arg=["--only_multi_gpu"])

if "t5" in test_list:
prepare_model_tests(model_name="t5",
python_exe=python_exe,
root_dir=root_dir,
resources_dir=resources_dir,
model_cache_arg=model_cache_arg,
only_multi_gpu_arg=['--tp', '4', '--pp', '1'])


def run_single_gpu_tests(build_dir: _pl.Path,
Expand Down
Loading