Skip to content

Commit 04fcf24

Browse files
authored
Fix LLM deploy test failure by defaulting expert parallelism to 1 (#1273)
### What does this PR do? Type of change: Bug fix Fixes TRT-LLM DeepEP kernel failures during LLM deployment on unsupported GPUs (e.g. Blackwell SM 12.0) by defaulting expert parallelism (`ep`) to 1 instead of auto-setting it to the GPU count for MoE models. Previously, when the model config contained expert-related keys, `ep` was automatically set to `torch.cuda.device_count()`, which triggered DeepEP kernel failures on GPUs that don't support it. Now `ep` defaults to 1 while still enabling attention data parallelism for MoE models. Expert parallelism can be enabled explicitly by the caller when the environment is known to support it. ### Testing - [x] Verified that the `llm_ptq` test passes with this fix on Blackwell GPUs. - [x] 2-gpu CI test triggered: https://github.com/NVIDIA/Model-Optimizer/actions/runs/24495054531/job/71588037727 ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
1 parent 6ded36b commit 04fcf24

1 file changed

Lines changed: 3 additions & 2 deletions

File tree

modelopt/deploy/llm/generate.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -109,12 +109,13 @@ def _find_max_position_embeddings(cfg: dict) -> int | None:
109109
if tp < 1:
110110
tp = torch.cuda.device_count()
111111

112-
# Check if any key in config contains both "num" and "experts"
112+
# Force ep=1 to avoid TRT-LLM DeepEP kernel failures on unsupported GPUs
113+
# (e.g. Blackwell SM 12.0). Expert parallelism can be enabled explicitly
114+
# by the caller when the environment is known to support it.
113115
ep = 1
114116
enable_attention_dp = False
115117
for k in config:
116118
if "num" in k and "experts" in k:
117-
ep = torch.cuda.device_count()
118119
enable_attention_dp = True
119120
break
120121

0 commit comments

Comments
 (0)