[NemotronH] Use ReplicatedLinear for fc1_latent_proj (#16569)
Signed-off-by: Roi Koren <roik@nvidia.com>
This commit is contained in:
@@ -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(
|
||||||
|
|||||||
Reference in New Issue
Block a user