Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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: 3 additions & 9 deletions jenkins/L0_Test.groovy
Original file line number Diff line number Diff line change
Expand Up @@ -1892,15 +1892,9 @@ def launchTestJobs(pipeline, testFilter, dockerNode=null)
]
fullSet += SBSASlurmTestConfigs.keySet()

multiNodesSBSAConfigs = [
// Each stage test 1 testcase with 8 GPUs and 2 nodes.
"GB200-8_GPUs-2_Nodes-PyTorch-Post-Merge-1": ["gb200-multi-node", "l0_gb200_multi_nodes", 1, 6, 8, 2],
"GB200-8_GPUs-2_Nodes-PyTorch-Post-Merge-2": ["gb200-multi-node", "l0_gb200_multi_nodes", 2, 6, 8, 2],
"GB200-8_GPUs-2_Nodes-PyTorch-Post-Merge-3": ["gb200-multi-node", "l0_gb200_multi_nodes", 3, 6, 8, 2],
"GB200-8_GPUs-2_Nodes-PyTorch-Post-Merge-4": ["gb200-multi-node", "l0_gb200_multi_nodes", 4, 6, 8, 2],
"GB200-8_GPUs-2_Nodes-PyTorch-Post-Merge-5": ["gb200-multi-node", "l0_gb200_multi_nodes", 5, 6, 8, 2],
"GB200-8_GPUs-2_Nodes-PyTorch-Post-Merge-6": ["gb200-multi-node", "l0_gb200_multi_nodes", 6, 6, 8, 2],
]
multiNodesSBSAConfigs = (1..7).collectEntries { i ->
["GB200-8_GPUs-2_Nodes-PyTorch-Post-Merge-${i}".toString(), ["gb200-multi-node", "l0_gb200_multi_nodes", i, 7, 8, 2]]
}
fullSet += multiNodesSBSAConfigs.keySet()

if (env.targetArch == AARCH64_TRIPLE) {
Expand Down
36 changes: 30 additions & 6 deletions tensorrt_llm/llmapi/llm_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,13 @@
import weakref
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Callable, List, Optional, Tuple, Union
from typing import Any, Callable, List, Optional, Tuple, Union

import torch
from tqdm import tqdm

from .._utils import (global_mpi_rank, mpi_barrier, mpi_broadcast, mpi_rank,
release_gc)
from .._utils import (global_mpi_rank, local_mpi_rank, mpi_barrier,
mpi_broadcast, mpi_rank, release_gc)
from ..auto_parallel import AutoParallelConfig
# yapf: disable
from ..bindings.executor import (BatchingType, CapacitySchedulerPolicy,
Expand Down Expand Up @@ -607,6 +607,17 @@ def workspace(self) -> Path:
self._workspace, tempfile.TemporaryDirectory) else Path(
self._workspace)

def _submit_to_all_workers(
self,
task: Callable[..., Any],
*args,
**kwargs,
) -> List[Any]:
if self.llm_args.parallel_config.is_multi_gpu:
return self.mpi_session.submit_sync(task, *args, **kwargs)
else:
return [task(*args, **kwargs)]

def __call__(self) -> Tuple[Path, Union[Path, None]]:

if self.llm_args.model_format is _ModelFormatKind.TLLM_ENGINE:
Expand All @@ -627,9 +638,11 @@ def __call__(self) -> Tuple[Path, Union[Path, None]]:
f'backend {self.llm_args.backend} is not supported.')

if self.model_loader.model_obj.is_hub_model:
self._hf_model_dir = download_hf_model(
self.model_loader.model_obj.model_name,
self.llm_args.revision)
hf_model_dirs = self._submit_to_all_workers(
CachedModelLoader._node_download_hf_model,
model=self.model_loader.model_obj.model_name,
revision=self.llm_args.revision)
self._hf_model_dir = hf_model_dirs[0]
else:
self._hf_model_dir = self.model_loader.model_obj.model_dir

Expand Down Expand Up @@ -806,6 +819,17 @@ def build_task(engine_dir: Path):

return self.get_engine_dir()

@print_traceback_on_error
@staticmethod
def _node_download_hf_model(
model: str,
revision: Optional[str] = None,
) -> Optional[Path]:
if local_mpi_rank() == 0:
return download_hf_model(model, revision)
else:
return None

@print_traceback_on_error
@staticmethod
def _node_build_task(
Expand Down
19 changes: 14 additions & 5 deletions tests/integration/defs/accuracy/test_llm_api_pytorch.py
Original file line number Diff line number Diff line change
Expand Up @@ -2028,16 +2028,25 @@ def test_fp8_block_scales(self, tp_size, pp_size, ep_size, attention_dp,
task.evaluate(llm)

@pytest.mark.parametrize(
"tp_size,pp_size,ep_size,attention_dp,cuda_graph,overlap_scheduler",
[(1, 1, 1, False, True, True)],
ids=["latency"])
"tp_size,pp_size,ep_size,attention_dp,cuda_graph,overlap_scheduler,is_cached",
[(1, 1, 1, False, True, True, True),
pytest.param(8,
1,
1,
False,
True,
True,
False,
marks=pytest.mark.skip_less_mpi_world_size(8))],
ids=["latency", "multi_gpus_no_cache"])
def test_bf16(self, tp_size, pp_size, ep_size, attention_dp, cuda_graph,
overlap_scheduler):
overlap_scheduler, is_cached):
pytorch_config = dict(
disable_overlap_scheduler=not overlap_scheduler,
cuda_graph_config=CudaGraphConfig() if cuda_graph else None)

with LLM(f"{llm_models_root()}/Qwen3/Qwen3-8B",
with LLM(f"{llm_models_root()}/Qwen3/Qwen3-8B"
if is_cached else "Qwen/Qwen3-8B",
tensor_parallel_size=tp_size,
pipeline_parallel_size=pp_size,
moe_expert_parallel_size=ep_size,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,3 +19,4 @@ l0_gb200_multi_nodes:
- accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4[latency_moe_cutlass] TIMEOUT (180)
- accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4[latency_moe_trtllm] TIMEOUT (180)
- accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4[latency_moe_trtllm_eagle3] TIMEOUT (180)
- accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_bf16[multi_gpus_no_cache] TIMEOUT (180)