diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 8bc9ce6154d9..c0fe3942739a 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -210,16 +210,12 @@ def __init__( ) if self.use_latent_moe: - # TODO: check if using ReplicatedLinear is better than - # ColumnParallelLinear + all_gather - self.fc1_latent_proj = ColumnParallelLinear( + self.fc1_latent_proj = ReplicatedLinear( input_size=config.hidden_size, output_size=self.moe_hidden_size, bias=config.mlp_bias, quant_config=quant_config, disable_tp=self.is_sequence_parallel, - # We need to gather the output to prepare input for moe - gather_output=True, prefix=f"{prefix}.fc1_latent_proj", ) self.fc2_latent_proj = ReplicatedLinear(