Fix Megatron GRPO vocab-parallel log-prob gradients

[taking-lying-flat]

What changed

• Replace the manual target-token gather and non-autograd TP all-reduce with Megatron Core's vocab-parallel cross entropy.
• Preserve -100 label masking by substituting a safe target before cross entropy and masking the resulting token log-probability.
• Compute the full vocab-parallel log-softmax only when entropy is requested.

Why

The previous forward path produced the correct target-token log-probability, but its ordinary torch.distributed.all_reduce did not propagate the remote autograd graph. When the target token belonged to another TP rank, non-owner vocabulary shards received zero or otherwise incorrect logits gradients.

This affects the differentiable current-policy forward used directly by Megatron GRPO policy loss when tensor parallelism is greater than one.

Megatron Core's vocab-parallel cross entropy provides the required custom backward: every vocabulary shard keeps its softmax gradient, while only the shard owning the target token subtracts the target contribution.