[NemotronH] Use ReplicatedLinear for fc1_latent_proj (#16569)

Signed-off-by: Roi Koren <roik@nvidia.com>
This commit is contained in:
roikoren755
2026-01-14 17:08:27 +08:00
committed by GitHub
parent ba625c2d90
commit 72bacc88c8
+1 -2
View File
@@ -191,12 +191,11 @@ class NemotronHMoE(nn.Module):
self.shared_experts = None self.shared_experts = None
if self.use_latent_moe: if self.use_latent_moe:
self.fc1_latent_proj = ColumnParallelLinear( self.fc1_latent_proj = ReplicatedLinear(
input_size=config.hidden_size, input_size=config.hidden_size,
output_size=self.moe_hidden_size, output_size=self.moe_hidden_size,
bias=config.mlp_bias, bias=config.mlp_bias,
quant_config=quant_config, quant_config=quant_config,
gather_output=True,
prefix=f"{prefix}.fc1_latent_proj", prefix=f"{prefix}.fc1_latent_proj",
) )
self.fc2_latent_proj = ReplicatedLinear( self.fc2_latent_proj = ReplicatedLinear(