[NemotronH] Use ReplicatedLinear for fc1_latent_proj

[roikoren755]

Purpose

Current implementation of latent MoE uses a ColumnParallelLinear for the fc1_latent_proj layer. After profiling, the synchronization overhead is quite substantial, and given that the layer itself is relatively small, replicating it instead improves performance while not costing too much in terms of VRAM usage.

Test Plan

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request replaces ColumnParallelLinear with ReplicatedLinear for the fc1_latent_proj layer in the latent MoE implementation. As described, this is a sound performance optimization for a relatively small layer, trading a small increase in VRAM for a reduction in synchronization overhead. The code change is correct and well-contained. I approve this change.

[tlrmchlsmth]

Anything you can share on the performance difference? Mainly just curious