From 68c2f79a0c2aab845fc0b9ce3c96085c76d84561 Mon Sep 17 00:00:00 2001 From: Dom Brown <3886319+DomBrown@users.noreply.github.com> Date: Mon, 12 May 2025 15:03:23 +0000 Subject: [PATCH 01/11] Fix padded vocab size for Llama Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com> --- cpp/tests/executor/disaggExecutorTest.cpp | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/cpp/tests/executor/disaggExecutorTest.cpp b/cpp/tests/executor/disaggExecutorTest.cpp index f2fd2d9b1f47..15c15319e1e9 100644 --- a/cpp/tests/executor/disaggExecutorTest.cpp +++ b/cpp/tests/executor/disaggExecutorTest.cpp @@ -29,6 +29,7 @@ namespace { auto constexpr LLAMA_INPUT_FILE = "input_tokens_llama.npy"; +auto constexpr LLAMA_VOCAB_SIZE_PADDED = 128256; auto constexpr LLAMA_END_ID = 128001; auto constexpr LLAMA_PAD_ID = 128001; @@ -520,6 +521,7 @@ TEST_P(DisaggParamsTest, DisaggTokenComparison) // set defaults and adjust if needed by different models fs::path inputPath = DATA_PATH / "input_tokens.npy"; ModelIds modelIds{50256, 50256}; + SizeType32 vocabSizePadded{50257}; // gpt vocabSizePadded bool isSpeculativeDecoding{false}; // NOTE: This can be used to disable checks for certain prompt batch entries @@ -557,6 +559,7 @@ TEST_P(DisaggParamsTest, DisaggTokenComparison) || modelName == "llama_tp1_pp2_cp1" || modelName == "llama_tp2_pp1_cp1" || modelName == "llama_tp1_pp1_cp1") { inputPath = DATA_PATH / LLAMA_INPUT_FILE; + vocabSizePadded = LLAMA_VOCAB_SIZE_PADDED; auto const resultsPath = LLAMA_DATA_PATH / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); @@ -619,8 +622,6 @@ TEST_P(DisaggParamsTest, DisaggTokenComparison) } } - SizeType32 constexpr vocabSizePadded{50257}; // gpt vocabSizePadded - // Returning logits will bring higher latency if (streaming && (outConfig.returnContextLogits || outConfig.returnGenerationLogits)) { @@ -756,6 +757,7 @@ TEST_P(DisaggOrchestratorParamsTest, DisaggTokenComparison) // set defaults and adjust if needed by different models fs::path inputPath = DATA_PATH / "input_tokens.npy"; ModelIds modelIds{50256, 50256}; + SizeType32 vocabSizePadded{50257}; // gpt vocabSizePadded bool isSpeculativeDecoding{false}; // NOTE: This can be used to disable checks for certain prompt batch entries @@ -764,6 +766,7 @@ TEST_P(DisaggOrchestratorParamsTest, DisaggTokenComparison) || modelName == "llama_tp1_pp2" || modelName == "llama_tp2_pp1" || modelName == "llama_tp1_pp1") { inputPath = DATA_PATH / LLAMA_INPUT_FILE; + vocabSizePadded = LLAMA_VOCAB_SIZE_PADDED; auto const resultsPath = LLAMA_DATA_PATH / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); @@ -827,8 +830,6 @@ TEST_P(DisaggOrchestratorParamsTest, DisaggTokenComparison) } } - SizeType32 constexpr vocabSizePadded{50257}; // gpt vocabSizePadded - // Returning logits will bring higher latency if (streaming && (outConfig.returnContextLogits || outConfig.returnGenerationLogits)) { @@ -920,6 +921,7 @@ TEST_P(ConditionalDisaggParamsTest, DisaggTokenComparison) // set defaults and adjust if needed by different models fs::path inputPath = DATA_PATH / "input_tokens.npy"; ModelIds modelIds{50256, 50256}; + SizeType32 vocabSizePadded{50257}; // gpt vocabSizePadded bool isSpeculativeDecoding{false}; // NOTE: This can be used to disable checks for certain prompt batch entries @@ -935,6 +937,7 @@ TEST_P(ConditionalDisaggParamsTest, DisaggTokenComparison) else if (modelName == "llama_tp1_pp1_cp1") { inputPath = DATA_PATH / LLAMA_INPUT_FILE; + vocabSizePadded = LLAMA_VOCAB_SIZE_PADDED; auto const resultsPath = LLAMA_DATA_PATH / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); @@ -948,8 +951,6 @@ TEST_P(ConditionalDisaggParamsTest, DisaggTokenComparison) TLLM_THROW("Unrecognized modelName"); } - SizeType32 constexpr vocabSizePadded{50257}; // gpt vocabSizePadded - auto executorConfig = ExecutorConfig(maxBeamWidth); FloatType freeGpuMemoryFraction = 0.9f; KvCacheConfig kvCacheConfig{true, std::nullopt, std::nullopt, std::nullopt, freeGpuMemoryFraction}; From eda828d5effd3ed2ee66dcdc249aa3129787451e Mon Sep 17 00:00:00 2001 From: Dom Brown <3886319+DomBrown@users.noreply.github.com> Date: Mon, 12 May 2025 16:34:08 +0000 Subject: [PATCH 02/11] Refactor multi GPU llama executor tests, and reuse the built model engines Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com> --- tests/integration/defs/cpp/test_multi_gpu.py | 138 ++++++++++++++---- .../test_lists/test-db/l0_dgx_h100.yml | 11 +- 2 files changed, 119 insertions(+), 30 deletions(-) diff --git a/tests/integration/defs/cpp/test_multi_gpu.py b/tests/integration/defs/cpp/test_multi_gpu.py index ee9860a6ce17..3099ae254652 100644 --- a/tests/integration/defs/cpp/test_multi_gpu.py +++ b/tests/integration/defs/cpp/test_multi_gpu.py @@ -77,7 +77,7 @@ def run_simple_multi_gpu_tests(build_dir: _pl.Path, timeout=1500): timeout=600) -def run_llama_executor_multi_gpu_tests(build_dir: _pl.Path, timeout=1500): +def run_llama_executor_leader_tests(build_dir: _pl.Path, timeout=1500): tests_dir = build_dir / "tests" cpp_env = {**_os.environ} @@ -94,12 +94,16 @@ def run_llama_executor_multi_gpu_tests(build_dir: _pl.Path, timeout=1500): "--gtest_filter=*LlamaExecutorTest*LeaderMode*:*LlamaMultiExecutorTest*LeaderMode*" ], leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - # https://nvbugspro.nvidia.com/bug/5026255 disable below tests for now. - if False: - _cpp.run_command(trt_model_test, - cwd=tests_dir, - env=mgpu_env, - timeout=1500) + + _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) + + +def run_llama_executor_orchestrator_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" #Executor test in orchestrator mode xml_output_file = build_dir / "results-multi-gpu-llama-exec-orch-mode.xml" @@ -110,18 +114,54 @@ def run_llama_executor_multi_gpu_tests(build_dir: _pl.Path, timeout=1500): ] _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) - #Logits processor and guided decoding test in leader mode + +def run_llama_executor_logits_proc_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" + + #Logits processor test in leader mode xml_output_file = build_dir / "results-multi-gpu-logits-proc.xml" + tp_pp_sizes = [(4, 1), (2, 2), (1, 4)] gtest_filter = [ f"LlamaExecutorTest/LogitsProcParamsTest*tp{tp}_pp{pp}*" for tp, pp in tp_pp_sizes ] - gtest_filter.extend([ + + gtest_filter = ":".join(gtest_filter) + + trt_model_test = _cpp.produce_mpirun_command( + global_commands=["mpirun", "--allow-run-as-root"], + nranks=4, + local_commands=[ + "executor/executorTest", f"--gtest_filter={gtest_filter}" + ], + leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) + + _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) + + +def run_llama_executor_guided_decoding_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" + + #Guided decoding test in leader mode + xml_output_file = build_dir / "results-multi-gpu-guided-decoding.xml" + + tp_pp_sizes = [(4, 1), (2, 2), (1, 4)] + gtest_filter = [ f"LlamaExecutorGuidedDecodingTest/GuidedDecodingParamsTest*tp{tp}_pp{pp}*" for tp, pp in tp_pp_sizes - ]) + ] + gtest_filter = ":".join(gtest_filter) + trt_model_test = _cpp.produce_mpirun_command( global_commands=["mpirun", "--allow-run-as-root"], nranks=4, @@ -129,10 +169,11 @@ def run_llama_executor_multi_gpu_tests(build_dir: _pl.Path, timeout=1500): "executor/executorTest", f"--gtest_filter={gtest_filter}" ], leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) + _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) -def run_t5_multi_gpu_tests(build_dir: _pl.Path, timeout=1500): +def run_enc_dec_multi_gpu_tests(build_dir: _pl.Path, timeout=1500): tests_dir = build_dir / "tests" cpp_env = {**_os.environ} @@ -399,6 +440,17 @@ def _prepare(model_name: str): return _prepare +# Use indirect parameterization to ensure that the model is built +# only once per pytest session +@pytest.fixture(scope="session") +def multi_gpu_model(request, prepare_model_multi_gpu): + + model_name = request.param + prepare_model_multi_gpu(model_name) + + return model_name + + @pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], indirect=True) def test_simple(build_google_tests, build_dir): @@ -411,46 +463,80 @@ def test_simple(build_google_tests, build_dir): @pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], indirect=True) -def test_t5(build_google_tests, prepare_model_multi_gpu, build_dir): +@pytest.mark.parametrize("multi_gpu_model", ["t5"], indirect=True) +def test_enc_dec(build_google_tests, multi_gpu_model, build_dir): if platform.system() != "Windows": - prepare_model_multi_gpu("t5") - run_t5_multi_gpu_tests(build_dir=build_dir, - timeout=_cpp.default_test_timeout) + run_enc_dec_multi_gpu_tests(build_dir=build_dir, + timeout=_cpp.default_test_timeout) @pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], indirect=True) -def test_llama_executor(build_google_tests, prepare_model_multi_gpu, lora_setup, +@pytest.mark.parametrize("mode", [ + "orchestrator", + pytest.param( + "leader", + marks=pytest.mark.skip("https://nvbugspro.nvidia.com/bug/5026255")) +]) +@pytest.mark.parametrize("multi_gpu_model", ["llama"], indirect=True) +def test_llama_executor(build_google_tests, multi_gpu_model, mode, lora_setup, build_dir): + if platform.system() == "Windows": + return + + if mode == "orchestrator": + run_llama_executor_orchestrator_tests(build_dir=build_dir, + timeout=_cpp.default_test_timeout) + elif mode == "leader": + run_llama_executor_leader_tests(build_dir=build_dir, + timeout=_cpp.default_test_timeout) + + +@pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], + indirect=True) +@pytest.mark.parametrize("multi_gpu_model", ["llama"], indirect=True) +def test_llama_executor_logits_proc(build_google_tests, multi_gpu_model, + lora_setup, build_dir): + + if platform.system() != "Windows": + run_llama_executor_logits_proc_tests(build_dir=build_dir, + timeout=_cpp.default_test_timeout) + + +@pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], + indirect=True) +@pytest.mark.parametrize("multi_gpu_model", ["llama"], indirect=True) +def test_llama_executor_guided_decoding(build_google_tests, multi_gpu_model, + lora_setup, build_dir): + if platform.system() != "Windows": - prepare_model_multi_gpu("llama") - run_llama_executor_multi_gpu_tests(build_dir=build_dir, - timeout=_cpp.default_test_timeout) + run_llama_executor_guided_decoding_tests( + build_dir=build_dir, timeout=_cpp.default_test_timeout) @pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], indirect=True) -def test_trt_gpt_real_decoder(build_google_tests, prepare_model_multi_gpu, - lora_setup, build_dir): +@pytest.mark.parametrize("multi_gpu_model", ["llama"], indirect=True) +def test_trt_gpt_real_decoder(build_google_tests, multi_gpu_model, lora_setup, + build_dir): if platform.system() != "Windows": - prepare_model_multi_gpu("llama") run_trt_gpt_model_real_decoder_multi_gpu_tests( build_dir=build_dir, timeout=_cpp.default_test_timeout) @pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], indirect=True) -def test_disagg(prepare_model, prepare_model_multi_gpu, build_google_tests, - build_dir): +@pytest.mark.parametrize("multi_gpu_model", ["llama"], indirect=True) +def test_disagg(prepare_model, multi_gpu_model, build_google_tests, build_dir): if platform.system() != "Windows": # Disagg tests need single + multi GPU llama models. - prepare_model("llama") - prepare_model_multi_gpu("llama") + # Disagg tests need only single GPU gpt model. + prepare_model("llama") prepare_model("gpt") run_disagg_multi_gpu_tests(build_dir=build_dir) diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index f8ba95fb838d..059f6f5c3dc0 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -115,10 +115,13 @@ l0_dgx_h100: tests: # ------------- CPP tests --------------- - cpp/test_multi_gpu.py::test_simple[90] - - cpp/test_multi_gpu.py::test_t5[90] - - cpp/test_multi_gpu.py::test_llama_executor[90] - - cpp/test_multi_gpu.py::test_trt_gpt_real_decoder[90] - - cpp/test_multi_gpu.py::test_disagg[90] + - cpp/test_multi_gpu.py::test_enc_dec[t5-90] + - cpp/test_multi_gpu.py::test_llama_executor[llama-orchestrator-90] + - cpp/test_multi_gpu.py::test_llama_executor[llama-leader-90] + - cpp/test_multi_gpu.py::test_llama_executor_guided_decoding[llama-90] + - cpp/test_multi_gpu.py::test_llama_logits_proc[llama-90] + - cpp/test_multi_gpu.py::test_trt_gpt_real_decoder[llama-90] + - cpp/test_multi_gpu.py::test_disagg[llama-90] - condition: ranges: system_gpu_count: From 31d28308c120f2484bb7af72639e3eb86aaf9226 Mon Sep 17 00:00:00 2001 From: Dom Brown <3886319+DomBrown@users.noreply.github.com> Date: Tue, 13 May 2025 11:53:20 +0000 Subject: [PATCH 03/11] Fix test list typo Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_dgx_h100.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index 059f6f5c3dc0..a527ddb87d60 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -119,7 +119,7 @@ l0_dgx_h100: - cpp/test_multi_gpu.py::test_llama_executor[llama-orchestrator-90] - cpp/test_multi_gpu.py::test_llama_executor[llama-leader-90] - cpp/test_multi_gpu.py::test_llama_executor_guided_decoding[llama-90] - - cpp/test_multi_gpu.py::test_llama_logits_proc[llama-90] + - cpp/test_multi_gpu.py::test_llama_executor_logits_proc[llama-90] - cpp/test_multi_gpu.py::test_trt_gpt_real_decoder[llama-90] - cpp/test_multi_gpu.py::test_disagg[llama-90] - condition: From 77465c6f07f97e6c504a1e7fc90ee33d7d898f97 Mon Sep 17 00:00:00 2001 From: Dom Brown <3886319+DomBrown@users.noreply.github.com> Date: Wed, 14 May 2025 16:08:21 +0000 Subject: [PATCH 04/11] WIP Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com> --- cpp/tests/executor/disaggExecutorTest.cpp | 7 +- tests/integration/defs/cpp/test_multi_gpu.py | 160 ++++++++++++------- 2 files changed, 111 insertions(+), 56 deletions(-) diff --git a/cpp/tests/executor/disaggExecutorTest.cpp b/cpp/tests/executor/disaggExecutorTest.cpp index 15c15319e1e9..3f1166932cb2 100644 --- a/cpp/tests/executor/disaggExecutorTest.cpp +++ b/cpp/tests/executor/disaggExecutorTest.cpp @@ -1116,8 +1116,11 @@ INSTANTIATE_TEST_SUITE_P(GptSingleDeviceDisaggSymmetricExecutorMixedTest, Disagg testing::Values(1)), generateTestNameDisaggParams); -INSTANTIATE_TEST_SUITE_P(ConditionalDisaggSymmetricExecutorTest, ConditionalDisaggParamsTest, - testing::Combine(testing::Values("gpt", "llama_tp1_pp1_cp1")), generateTestNameCondDisaggParams); +INSTANTIATE_TEST_SUITE_P(GptConditionalDisaggSymmetricExecutorTest, ConditionalDisaggParamsTest, + testing::Combine(testing::Values("gpt")), generateTestNameCondDisaggParams); + +INSTANTIATE_TEST_SUITE_P(LlamaConditionalDisaggSymmetricExecutorTest, ConditionalDisaggParamsTest, + testing::Combine(testing::Values("llama_tp1_pp1_cp1")), generateTestNameCondDisaggParams); INSTANTIATE_TEST_SUITE_P(LlamaTP2DisaggSymmetricExecutorTest, DisaggParamsTest, testing::Combine(testing::Values(4), diff --git a/tests/integration/defs/cpp/test_multi_gpu.py b/tests/integration/defs/cpp/test_multi_gpu.py index 3099ae254652..e13b87aac7f6 100644 --- a/tests/integration/defs/cpp/test_multi_gpu.py +++ b/tests/integration/defs/cpp/test_multi_gpu.py @@ -211,49 +211,81 @@ def run_trt_gpt_model_real_decoder_multi_gpu_tests(build_dir: _pl.Path, timeout=timeout) # expecting ~ 1200s -def run_disagg_multi_gpu_tests(build_dir: _pl.Path): - +def run_disagg_symmetric_executor_tests(build_dir: _pl.Path, + model: str, + use_ucx_kvcache=False, + timeout=1500): tests_dir = build_dir / "tests" cpp_env = {**_os.environ} + match model: + case "llama": + prefix = "Llama" + case "gpt": + prefix = "Gpt" + case _: + raise ValueError(f"Unsupported model: {model}") + new_env = copy.copy(cpp_env) - new_env["TRTLLM_USE_MPI_KVCACHE"] = "1" + if use_ucx_kvcache: + new_env["TRTLLM_USE_UCX_KVCACHE"] = "1" + else: + new_env["TRTLLM_USE_MPI_KVCACHE"] = "1" + xml_output_file = build_dir / "results-multi-gpu-disagg-executor-2-process.xml" trt_model_test = _cpp.produce_mpirun_command( global_commands=["mpirun", "--allow-run-as-root"], nranks=2, local_commands=[ "executor/disaggExecutorTest", - "--gtest_filter=*DisaggSymmetricExecutorTest*" + f"--gtest_filter=*{prefix}*DisaggSymmetricExecutorTest*" ], leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) _cpp.run_command(trt_model_test, cwd=tests_dir, env=new_env, timeout=1500) mgpu_env = copy.copy(cpp_env) mgpu_env["RUN_LLAMA_MULTI_GPU"] = "true" - mgpu_env["TRTLLM_USE_MPI_KVCACHE"] = "1" + + if use_ucx_kvcache: + mgpu_env["TRTLLM_USE_UCX_KVCACHE"] = "1" + else: + mgpu_env["TRTLLM_USE_MPI_KVCACHE"] = "1" + xml_output_file = build_dir / "results-multi-gpu-disagg-executor-4-process.xml" trt_model_test = _cpp.produce_mpirun_command( global_commands=["mpirun", "--allow-run-as-root"], nranks=4, local_commands=[ "executor/disaggExecutorTest", - "--gtest_filter=*DisaggSymmetricExecutorTest*" + f"--gtest_filter=*{prefix}*DisaggSymmetricExecutorTest*" ], leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - # https://nvbugspro.nvidia.com/bug/5026255 disable below tests for now. _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) - xml_output_file = build_dir / "results-multi-gpu-disagg-executor-8-process.xml" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=8, - local_commands=[ - "executor/disaggExecutorTest", - "--gtest_filter=*LlamaTP2PP2DisaggSymmetricExecutorTest*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) + if model == "llama": + xml_output_file = build_dir / "results-multi-gpu-disagg-executor-8-process.xml" + trt_model_test = _cpp.produce_mpirun_command( + global_commands=["mpirun", "--allow-run-as-root"], + nranks=8, + local_commands=[ + "executor/disaggExecutorTest", + "--gtest_filter=*LlamaTP2PP2DisaggSymmetricExecutorTest*" + ], + leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) + _cpp.run_command(trt_model_test, + cwd=tests_dir, + env=mgpu_env, + timeout=1500) + + +def run_disagg_multi_gpu_tests(build_dir: _pl.Path): + + tests_dir = build_dir / "tests" + cpp_env = {**_os.environ} + + mgpu_env = copy.copy(cpp_env) + mgpu_env["RUN_LLAMA_MULTI_GPU"] = "true" + mgpu_env["TRTLLM_USE_MPI_KVCACHE"] = "1" xml_output_file = build_dir / "results-multi-gpu-disagg-asymmetric-executor-4-process.xml" trt_model_test = _cpp.produce_mpirun_command( @@ -299,45 +331,9 @@ def run_disagg_multi_gpu_tests(build_dir: _pl.Path): leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) _cpp.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 - new_env = copy.copy(cpp_env) - new_env["TRTLLM_USE_UCX_KVCACHE"] = "1" - xml_output_file = build_dir / "results-multi-gpu-disagg-executor-2-process.xml" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=2, - local_commands=[ - "executor/disaggExecutorTest", - "--gtest_filter=*DisaggSymmetricExecutorTest*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, cwd=tests_dir, env=new_env, timeout=1500) - mgpu_env = copy.copy(cpp_env) mgpu_env["RUN_LLAMA_MULTI_GPU"] = "true" mgpu_env["TRTLLM_USE_UCX_KVCACHE"] = "1" - xml_output_file = build_dir / "results-multi-gpu-disagg-executor-4-process.xml" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=4, - local_commands=[ - "executor/disaggExecutorTest", - "--gtest_filter=*DisaggSymmetricExecutorTest*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - # https://nvbugspro.nvidia.com/bug/5026255 disable below tests for now. - _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) - - xml_output_file = build_dir / "results-multi-gpu-disagg-executor-8-process.xml" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=8, - local_commands=[ - "executor/disaggExecutorTest", - "--gtest_filter=*LlamaTP2PP2DisaggSymmetricExecutorTest*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) xml_output_file = build_dir / "results-multi-gpu-disagg-asymmetric-executor-4-process.xml" trt_model_test = _cpp.produce_mpirun_command( @@ -440,6 +436,39 @@ def _prepare(model_name: str): return _prepare +@pytest.fixture(scope="session") +def gpt_single_gpu_model(prepare_model): + prepare_model("gpt") + return "gpt" + + +@pytest.fixture(scope="session") +def llama_single_gpu_model(prepare_model): + prepare_model("llama") + return "llama" + + +@pytest.fixture(scope="session") +def llama_multi_gpu_model(prepare_model_multi_gpu): + prepare_model_multi_gpu("llama") + return "llama" + + +# Allow us to dynamically choose a fixture at runtime +# Combined with session scope to ensure that the model is built only once +# per pytest session +@pytest.fixture +def run_fixture(request): + + def _run(fixture_name): + fixture = request.getfixturevalue(fixture_name) + if fixture is None: + raise ValueError(f"Fixture {fixture_name} not found") + return fixture + + return _run + + # Use indirect parameterization to ensure that the model is built # only once per pytest session @pytest.fixture(scope="session") @@ -530,7 +559,7 @@ def test_trt_gpt_real_decoder(build_google_tests, multi_gpu_model, lora_setup, @pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], indirect=True) @pytest.mark.parametrize("multi_gpu_model", ["llama"], indirect=True) -def test_disagg(prepare_model, multi_gpu_model, build_google_tests, build_dir): +def test_disagg(build_google_tests, prepare_model, multi_gpu_model, build_dir): if platform.system() != "Windows": # Disagg tests need single + multi GPU llama models. @@ -540,3 +569,26 @@ def test_disagg(prepare_model, multi_gpu_model, build_google_tests, build_dir): prepare_model("gpt") run_disagg_multi_gpu_tests(build_dir=build_dir) + + +@pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], + indirect=True) +@pytest.mark.parametrize("use_ucx_kvcache", [False, True], + ids=["mpi_kvcache", "ucx_kvcache"]) +@pytest.mark.parametrize("model", ["gpt", "llama"]) +def test_disagg_symmetric_executor(build_google_tests, model, use_ucx_kvcache, + run_fixture, build_dir): + + if platform.system() != "Windows": + # Disagg tests need single + multi GPU llama models. + # Disagg tests need only single GPU gpt model. + + if model == "llama": + run_fixture("llama_single_gpu_model") + run_fixture("llama_multi_gpu_model") + elif model == "gpt": + run_fixture("gpt_single_gpu_model") + + run_disagg_symmetric_executor_tests(build_dir=build_dir, + model="gpt", + use_ucx_kvcache=use_ucx_kvcache) From c619cc908123a9d2467b7abc58902bc60997c357 Mon Sep 17 00:00:00 2001 From: Dom Brown <3886319+DomBrown@users.noreply.github.com> Date: Thu, 15 May 2025 10:57:08 +0000 Subject: [PATCH 05/11] Further WIP Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com> --- tests/integration/defs/cpp/test_multi_gpu.py | 233 ++++++++++++------- 1 file changed, 154 insertions(+), 79 deletions(-) diff --git a/tests/integration/defs/cpp/test_multi_gpu.py b/tests/integration/defs/cpp/test_multi_gpu.py index e13b87aac7f6..96134feffb52 100644 --- a/tests/integration/defs/cpp/test_multi_gpu.py +++ b/tests/integration/defs/cpp/test_multi_gpu.py @@ -9,6 +9,16 @@ import pytest +# Helper filter for disagg google tests +def get_model_test_filter_prefix(model: str) -> str: + if model == "llama": + return "Llama" + elif model == "gpt": + return "Gpt" + else: + raise ValueError(f"Unsupported model: {model}") + + def run_simple_multi_gpu_tests(build_dir: _pl.Path, timeout=1500): tests_dir = build_dir / "tests" cpp_env = {**_os.environ} @@ -218,13 +228,7 @@ def run_disagg_symmetric_executor_tests(build_dir: _pl.Path, tests_dir = build_dir / "tests" cpp_env = {**_os.environ} - match model: - case "llama": - prefix = "Llama" - case "gpt": - prefix = "Gpt" - case _: - raise ValueError(f"Unsupported model: {model}") + prefix = get_model_test_filter_prefix(model) new_env = copy.copy(cpp_env) if use_ucx_kvcache: @@ -241,7 +245,10 @@ def run_disagg_symmetric_executor_tests(build_dir: _pl.Path, f"--gtest_filter=*{prefix}*DisaggSymmetricExecutorTest*" ], leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, cwd=tests_dir, env=new_env, timeout=1500) + _cpp.run_command(trt_model_test, + cwd=tests_dir, + env=new_env, + timeout=timeout) mgpu_env = copy.copy(cpp_env) mgpu_env["RUN_LLAMA_MULTI_GPU"] = "true" @@ -260,7 +267,10 @@ def run_disagg_symmetric_executor_tests(build_dir: _pl.Path, f"--gtest_filter=*{prefix}*DisaggSymmetricExecutorTest*" ], leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) + _cpp.run_command(trt_model_test, + cwd=tests_dir, + env=mgpu_env, + timeout=timeout) if model == "llama": xml_output_file = build_dir / "results-multi-gpu-disagg-executor-8-process.xml" @@ -275,17 +285,26 @@ def run_disagg_symmetric_executor_tests(build_dir: _pl.Path, _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, - timeout=1500) + timeout=timeout) -def run_disagg_multi_gpu_tests(build_dir: _pl.Path): +def run_disagg_asymmetric_executor_tests(build_dir: _pl.Path, + model: str, + use_ucx_kvcache=False, + timeout=1500): tests_dir = build_dir / "tests" - cpp_env = {**_os.environ} + prefix = get_model_test_filter_prefix(model) + + cpp_env = {**_os.environ} mgpu_env = copy.copy(cpp_env) mgpu_env["RUN_LLAMA_MULTI_GPU"] = "true" - mgpu_env["TRTLLM_USE_MPI_KVCACHE"] = "1" + + if use_ucx_kvcache: + mgpu_env["TRTLLM_USE_UCX_KVCACHE"] = "1" + else: + mgpu_env["TRTLLM_USE_MPI_KVCACHE"] = "1" xml_output_file = build_dir / "results-multi-gpu-disagg-asymmetric-executor-4-process.xml" trt_model_test = _cpp.produce_mpirun_command( @@ -293,10 +312,13 @@ def run_disagg_multi_gpu_tests(build_dir: _pl.Path): nranks=4, local_commands=[ "executor/disaggExecutorTest", - "--gtest_filter=*DisaggAsymmetricExecutorTest*" + f"--gtest_filter=*{prefix}*DisaggAsymmetricExecutorTest*" ], leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) + _cpp.run_command(trt_model_test, + cwd=tests_dir, + env=mgpu_env, + timeout=timeout) xml_output_file = build_dir / "results-multi-gpu-disagg-asymmetric-executor-6-process.xml" trt_model_test = _cpp.produce_mpirun_command( @@ -304,10 +326,13 @@ def run_disagg_multi_gpu_tests(build_dir: _pl.Path): nranks=6, local_commands=[ "executor/disaggExecutorTest", - "--gtest_filter=*DisaggAsymmetricExecutorTest*" + f"--gtest_filter=*{prefix}*DisaggAsymmetricExecutorTest*" ], leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) + _cpp.run_command(trt_model_test, + cwd=tests_dir, + env=mgpu_env, + timeout=timeout) xml_output_file = build_dir / "results-multi-gpu-disagg-asymmetric-executor-8-process.xml" trt_model_test = _cpp.produce_mpirun_command( @@ -315,58 +340,32 @@ def run_disagg_multi_gpu_tests(build_dir: _pl.Path): nranks=8, local_commands=[ "executor/disaggExecutorTest", - "--gtest_filter=*DisaggAsymmetricExecutorTest*" + f"--gtest_filter=*{prefix}*DisaggAsymmetricExecutorTest*" ], leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) + _cpp.run_command(trt_model_test, + cwd=tests_dir, + env=mgpu_env, + timeout=timeout) - xml_output_file = build_dir / "results-multi-gpu-disagg-asymmetric-orchestrator-executor-7-process.xml" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=7, - local_commands=[ - "executor/disaggExecutorTest", - "--gtest_filter=*DisaggOrchestratorParamsTest*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) - mgpu_env = copy.copy(cpp_env) - mgpu_env["RUN_LLAMA_MULTI_GPU"] = "true" - mgpu_env["TRTLLM_USE_UCX_KVCACHE"] = "1" +def run_disagg_orchestrator_params_tests(build_dir: _pl.Path, + model: str, + use_ucx_kvcache=False, + timeout=1500): - xml_output_file = build_dir / "results-multi-gpu-disagg-asymmetric-executor-4-process.xml" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=4, - local_commands=[ - "executor/disaggExecutorTest", - "--gtest_filter=*DisaggAsymmetricExecutorTest*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) + tests_dir = build_dir / "tests" - xml_output_file = build_dir / "results-multi-gpu-disagg-asymmetric-executor-6-process.xml" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=6, - local_commands=[ - "executor/disaggExecutorTest", - "--gtest_filter=*DisaggAsymmetricExecutorTest*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) + prefix = get_model_test_filter_prefix(model) - xml_output_file = build_dir / "results-multi-gpu-disagg-asymmetric-executor-8-process.xml" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=8, - local_commands=[ - "executor/disaggExecutorTest", - "--gtest_filter=*DisaggAsymmetricExecutorTest*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) + cpp_env = {**_os.environ} + mgpu_env = copy.copy(cpp_env) + mgpu_env["RUN_LLAMA_MULTI_GPU"] = "true" + + if use_ucx_kvcache: + mgpu_env["TRTLLM_USE_UCX_KVCACHE"] = "1" + else: + mgpu_env["TRTLLM_USE_MPI_KVCACHE"] = "1" xml_output_file = build_dir / "results-multi-gpu-disagg-asymmetric-orchestrator-executor-7-process.xml" trt_model_test = _cpp.produce_mpirun_command( @@ -374,17 +373,41 @@ def run_disagg_multi_gpu_tests(build_dir: _pl.Path): nranks=7, local_commands=[ "executor/disaggExecutorTest", - "--gtest_filter=*DisaggOrchestratorParamsTest*" + f"--gtest_filter=*{prefix}*DisaggOrchestratorParamsTest*" ], leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) + _cpp.run_command(trt_model_test, + cwd=tests_dir, + env=mgpu_env, + timeout=timeout) + + +def run_disagg_spawn_orchestrator_tests(build_dir: _pl.Path, + model: str, + use_ucx_kvcache=False, + timeout=1500): + + tests_dir = build_dir / "tests" + + prefix = get_model_test_filter_prefix(model) + + cpp_env = {**_os.environ} + mgpu_env = copy.copy(cpp_env) + mgpu_env["RUN_LLAMA_MULTI_GPU"] = "true" + + if use_ucx_kvcache: + mgpu_env["TRTLLM_USE_UCX_KVCACHE"] = "1" + else: + mgpu_env["TRTLLM_USE_MPI_KVCACHE"] = "1" xml_output_file = build_dir / "results-multi-gpu-disagg-spawn-asymmetric-orchestrator-executor-1-process.xml" + comms = [ "executor/disaggExecutorTest", - "--gtest_filter=*DisaaggSpawnOrchestrator*" + f"--gtest_filter=*{prefix}*DisaaggSpawnOrchestrator*", + f"--gtest_output=xml:{xml_output_file}" ] - _cpp.run_command(comms, cwd=tests_dir, env=mgpu_env, timeout=1500) + _cpp.run_command(comms, cwd=tests_dir, env=mgpu_env, timeout=timeout) def prepare_multi_gpu_model_tests(test_list: List[str], @@ -521,6 +544,8 @@ def test_llama_executor(build_google_tests, multi_gpu_model, mode, lora_setup, elif mode == "leader": run_llama_executor_leader_tests(build_dir=build_dir, timeout=_cpp.default_test_timeout) + else: + raise ValueError(f"Unsupported mode: {mode}") @pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], @@ -558,37 +583,87 @@ def test_trt_gpt_real_decoder(build_google_tests, multi_gpu_model, lora_setup, @pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], indirect=True) -@pytest.mark.parametrize("multi_gpu_model", ["llama"], indirect=True) -def test_disagg(build_google_tests, prepare_model, multi_gpu_model, build_dir): +@pytest.mark.parametrize("use_ucx_kvcache", [False, True], + ids=["mpi_kvcache", "ucx_kvcache"]) +@pytest.mark.parametrize("model", ["gpt", "llama"]) +def test_disagg_symmetric_executor(build_google_tests, model, use_ucx_kvcache, + run_fixture, build_dir): if platform.system() != "Windows": # Disagg tests need single + multi GPU llama models. # Disagg tests need only single GPU gpt model. - prepare_model("llama") - prepare_model("gpt") + if model == "llama": + run_fixture("llama_single_gpu_model") + run_fixture("llama_multi_gpu_model") + elif model == "gpt": + run_fixture("gpt_single_gpu_model") + else: + raise ValueError(f"Unsupported model: {model}") - run_disagg_multi_gpu_tests(build_dir=build_dir) + run_disagg_symmetric_executor_tests(build_dir=build_dir, + model=model, + use_ucx_kvcache=use_ucx_kvcache) @pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], indirect=True) @pytest.mark.parametrize("use_ucx_kvcache", [False, True], ids=["mpi_kvcache", "ucx_kvcache"]) -@pytest.mark.parametrize("model", ["gpt", "llama"]) -def test_disagg_symmetric_executor(build_google_tests, model, use_ucx_kvcache, - run_fixture, build_dir): +@pytest.mark.parametrize("model", ["llama"]) +def test_disagg_asymmetric_executor(build_google_tests, model, use_ucx_kvcache, + run_fixture, build_dir): if platform.system() != "Windows": # Disagg tests need single + multi GPU llama models. - # Disagg tests need only single GPU gpt model. if model == "llama": run_fixture("llama_single_gpu_model") run_fixture("llama_multi_gpu_model") - elif model == "gpt": - run_fixture("gpt_single_gpu_model") + else: + raise ValueError(f"Unsupported model: {model}") - run_disagg_symmetric_executor_tests(build_dir=build_dir, - model="gpt", + run_disagg_asymmetric_executor_tests(build_dir=build_dir, + model=model, + use_ucx_kvcache=use_ucx_kvcache) + + +@pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], + indirect=True) +@pytest.mark.parametrize("use_ucx_kvcache", [False, True], + ids=["mpi_kvcache", "ucx_kvcache"]) +@pytest.mark.parametrize("model", ["llama"]) +def test_disagg_orchestrator_params(build_google_tests, model, use_ucx_kvcache, + run_fixture, build_dir): + + if platform.system() != "Windows": + # Disagg tests need single + multi GPU llama models. + if model == "llama": + run_fixture("llama_single_gpu_model") + run_fixture("llama_multi_gpu_model") + else: + raise ValueError(f"Unsupported model: {model}") + + run_disagg_orchestrator_params_tests(build_dir=build_dir, + model=model, + use_ucx_kvcache=use_ucx_kvcache) + + +@pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], + indirect=True) +@pytest.mark.parametrize("use_ucx_kvcache", [True], ids=["ucx_kvcache"]) +@pytest.mark.parametrize("model", ["llama"]) +def test_disagg_spawn_orchestrator(build_google_tests, model, use_ucx_kvcache, + run_fixture, build_dir): + + if platform.system() != "Windows": + # Disagg tests need single + multi GPU llama models. + if model == "llama": + run_fixture("llama_single_gpu_model") + run_fixture("llama_multi_gpu_model") + else: + raise ValueError(f"Unsupported model: {model}") + + run_disagg_spawn_orchestrator_tests(build_dir=build_dir, + model=model, use_ucx_kvcache=use_ucx_kvcache) From a9db385b13a740d39d1c293292680f3a2815ac0b Mon Sep 17 00:00:00 2001 From: Dom Brown <3886319+DomBrown@users.noreply.github.com> Date: Thu, 15 May 2025 15:15:01 +0000 Subject: [PATCH 06/11] WIP Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com> --- tests/integration/defs/cpp/test_multi_gpu.py | 148 +++++++++---------- 1 file changed, 67 insertions(+), 81 deletions(-) diff --git a/tests/integration/defs/cpp/test_multi_gpu.py b/tests/integration/defs/cpp/test_multi_gpu.py index 96134feffb52..790261620da7 100644 --- a/tests/integration/defs/cpp/test_multi_gpu.py +++ b/tests/integration/defs/cpp/test_multi_gpu.py @@ -478,18 +478,30 @@ def llama_multi_gpu_model(prepare_model_multi_gpu): # Allow us to dynamically choose a fixture at runtime -# Combined with session scope to ensure that the model is built only once -# per pytest session +# Combined with session scope fixtures above to ensure +# that the model is built only once per pytest session @pytest.fixture -def run_fixture(request): +def prepare_models_disagg(request): - def _run(fixture_name): - fixture = request.getfixturevalue(fixture_name) - if fixture is None: - raise ValueError(f"Fixture {fixture_name} not found") - return fixture + def _prepare(model_name: str): + if model_name == "llama": + fixture_names = [ + "llama_single_gpu_model", + "llama_multi_gpu_model", + ] + elif model_name == "gpt": + fixture_names = [ + "gpt_single_gpu_model", + ] + else: + raise ValueError(f"Disagg tests don't support model: {model_name}") - return _run + print(f"Preparing models for disagg tests: {fixture_names}") + # Run the fixtures + for fixture_name in fixture_names: + request.getfixturevalue(fixture_name) + + return _prepare # Use indirect parameterization to ensure that the model is built @@ -583,87 +595,61 @@ def test_trt_gpt_real_decoder(build_google_tests, multi_gpu_model, lora_setup, @pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], indirect=True) -@pytest.mark.parametrize("use_ucx_kvcache", [False, True], - ids=["mpi_kvcache", "ucx_kvcache"]) -@pytest.mark.parametrize("model", ["gpt", "llama"]) -def test_disagg_symmetric_executor(build_google_tests, model, use_ucx_kvcache, - run_fixture, build_dir): - - if platform.system() != "Windows": - # Disagg tests need single + multi GPU llama models. - # Disagg tests need only single GPU gpt model. - - if model == "llama": - run_fixture("llama_single_gpu_model") - run_fixture("llama_multi_gpu_model") - elif model == "gpt": - run_fixture("gpt_single_gpu_model") - else: - raise ValueError(f"Unsupported model: {model}") +class TestDisagg: - run_disagg_symmetric_executor_tests(build_dir=build_dir, - model=model, - use_ucx_kvcache=use_ucx_kvcache) + @pytest.mark.parametrize("use_ucx_kvcache", [False, True], + ids=["mpi_kvcache", "ucx_kvcache"]) + @pytest.mark.parametrize("model", ["gpt", "llama"]) + def test_symmetric_executor(self, build_google_tests, model, + use_ucx_kvcache, prepare_models_disagg, + build_dir): + if platform.system() != "Windows": + prepare_models_disagg(model) -@pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], - indirect=True) -@pytest.mark.parametrize("use_ucx_kvcache", [False, True], - ids=["mpi_kvcache", "ucx_kvcache"]) -@pytest.mark.parametrize("model", ["llama"]) -def test_disagg_asymmetric_executor(build_google_tests, model, use_ucx_kvcache, - run_fixture, build_dir): - - if platform.system() != "Windows": - # Disagg tests need single + multi GPU llama models. + run_disagg_symmetric_executor_tests(build_dir=build_dir, + model=model, + use_ucx_kvcache=use_ucx_kvcache) - if model == "llama": - run_fixture("llama_single_gpu_model") - run_fixture("llama_multi_gpu_model") - else: - raise ValueError(f"Unsupported model: {model}") - - run_disagg_asymmetric_executor_tests(build_dir=build_dir, - model=model, - use_ucx_kvcache=use_ucx_kvcache) + @pytest.mark.parametrize("use_ucx_kvcache", [False, True], + ids=["mpi_kvcache", "ucx_kvcache"]) + @pytest.mark.parametrize("model", ["llama"]) + def test_asymmetric_executor(self, build_google_tests, model, + use_ucx_kvcache, prepare_models_disagg, + build_dir): + if platform.system() != "Windows": + prepare_models_disagg(model_name=model) -@pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], - indirect=True) -@pytest.mark.parametrize("use_ucx_kvcache", [False, True], - ids=["mpi_kvcache", "ucx_kvcache"]) -@pytest.mark.parametrize("model", ["llama"]) -def test_disagg_orchestrator_params(build_google_tests, model, use_ucx_kvcache, - run_fixture, build_dir): + run_disagg_asymmetric_executor_tests( + build_dir=build_dir, + model=model, + use_ucx_kvcache=use_ucx_kvcache) - if platform.system() != "Windows": - # Disagg tests need single + multi GPU llama models. - if model == "llama": - run_fixture("llama_single_gpu_model") - run_fixture("llama_multi_gpu_model") - else: - raise ValueError(f"Unsupported model: {model}") + @pytest.mark.parametrize("use_ucx_kvcache", [False, True], + ids=["mpi_kvcache", "ucx_kvcache"]) + @pytest.mark.parametrize("model", ["llama"]) + def test_orchestrator_params(self, build_google_tests, model, + use_ucx_kvcache, prepare_models_disagg, + build_dir): - run_disagg_orchestrator_params_tests(build_dir=build_dir, - model=model, - use_ucx_kvcache=use_ucx_kvcache) + if platform.system() != "Windows": + prepare_models_disagg(model) + run_disagg_orchestrator_params_tests( + build_dir=build_dir, + model=model, + use_ucx_kvcache=use_ucx_kvcache) -@pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], - indirect=True) -@pytest.mark.parametrize("use_ucx_kvcache", [True], ids=["ucx_kvcache"]) -@pytest.mark.parametrize("model", ["llama"]) -def test_disagg_spawn_orchestrator(build_google_tests, model, use_ucx_kvcache, - run_fixture, build_dir): + @pytest.mark.parametrize("use_ucx_kvcache", [True], ids=["ucx_kvcache"]) + @pytest.mark.parametrize("model", ["llama"]) + def test_spawn_orchestrator(self, build_google_tests, model, + use_ucx_kvcache, prepare_models_disagg, + build_dir): - if platform.system() != "Windows": - # Disagg tests need single + multi GPU llama models. - if model == "llama": - run_fixture("llama_single_gpu_model") - run_fixture("llama_multi_gpu_model") - else: - raise ValueError(f"Unsupported model: {model}") + if platform.system() != "Windows": + prepare_models_disagg(model) - run_disagg_spawn_orchestrator_tests(build_dir=build_dir, - model=model, - use_ucx_kvcache=use_ucx_kvcache) + run_disagg_spawn_orchestrator_tests(build_dir=build_dir, + model=model, + use_ucx_kvcache=use_ucx_kvcache) From f8329568a798844ee7b1d9d8d96901d66273ff50 Mon Sep 17 00:00:00 2001 From: Dom Brown <3886319+DomBrown@users.noreply.github.com> Date: Thu, 15 May 2025 15:23:38 +0000 Subject: [PATCH 07/11] Update test lists and readme Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com> --- cpp/tests/README.md | 2 +- tests/integration/test_lists/test-db/l0_dgx_h100.yml | 10 +++++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/cpp/tests/README.md b/cpp/tests/README.md index f12a9e73cba6..8ce1c354afe0 100644 --- a/cpp/tests/README.md +++ b/cpp/tests/README.md @@ -31,7 +31,7 @@ pytest tests/integration/defs/cpp/test_e2e.py::test_model[llama-90] pytest tests/integration/defs/cpp/test_e2e.py::test_benchmarks[gpt-90] -pytest tests/integration/defs/cpp/test_multi_gpu.py::test_disagg[90] +pytest tests/integration/defs/cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[gpt-mpi_kvcache-90] ``` ## Manual steps diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index a527ddb87d60..e0c92e2d66e6 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -121,7 +121,15 @@ l0_dgx_h100: - cpp/test_multi_gpu.py::test_llama_executor_guided_decoding[llama-90] - cpp/test_multi_gpu.py::test_llama_executor_logits_proc[llama-90] - cpp/test_multi_gpu.py::test_trt_gpt_real_decoder[llama-90] - - cpp/test_multi_gpu.py::test_disagg[llama-90] + - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[gpt-mpi_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[gpt-ucx_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-mpi_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-ucx_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-mpi_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-ucx_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_orchestrator_params[llama-mpi_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_orchestrator_params[llama-ucx_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_spawn_orchestrator[llama-ucx_kvcache-90] - condition: ranges: system_gpu_count: From b27e09c9e8397542b07e7bef58e2732a3af130e4 Mon Sep 17 00:00:00 2001 From: Dom Brown <3886319+DomBrown@users.noreply.github.com> Date: Fri, 16 May 2025 10:20:33 +0000 Subject: [PATCH 08/11] Try parametrize for asymmetric Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com> --- tests/integration/defs/cpp/test_multi_gpu.py | 36 ++++---------------- 1 file changed, 7 insertions(+), 29 deletions(-) diff --git a/tests/integration/defs/cpp/test_multi_gpu.py b/tests/integration/defs/cpp/test_multi_gpu.py index 790261620da7..95a8baba823f 100644 --- a/tests/integration/defs/cpp/test_multi_gpu.py +++ b/tests/integration/defs/cpp/test_multi_gpu.py @@ -290,6 +290,7 @@ def run_disagg_symmetric_executor_tests(build_dir: _pl.Path, def run_disagg_asymmetric_executor_tests(build_dir: _pl.Path, model: str, + nprocs=4, use_ucx_kvcache=False, timeout=1500): @@ -306,43 +307,17 @@ def run_disagg_asymmetric_executor_tests(build_dir: _pl.Path, else: mgpu_env["TRTLLM_USE_MPI_KVCACHE"] = "1" - xml_output_file = build_dir / "results-multi-gpu-disagg-asymmetric-executor-4-process.xml" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=4, - local_commands=[ - "executor/disaggExecutorTest", - f"--gtest_filter=*{prefix}*DisaggAsymmetricExecutorTest*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, - cwd=tests_dir, - env=mgpu_env, - timeout=timeout) + xml_output_file = build_dir / f"results-multi-gpu-disagg-asymmetric-executor-{nprocs}-process.xml" - xml_output_file = build_dir / "results-multi-gpu-disagg-asymmetric-executor-6-process.xml" trt_model_test = _cpp.produce_mpirun_command( global_commands=["mpirun", "--allow-run-as-root"], - nranks=6, + nranks=nprocs, local_commands=[ "executor/disaggExecutorTest", f"--gtest_filter=*{prefix}*DisaggAsymmetricExecutorTest*" ], leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, - cwd=tests_dir, - env=mgpu_env, - timeout=timeout) - xml_output_file = build_dir / "results-multi-gpu-disagg-asymmetric-executor-8-process.xml" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=8, - local_commands=[ - "executor/disaggExecutorTest", - f"--gtest_filter=*{prefix}*DisaggAsymmetricExecutorTest*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, @@ -613,8 +588,10 @@ def test_symmetric_executor(self, build_google_tests, model, @pytest.mark.parametrize("use_ucx_kvcache", [False, True], ids=["mpi_kvcache", "ucx_kvcache"]) + @pytest.mark.parametrize("nprocs", [4, 6, 8], + ids=["4proc", "6proc", "8proc"]) @pytest.mark.parametrize("model", ["llama"]) - def test_asymmetric_executor(self, build_google_tests, model, + def test_asymmetric_executor(self, build_google_tests, model, nprocs, use_ucx_kvcache, prepare_models_disagg, build_dir): @@ -624,6 +601,7 @@ def test_asymmetric_executor(self, build_google_tests, model, run_disagg_asymmetric_executor_tests( build_dir=build_dir, model=model, + nprocs=nprocs, use_ucx_kvcache=use_ucx_kvcache) @pytest.mark.parametrize("use_ucx_kvcache", [False, True], From 3f6bd083d7c3553ea167120f22b49c568c60a166 Mon Sep 17 00:00:00 2001 From: domb <3886319+DomBrown@users.noreply.github.com> Date: Fri, 16 May 2025 14:06:58 +0000 Subject: [PATCH 09/11] Parametrize + skip unsupported combinations Signed-off-by: domb <3886319+DomBrown@users.noreply.github.com> --- tests/integration/defs/cpp/test_multi_gpu.py | 54 +++++--------------- 1 file changed, 14 insertions(+), 40 deletions(-) diff --git a/tests/integration/defs/cpp/test_multi_gpu.py b/tests/integration/defs/cpp/test_multi_gpu.py index 95a8baba823f..ce40721571eb 100644 --- a/tests/integration/defs/cpp/test_multi_gpu.py +++ b/tests/integration/defs/cpp/test_multi_gpu.py @@ -223,6 +223,7 @@ def run_trt_gpt_model_real_decoder_multi_gpu_tests(build_dir: _pl.Path, def run_disagg_symmetric_executor_tests(build_dir: _pl.Path, model: str, + nprocs=2, use_ucx_kvcache=False, timeout=1500): tests_dir = build_dir / "tests" @@ -230,63 +231,29 @@ def run_disagg_symmetric_executor_tests(build_dir: _pl.Path, prefix = get_model_test_filter_prefix(model) - new_env = copy.copy(cpp_env) - if use_ucx_kvcache: - new_env["TRTLLM_USE_UCX_KVCACHE"] = "1" - else: - new_env["TRTLLM_USE_MPI_KVCACHE"] = "1" - - xml_output_file = build_dir / "results-multi-gpu-disagg-executor-2-process.xml" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=2, - local_commands=[ - "executor/disaggExecutorTest", - f"--gtest_filter=*{prefix}*DisaggSymmetricExecutorTest*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, - cwd=tests_dir, - env=new_env, - timeout=timeout) - mgpu_env = copy.copy(cpp_env) - mgpu_env["RUN_LLAMA_MULTI_GPU"] = "true" - if use_ucx_kvcache: mgpu_env["TRTLLM_USE_UCX_KVCACHE"] = "1" else: mgpu_env["TRTLLM_USE_MPI_KVCACHE"] = "1" - xml_output_file = build_dir / "results-multi-gpu-disagg-executor-4-process.xml" + mgpu_env["RUN_LLAMA_MULTI_GPU"] = "true" + + xml_output_file = build_dir / f"results-multi-gpu-disagg-executor-{nprocs}-process.xml" trt_model_test = _cpp.produce_mpirun_command( global_commands=["mpirun", "--allow-run-as-root"], - nranks=4, + nranks=nprocs, local_commands=[ "executor/disaggExecutorTest", f"--gtest_filter=*{prefix}*DisaggSymmetricExecutorTest*" ], leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) + _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=timeout) - if model == "llama": - xml_output_file = build_dir / "results-multi-gpu-disagg-executor-8-process.xml" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=8, - local_commands=[ - "executor/disaggExecutorTest", - "--gtest_filter=*LlamaTP2PP2DisaggSymmetricExecutorTest*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, - cwd=tests_dir, - env=mgpu_env, - timeout=timeout) - def run_disagg_asymmetric_executor_tests(build_dir: _pl.Path, model: str, @@ -574,16 +541,23 @@ class TestDisagg: @pytest.mark.parametrize("use_ucx_kvcache", [False, True], ids=["mpi_kvcache", "ucx_kvcache"]) + @pytest.mark.parametrize("nprocs", [2, 4, 8], + ids=["2proc", "4proc", "8proc"]) @pytest.mark.parametrize("model", ["gpt", "llama"]) - def test_symmetric_executor(self, build_google_tests, model, + def test_symmetric_executor(self, build_google_tests, model, nprocs, use_ucx_kvcache, prepare_models_disagg, build_dir): + if model == "gpt" and nprocs > 2: + pytest.skip( + "test_symmetric_executor only supports 2 processes for gpt") + if platform.system() != "Windows": prepare_models_disagg(model) run_disagg_symmetric_executor_tests(build_dir=build_dir, model=model, + nprocs=nprocs, use_ucx_kvcache=use_ucx_kvcache) @pytest.mark.parametrize("use_ucx_kvcache", [False, True], From 6e7010c153b71aa78e202261ca8a3ffbff8f04eb Mon Sep 17 00:00:00 2001 From: domb <3886319+DomBrown@users.noreply.github.com> Date: Fri, 16 May 2025 14:18:20 +0000 Subject: [PATCH 10/11] Update test list Signed-off-by: domb <3886319+DomBrown@users.noreply.github.com> --- .../test_lists/test-db/l0_dgx_h100.yml | 20 +++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index e0c92e2d66e6..7b3adaf22198 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -121,12 +121,20 @@ l0_dgx_h100: - cpp/test_multi_gpu.py::test_llama_executor_guided_decoding[llama-90] - cpp/test_multi_gpu.py::test_llama_executor_logits_proc[llama-90] - cpp/test_multi_gpu.py::test_trt_gpt_real_decoder[llama-90] - - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[gpt-mpi_kvcache-90] - - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[gpt-ucx_kvcache-90] - - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-mpi_kvcache-90] - - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-ucx_kvcache-90] - - cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-mpi_kvcache-90] - - cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-ucx_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[gpt-2proc-mpi_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[gpt-2proc-ucx_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-2proc-mpi_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-4proc-mpi_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-8proc-mpi_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-2proc-ucx_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-4proc-ucx_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-8proc-ucx_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-4proc-mpi_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-6proc-mpi_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-8proc-mpi_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-4proc-ucx_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-6proc-ucx_kvcache-90] + - cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-8proc-ucx_kvcache-90] - cpp/test_multi_gpu.py::TestDisagg::test_orchestrator_params[llama-mpi_kvcache-90] - cpp/test_multi_gpu.py::TestDisagg::test_orchestrator_params[llama-ucx_kvcache-90] - cpp/test_multi_gpu.py::TestDisagg::test_spawn_orchestrator[llama-ucx_kvcache-90] From e4998ee46f02961ee0cc8acd4fbf654d41472146 Mon Sep 17 00:00:00 2001 From: domb <3886319+DomBrown@users.noreply.github.com> Date: Fri, 16 May 2025 15:37:26 +0000 Subject: [PATCH 11/11] Reduce environment duplicated code Signed-off-by: domb <3886319+DomBrown@users.noreply.github.com> --- tests/integration/defs/cpp/test_multi_gpu.py | 150 +++++++++---------- 1 file changed, 67 insertions(+), 83 deletions(-) diff --git a/tests/integration/defs/cpp/test_multi_gpu.py b/tests/integration/defs/cpp/test_multi_gpu.py index ce40721571eb..be15906cae90 100644 --- a/tests/integration/defs/cpp/test_multi_gpu.py +++ b/tests/integration/defs/cpp/test_multi_gpu.py @@ -1,8 +1,8 @@ -import copy import os as _os import pathlib as _pl import platform import time +from enum import Enum, auto from typing import List, Optional import defs.cpp.cpp_common as _cpp @@ -19,6 +19,31 @@ def get_model_test_filter_prefix(model: str) -> str: raise ValueError(f"Unsupported model: {model}") +class KVCacheType(Enum): + NONE = auto() + MPI = auto() + UCX = auto() + + +def get_multi_gpu_env(kv_cache_type=KVCacheType.NONE, llama_multi_gpu=False): + env = {**_os.environ} + + match kv_cache_type: + case KVCacheType.MPI: + env["TRTLLM_USE_MPI_KVCACHE"] = "1" + case KVCacheType.UCX: + env["TRTLLM_USE_UCX_KVCACHE"] = "1" + case KVCacheType.NONE: + pass + case _: + raise ValueError(f"Unsupported KVCacheType: {kv_cache_type}") + + if llama_multi_gpu: + env["RUN_LLAMA_MULTI_GPU"] = "true" + + return env + + def run_simple_multi_gpu_tests(build_dir: _pl.Path, timeout=1500): tests_dir = build_dir / "tests" cpp_env = {**_os.environ} @@ -32,9 +57,9 @@ def run_simple_multi_gpu_tests(build_dir: _pl.Path, timeout=1500): ] _cpp.run_command(mpi_utils_test, cwd=tests_dir, env=cpp_env, timeout=300) - # Cache transceiver tests - new_env = copy.copy(cpp_env) - new_env["TRTLLM_USE_MPI_KVCACHE"] = "1" + # Cache transceiver MPI tests + new_env = get_multi_gpu_env(kv_cache_type=KVCacheType.MPI) + cache_trans_test = [ "mpirun", "-n", @@ -44,9 +69,6 @@ def run_simple_multi_gpu_tests(build_dir: _pl.Path, timeout=1500): ] _cpp.run_command(cache_trans_test, cwd=tests_dir, env=new_env, timeout=300) - new_env = copy.copy(cpp_env) - new_env["TRTLLM_USE_MPI_KVCACHE"] = "1" - # Cache transceiver tests cache_trans_test_8_proc = [ "mpirun", "-n", @@ -60,8 +82,8 @@ def run_simple_multi_gpu_tests(build_dir: _pl.Path, timeout=1500): timeout=600) # Cache transceiver tests with UCX - new_env = copy.copy(cpp_env) - new_env["TRTLLM_USE_UCX_KVCACHE"] = "1" + new_env = get_multi_gpu_env(kv_cache_type=KVCacheType.UCX) + cache_trans_test = [ "mpirun", "-n", @@ -71,8 +93,6 @@ def run_simple_multi_gpu_tests(build_dir: _pl.Path, timeout=1500): ] _cpp.run_command(cache_trans_test, cwd=tests_dir, env=new_env, timeout=300) - new_env = copy.copy(cpp_env) - new_env["TRTLLM_USE_UCX_KVCACHE"] = "1" # Cache transceiver tests cache_trans_test_8_proc = [ "mpirun", @@ -89,10 +109,8 @@ def run_simple_multi_gpu_tests(build_dir: _pl.Path, timeout=1500): def run_llama_executor_leader_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" + mgpu_env = get_multi_gpu_env(llama_multi_gpu=True) #Executor test in leader mode xml_output_file = build_dir / "results-multi-gpu-llama-exec-leader-mode.xml" @@ -110,10 +128,8 @@ def run_llama_executor_leader_tests(build_dir: _pl.Path, timeout=1500): def run_llama_executor_orchestrator_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" + mgpu_env = get_multi_gpu_env(llama_multi_gpu=True) #Executor test in orchestrator mode xml_output_file = build_dir / "results-multi-gpu-llama-exec-orch-mode.xml" @@ -127,10 +143,8 @@ def run_llama_executor_orchestrator_tests(build_dir: _pl.Path, timeout=1500): def run_llama_executor_logits_proc_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" + mgpu_env = get_multi_gpu_env(llama_multi_gpu=True) #Logits processor test in leader mode xml_output_file = build_dir / "results-multi-gpu-logits-proc.xml" @@ -156,10 +170,8 @@ def run_llama_executor_logits_proc_tests(build_dir: _pl.Path, timeout=1500): def run_llama_executor_guided_decoding_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" + mgpu_env = get_multi_gpu_env(llama_multi_gpu=True) #Guided decoding test in leader mode xml_output_file = build_dir / "results-multi-gpu-guided-decoding.xml" @@ -224,20 +236,14 @@ def run_trt_gpt_model_real_decoder_multi_gpu_tests(build_dir: _pl.Path, def run_disagg_symmetric_executor_tests(build_dir: _pl.Path, model: str, nprocs=2, - use_ucx_kvcache=False, + kvcache_type=KVCacheType.MPI, timeout=1500): tests_dir = build_dir / "tests" - cpp_env = {**_os.environ} prefix = get_model_test_filter_prefix(model) - mgpu_env = copy.copy(cpp_env) - if use_ucx_kvcache: - mgpu_env["TRTLLM_USE_UCX_KVCACHE"] = "1" - else: - mgpu_env["TRTLLM_USE_MPI_KVCACHE"] = "1" - - mgpu_env["RUN_LLAMA_MULTI_GPU"] = "true" + mgpu_env = get_multi_gpu_env(kv_cache_type=kvcache_type, + llama_multi_gpu=True) xml_output_file = build_dir / f"results-multi-gpu-disagg-executor-{nprocs}-process.xml" trt_model_test = _cpp.produce_mpirun_command( @@ -258,21 +264,15 @@ def run_disagg_symmetric_executor_tests(build_dir: _pl.Path, def run_disagg_asymmetric_executor_tests(build_dir: _pl.Path, model: str, nprocs=4, - use_ucx_kvcache=False, + kvcache_type=KVCacheType.MPI, timeout=1500): tests_dir = build_dir / "tests" prefix = get_model_test_filter_prefix(model) - cpp_env = {**_os.environ} - mgpu_env = copy.copy(cpp_env) - mgpu_env["RUN_LLAMA_MULTI_GPU"] = "true" - - if use_ucx_kvcache: - mgpu_env["TRTLLM_USE_UCX_KVCACHE"] = "1" - else: - mgpu_env["TRTLLM_USE_MPI_KVCACHE"] = "1" + mgpu_env = get_multi_gpu_env(kv_cache_type=kvcache_type, + llama_multi_gpu=True) xml_output_file = build_dir / f"results-multi-gpu-disagg-asymmetric-executor-{nprocs}-process.xml" @@ -293,21 +293,15 @@ def run_disagg_asymmetric_executor_tests(build_dir: _pl.Path, def run_disagg_orchestrator_params_tests(build_dir: _pl.Path, model: str, - use_ucx_kvcache=False, + kvcache_type=KVCacheType.MPI, timeout=1500): tests_dir = build_dir / "tests" prefix = get_model_test_filter_prefix(model) - cpp_env = {**_os.environ} - mgpu_env = copy.copy(cpp_env) - mgpu_env["RUN_LLAMA_MULTI_GPU"] = "true" - - if use_ucx_kvcache: - mgpu_env["TRTLLM_USE_UCX_KVCACHE"] = "1" - else: - mgpu_env["TRTLLM_USE_MPI_KVCACHE"] = "1" + mgpu_env = get_multi_gpu_env(kv_cache_type=kvcache_type, + llama_multi_gpu=True) xml_output_file = build_dir / "results-multi-gpu-disagg-asymmetric-orchestrator-executor-7-process.xml" trt_model_test = _cpp.produce_mpirun_command( @@ -326,21 +320,15 @@ def run_disagg_orchestrator_params_tests(build_dir: _pl.Path, def run_disagg_spawn_orchestrator_tests(build_dir: _pl.Path, model: str, - use_ucx_kvcache=False, + kvcache_type=False, timeout=1500): tests_dir = build_dir / "tests" prefix = get_model_test_filter_prefix(model) - cpp_env = {**_os.environ} - mgpu_env = copy.copy(cpp_env) - mgpu_env["RUN_LLAMA_MULTI_GPU"] = "true" - - if use_ucx_kvcache: - mgpu_env["TRTLLM_USE_UCX_KVCACHE"] = "1" - else: - mgpu_env["TRTLLM_USE_MPI_KVCACHE"] = "1" + mgpu_env = get_multi_gpu_env(kv_cache_type=kvcache_type, + llama_multi_gpu=True) xml_output_file = build_dir / "results-multi-gpu-disagg-spawn-asymmetric-orchestrator-executor-1-process.xml" @@ -539,14 +527,13 @@ def test_trt_gpt_real_decoder(build_google_tests, multi_gpu_model, lora_setup, indirect=True) class TestDisagg: - @pytest.mark.parametrize("use_ucx_kvcache", [False, True], + @pytest.mark.parametrize("kvcache_type", [KVCacheType.MPI, KVCacheType.UCX], ids=["mpi_kvcache", "ucx_kvcache"]) @pytest.mark.parametrize("nprocs", [2, 4, 8], ids=["2proc", "4proc", "8proc"]) @pytest.mark.parametrize("model", ["gpt", "llama"]) def test_symmetric_executor(self, build_google_tests, model, nprocs, - use_ucx_kvcache, prepare_models_disagg, - build_dir): + kvcache_type, prepare_models_disagg, build_dir): if model == "gpt" and nprocs > 2: pytest.skip( @@ -558,50 +545,47 @@ def test_symmetric_executor(self, build_google_tests, model, nprocs, run_disagg_symmetric_executor_tests(build_dir=build_dir, model=model, nprocs=nprocs, - use_ucx_kvcache=use_ucx_kvcache) + kvcache_type=kvcache_type) - @pytest.mark.parametrize("use_ucx_kvcache", [False, True], + @pytest.mark.parametrize("kvcache_type", [KVCacheType.MPI, KVCacheType.UCX], ids=["mpi_kvcache", "ucx_kvcache"]) @pytest.mark.parametrize("nprocs", [4, 6, 8], ids=["4proc", "6proc", "8proc"]) @pytest.mark.parametrize("model", ["llama"]) def test_asymmetric_executor(self, build_google_tests, model, nprocs, - use_ucx_kvcache, prepare_models_disagg, + kvcache_type, prepare_models_disagg, build_dir): if platform.system() != "Windows": prepare_models_disagg(model_name=model) - run_disagg_asymmetric_executor_tests( - build_dir=build_dir, - model=model, - nprocs=nprocs, - use_ucx_kvcache=use_ucx_kvcache) + run_disagg_asymmetric_executor_tests(build_dir=build_dir, + model=model, + nprocs=nprocs, + kvcache_type=kvcache_type) - @pytest.mark.parametrize("use_ucx_kvcache", [False, True], + @pytest.mark.parametrize("kvcache_type", [KVCacheType.MPI, KVCacheType.UCX], ids=["mpi_kvcache", "ucx_kvcache"]) @pytest.mark.parametrize("model", ["llama"]) - def test_orchestrator_params(self, build_google_tests, model, - use_ucx_kvcache, prepare_models_disagg, - build_dir): + def test_orchestrator_params(self, build_google_tests, model, kvcache_type, + prepare_models_disagg, build_dir): if platform.system() != "Windows": prepare_models_disagg(model) - run_disagg_orchestrator_params_tests( - build_dir=build_dir, - model=model, - use_ucx_kvcache=use_ucx_kvcache) + run_disagg_orchestrator_params_tests(build_dir=build_dir, + model=model, + kvcache_type=kvcache_type) - @pytest.mark.parametrize("use_ucx_kvcache", [True], ids=["ucx_kvcache"]) + @pytest.mark.parametrize("kvcache_type", [KVCacheType.UCX], + ids=["ucx_kvcache"]) @pytest.mark.parametrize("model", ["llama"]) - def test_spawn_orchestrator(self, build_google_tests, model, - use_ucx_kvcache, prepare_models_disagg, - build_dir): + def test_spawn_orchestrator(self, build_google_tests, model, kvcache_type, + prepare_models_disagg, build_dir): if platform.system() != "Windows": prepare_models_disagg(model) run_disagg_spawn_orchestrator_tests(build_dir=build_dir, model=model, - use_ucx_kvcache=use_ucx_kvcache) + kvcache_type=kvcache_type)