NVIDIA-NeMo / NVIDIA-NeMo/RL

Add nightly test coverage for TRT-LLM generation expert parallelism (EP)

Open
#3,304 0 comments 0 reactions 1 assignee Claimed by @shuyixiong View on GitHub
Automation Feature
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

## Context

During review of #2420 (TensorRT-LLM rollout backend), [it was asked](https://github.com/NVIDIA-NeMo/RL/pull/2420#discussion_r3580322964) whether expert parallelism (EP) was tested for the TRT-LLM generation backend and, if so, to cover it in a nightly test. The author [proposed deferring](https://github.com/NVIDIA-NeMo/RL/pull/2420#discussion_r3580322964) EP coverage to a follow-up PR so #2420 can land with the minimal test set. This issue tracks that follow-up.

## Current state (as of #2420)

- `TrtllmSpecificArgs` supports generation-side MoE parallelism via `moe_tensor_parallel_size` / `moe_expert_parallel_size` (`nemo_rl/models/generation/trtllm/config.py`), with the guard `moe_tp * moe_ep == tensor_parallel_size` in `TrtllmGeneration`.
- No recipe sets the generation-side knobs: they are `null` in the exemplar (`examples/configs/grpo_math_1B_trtllm.yaml`) and unset in all three trtllm recipes.
- The only EP in the suite is **training-side** megatron EP (`expert_model_parallel_size: 8`) in the release recipe `grpo-dapomath17k-nanov3-30BA3B-8n4g-megatron-trtllm.yaml`.

## Task

- [ ] Validate TRT-LLM generation with `moe_expert_parallel_size > 1` (MoE model, e.g. the nanov3-30BA3B setup)
- [ ] Add a nightly test/recipe exercising generation-side EP (register in the appropriate `tests/test_suites/nightly*.txt`)

## References
- Originating thread: https://github.com/NVIDIA-NeMo/RL/pull/2420#discussion_r3580322964
- TRT-LLM backend PR: #2420

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.