fix: use tied embedding for linear CE fusion output weight

[jthomson04]

Summary

• _gpt_forward_with_linear_ce_fusion set output_weight_layer = self.output_layer.weight, which is None for models with tied embeddings (e.g. Qwen3 with share_embeddings_and_output_weights=True). The first positional argument to from_parallel_hidden_states_to_logprobs was therefore None, crashing logprob computation when policy.use_linear_ce_fusion_loss=True is set on a tied-embedding model.
• Fetch the weight via shared_embedding_or_output_weight() when embeddings are tied, matching the logic already used for the second positional argument, and reuse the same tensor for both args.

Test plan

• Run a logprob/training step on a Qwen3 model with policy.use_linear_ce_fusion_loss=True and confirm no crash and matching logprobs vs. the non-fusion path.
• Confirm untied-embedding models (e.g. Llama) still produce identical logprobs (the non-tied branch is unchanged).

🤖 Generated with Claude Code

[copy-pr-bot]

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[jthomson04]

/ok to test f3466a3

[terrykong]

LGTM — clean, minimal bug fix. The tied-embedding path was passing None (self.output_layer.weight) as the first positional arg to from_parallel_hidden_states_to_logprobs, while the correct weight was only computed for the second arg (which turns out to be unused in the callee). The fix correctly resolves the weight via shared_embedding_or_output_weight() for tied-embedding models.

Minor observation (non-blocking): from_parallel_hidden_states_to_logprobs accepts an output_weight parameter (its 3rd positional arg) that is never referenced in the function body — only output_weight_layer (2nd positional arg) gets passed into ChunkedDistributedHiddenStatesToLogprobs.apply(). This dead parameter pre-dates this PR, but could be cleaned up in a follow-up to avoid future confusion.

Generated by Claude Code

[terrykong]

nemo_rl/distributed/model_utils.py:2165

Nit: both args now pass the same output_weight_layer, which is correct since the callee's output_weight parameter (line 1836) is never actually read — only output_weight_layer is forwarded to ChunkedDistributedHiddenStatesToLogprobs.apply(). Consider removing the dead output_weight parameter in a follow-up to prevent future confusion.

[jthomson04]

Good catch — dropped the dead output_weight parameter in 68e1b23.