A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 16:30 UTC
[NemotronH] Use ReplicatedLinear for fc1_latent_proj (#31807)
Signed-off-by: Roi Koren <roik@nvidia.com>
roikoren755 committed
Jan 6, 2026 at 18:00 UTC
28c94770adfcb9cfbc78a3221f915bfc830c6582
1 file changed
+1
-5
vllm/model_executor/models/nemotron_h.py
+1
-5
@@ -210,16 +210,12 @@ class NemotronHMoE(nn.Module):
210
)
211
212
if self.use_latent_moe:
213
- # TODO: check if using ReplicatedLinear is better than
214
- # ColumnParallelLinear + all_gather
215
- self.fc1_latent_proj = ColumnParallelLinear(
213
+ self.fc1_latent_proj = ReplicatedLinear(
214
input_size=config.hidden_size,
215
output_size=self.moe_hidden_size,
216
bias=config.mlp_bias,
217
quant_config=quant_config,
218
disable_tp=self.is_sequence_parallel,
221
- # We need to gather the output to prepare input for moe
222
- gather_output=True,
219
prefix=f"{prefix}.fc1_latent_proj",
220
)
221
self.fc2_latent_proj = ReplicatedLinear(